Stanford-Berkeley TokenGate: 11.3x Faster Prefill

TokenGate, a learned token-pruning module from Stanford and UC Berkeley researchers, cuts long-context prefill by up to 11.3× and peak prefill memory by up to 3.4×, according to the team's self-reported benchmarks. The Apache-2.0 release includes a GitHub repository, benchmark scripts, and named Hugging Face checkpoints, which makes the claims easier to reproduce than a paper-only result.

How TokenGate prunes tokens without breaking gradients

A gate at layer L preserves standard attention below L for all N tokens. At L, a small MLP scores every token in the prunable {state} document from its post-attention residual, and a top-K selector keeps the highest-scoring tokens. Above L, only those K tokens receive compute and KV-cache entries, so pruned tokens never consume time-to-first-token or peak-memory above the gate. Instructions, questions, and generated answers are never pruned.

Top-K is non-differentiable. TokenGate's token-gated attention borrows the MoE-router trick: kept tokens' importance scores become additive biases on attention logits above the gate, giving gradients a path to the scorer. A learnable per-head gain and a learnable temperature refine the mechanism; the temperature corrects for queries seeing fewer keys. The researchers report that removing either lowers BioDEX accuracy. Training uses standard next-token prediction and averages a full-context loss with a random-budget loss at each step, so one gate works across budgets. The frozen backbone gains a rank-16 LoRA on q_proj and v_proj.

Benchmark results

All paper timings use one A100 at batch size 1; training used four A100s. Qwen3-1.7B was tested on RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX, and needle-in-a-haystack (NIAH), against FastKV, SlimInfer, DASH, GemFilter, FlexPrefill, and MInference on the same LoRA backbone. Figures are researcher-reported.

Source: TokenGate GitHub repository, examples/niah_demo.py validated output; Qwen3-1.7B NIAH gate, ~32k-token haystacks, 5 samples, 1 GPU — researcher-reported.

The headline 11.3× result is NIAH at 16k context. MarkTechPost reports that a gate trained at 16k reaches 12.8× at 32k and 13.0× at 64k and 128k, while memory savings stay near 3.43×. Against the next-best baseline, the team reports 1.2× to 6.4× faster TTFT and up to 2.1× lower memory; the worst-case speedup across tasks is 1.4×. On 2WikiMultihopQA, reaching 95% of full-context accuracy runs the prefill 3.0× faster with 2.7× less memory. On hybrid Qwen3.5-2B, which mixes Gated DeltaNet and full-attention layers, TokenGate at gate layer 3 matched full-prefill accuracy on NIAH 32k at 3.8× faster TTFT and 2.1× faster than FastKV. Gating at layers 1 or 2 reached 7.8× speedup within 6% accuracy.

Deployment and training overhead

The repository requires Python 3.10+, CUDA, and PyTorch 2.9+. An optional FlashAttention-2 key-bias kernel builds from python kernels/fa2bias/build.py; without it inference falls back to FlexAttention automatically. The kernel costs 1.0× to 1.1× a plain forward call, according to the researchers. Inference marks the prunable document with {state} and takes max_kept_prefill_tokens in generate. Released checkpoints cover all five Qwen3-1.7B tasks and a Qwen3.5-2B NIAH gate. Training runs 1.3× slower than plain LoRA per step, but converged faster on four of five tasks; the team reports mean training time of 1.8 A100-hours versus 2.2 for the baseline.

AI Mastery analysis

The most important result is not the headline speedup but the generalisation from 16k to 128k without retraining. If it holds beyond NIAH—a task where a single planted sentence is unusually gate-friendly—it would shrink the retraining surface for teams facing long-context pricing and quality tradeoffs. The 1.4× worst case implies task-specific gate plus LoRA pairs, not a universal gate.

The co-adaptation claim deserves scrutiny. The researchers report that even random-K selection on the TokenGate backbone scores higher than FastKV's heuristic, and credit backbone co-adaptation. If the LoRA reshapes residual representations to be more gate-friendly, part of the measured speedup comes from backbone adaptation, not the selector alone. This determines whether TokenGate checkpoints transfer across inference stacks or lock teams into a specific one.

Batch-size 1 numbers are clean, but production batches are not answered by this material: multiple sequences sharing KV cache may alter the memory-reduction benefit when K varies per request. The open licence and public checkpoints make that test practical. The broader pattern matches infrastructure-layer gains, not model-weight changes, driving 2026 inference improvements.

Sources

• Primary: GitHub - token-gate-llms/TokenGate: https://arxiv.org/abs/2610.26754 Process all your context faster with LLMs that can skim and read what m
• Reported by: MarkTechPost: Researchers from Stanford and UC Berkeley Introduce TokenGate: Token-Gated Attention Makes Long-Context Prefill Up to 11.3x Faster