Prune prefill tokens to cut TTFT up to 11.3× with TokenGate — Token-gated attention — What does it mean?
The news. On October 19, 2026, researchers at Stanford and UC Berkeley released TokenGate, a drop-in component for pretrained transformers that learns which long-context tokens to prune during prefill. They report 1.4–11.3× faster time to first token and 1.5–3.4× lower peak prefill memory with no accuracy degradation against the full-context fine-tuned baseline, on Qwen3-1.7B and a hybrid-attention Qwen3.5 model. Read the paper →
Picture a review office handed a 300-page stack and one question. A careful office does not make every reader read every page at every floor. A triage clerk skims the stack once, highlights what matters, and only the highlighted pages go upstairs. A transformer today does the opposite: during prefill every prompt token passes through every layer, whether or not it matters to the question. The bill has three parts: attention compute that grows with the square of the prompt length, MLP compute that grows linearly, and KV-cache memory that also grows linearly.
TokenGate puts the clerk at one early layer, called L. Layers below L run as usual on every token. At L, a small scoring network (an MLP) gives each prunable token an importance score, a top-K selector keeps the best K, and every layer above L processes only those K tokens. The pruned tokens are never forwarded again and never create KV-cache entries in the upper layers. Tokens such as the user's question can be marked non-prunable, so the clerk never discards the question itself.
The hard part is training the clerk. A top-K cut is a yes-or-no choice, so a plain loss gives the scoring network no direction to improve. Token-gated attention fixes this by wiring the score into the upstairs readers. Each kept token's score becomes a gate value g, a softmax over the kept tokens, and the attention weight of that token is multiplied by g inside every layer above L. Now the loss can tell the scorer which way to move. The paper's intuition is that a token whose extra attention mass lowers the loss receives a larger score, and a token that distracts receives a smaller one. This is the same move a Mixture-of-Experts router makes when its score multiplies the output of the experts it picks, applied to tokens instead of experts. The paper adds a learnable gain per layer and head, plus a temperature, so the gate keeps working at different pruning ratios. At the start of training all scores are zero, so the model behaves exactly like plain top-K attention on the unmodified backbone.
The training recipe is ordinary fine-tuning, not a new pre-training run. The scorer and the gains are trained together with LoRA on the backbone using the usual next-token loss, with no auxiliary loss. Each step averages two losses: one with all tokens kept, which protects full-context accuracy and gives every token a gradient, and one with a budget K drawn log-uniformly from 1 up to half the prunable tokens, which teaches the gate to work across budgets. The result is a model that decides per task what to skim, and the paper's visualizations show different kept-token patterns for multi-hop question answering and for summarization.
| Approach | Needs training? | What it changes | Limit named in the paper |
|---|---|---|---|
| Training-free token pruning (FastKV, SlimInfer and similar) | No | Scores tokens with attention heuristics and forwards only the top ones | Limited speedups before accuracy degrades |
| Architecture changes (state-space, mixture-of-depths and similar) | Yes, often from scratch | Replaces the standard transformer design | Expensive to train, needs different hardware and software |
| TokenGate | Light adaptation (LoRA plus gate) | Adds a learned gate at one layer of an existing model | Tested on Qwen3-1.7B and a Qwen3.5 hybrid; pruning needs a training run per model |
Here is how the savings compose, using round numbers (illustrative, not from the paper). Take a 28-layer model, a 32,000-token document, a gate after layer 1, and K = 3,200 kept tokens (10%). Without pruning, the token-layer passes are 28 × 32,000 = 896,000. With the gate, the first two layers still see everything and the other 26 see only the kept tokens: 2 × 32,000 + 26 × 3,200 = 64,000 + 83,200 = 147,200. That is about 6.1× fewer token-layer passes for the MLPs and about 6.1× fewer KV-cache entries. Attention is steeper still: pairwise attention work in each upper layer falls from 32,000² to 3,200², a 100× cut per layer. Real speedups land below the raw count because the lower layers, the scorer and the question tokens still cost something, which is consistent with the paper's measured range of 1.4–11.3× TTFT rather than a single headline number.
Fewer tokens in the upper layers means fewer matrix multiplies and fewer KV entries to store, and that is the whole saving. The scope is narrow, though. TokenGate speeds prefill for inputs you cannot cache in advance, such as a document uploaded at request time, and it needs a trained gate for each backbone. The paper reports results on two model families and five task types, so treat the gains as shown for those settings.
Goes deeper in: LLM Internals → Attention → Computing Attention Scores
Related explainers
- Self-pruned KV cache — prunes tokens from the cache instead of from the layers, a different point in the pipeline.
- Prefill-only LoRA adapters — another way to change only how the prompt is read.