Stanford and UC Berkeley propose TokenGate: token-gated attention makes long-context prefill up to 11.3x faster

Add as a preferredsource on Google

Stanford and UC Berkeley researchers released TokenGate, a module for long-context prefill. TokenGate learned token pruning lets an existing LLM skim its input. Only the tokens that matter get full processing. The research team reports 1.4x to 11.3x faster time-to-first-token (TTFT) with no accuracy loss. Peak prefill memory drops 1.5x to 3.4x. Code and checkpoints are on GitHub under Apache-2.0.

TL;DR

MLP gate: A small MLP gate and LoRA (r=16) added to Qwen3-1.7B or Qwen3.5-2B. Tested on contexts up to 128k tokens.

Runs on: Paper timings use 1 A100 at batch size 1. Training used 4 A100s.

Performance: Pareto-optimal on 5 tasks against 6 baselines.

Best: NIAH 16k at 11.3x faster TTFT and 3.42x lower memory. TTFT gain hits 13.0x at 128k.

Worst: Speedup drops to 1.4x on some tasks.

What is TokenGate?

TokenGate is a drop-in component that adds trainable token pruning to pretrained LLMs. Transformers push every token through every layer. Attention cost grows quadratically with length. Training-free pruning is cheap but loses accuracy at high pruning rates. Architectural changes often need pre-training from scratch.

TEXTTABLELAYOUTIs your PDF parser actually getting it right?omni_parse_bench · scores text, tables & layout · open sourceStar on GitHub→Sponsored

How does TokenGate learned token pruning work?

The forward pass splits into 3 phases:

Standard forward: layers below gate layer L process all N tokens as usual.

Gate module: an MLP scores each token from its post-attention residual at layer L. A top-K selector keeps the highest-scoring tokens.

Pruned forward: layers above L process only those K tokens and never create KV entries for pruned ones.

Only the document part of the prompt is prunable. The question is never pruned.

What is token-gated attention?

Top-K selection is not differentiable, which blocks end-to-end training. TokenGate borrows the Mixture-of-Experts router trick. Each kept token’s score becomes a bias on attention logits above the gate. That bias gives gradients a path back to the scorer.

A learnable per-head gain and a learnable temperature refine this. The temperature corrects for each query seeing fewer keys. Training uses plain next-token loss. Each step averages a full-context loss and a random-budget loss. So 1 trained gate works across budgets.

How does TokenGate perform on benchmarks?

The research team tested Qwen3-1.7B on RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX and needle-in-a-haystack (NIAH). Baselines were FastKV, SlimInfer, DASH, GemFilter, FlexPrefill and MInference, on the same LoRA backbone.

Versus full prefill: 1.4x to 11.3x faster TTFT and 1.5x to 3.4x lower memory.

Versus the next-best baseline: 1.2x to 6.4x faster TTFT and up to 2.1x lower memory.

2WikiMultihopQA: at 95% of full-context accuracy, prefill runs 3.0x faster with 2.7x less memory.

Longer contexts: a gate trained at 16k reaches 12.8x at 32k and 13.0x at 64k and 128k. Memory savings stay near 3.43x.

The learned selector beats FastKV’s heuristic at the same layer. Even random-K selection scores higher on the TokenGate backbone. The research team credits backbone co-adaptation.

Does it work on hybrid-attention models?

Yes, Qwen3.5-2B mixes Gated DeltaNet and full-attention layers. On NIAH 32k, TokenGate at layer 3 matched full-prefill accuracy. It ran 3.8x faster than full prefill and 2.1x faster than FastKV. Gating at layers 1 or 2 reached 7.8x, within 6% accuracy. FastKV can only prune attention layers.

FeatureTokenGateFastKVSlimInferGemFilterDASH

Token selectionLearned MLP gateAttention from last queriesBlock attention, multi-stageLast-query attention at filter layerHidden-state stability

Training neededYes (gate + LoRA r=16)NoNoNoNo

Pruning layer in TokenGate’s Qwen3-1.7B runs113 (default)8, 17, 251211

Prunes at Gated DeltaNet layersYesNo (attention layers only)Not disclosedNot disclosedNot disclosed

Best self-reported prefill gain (own setup)11.3x TTFT, 3.4x memory1.82x prefill2.53x TTFT (LLaMA3.1-8B, RTX 4090)2.4x speed, 30% less GPU memory vs SOTANot disclosed

CodeGitHubGitHubGitHubGitHubAnnounced

Self-reported gains use different setups, so they are not directly comparable. Pruning layers come from TokenGate’s Appendix C.2.

How do you deploy TokenGate?

The GitHub Repo needs Python 3.10+ and CUDA PyTorch 2.9+. An optional FlashAttention-2 key-bias kernel speeds up dense Qwen3. It costs 1.0x to 1.1x a plain call. Without it, inference falls back to FlexAttention.

A {state} placeholder marks the prunable document. generate takes a max_kept_prefill_tokens budget. Checkpoints cover all 5 tasks on Qwen3-1.7B, plus a Qwen3.5-2B NIAH gate. Example: liana313/tokengate-niah-qwen3-1.7b.

Training mirrors a TRL/PEFT recipe. Each step is 1.3x slower than plain LoRA. Runs converged faster on 4 of 5 tasks. Mean training time was 1.8 versus 2.2 A100 hours.

Key Takeaways

TokenGate adds learned token pruning to existing LLMs via a small gate and LoRA.

Prefill runs 1.4x to 11.3x faster with 1.5x to 3.4x lower peak memory.

Token-gated attention makes top-K selection trainable, MoE-router style.

It beats the next-best pruning baseline by up to 6.4x on TTFT.

It works on hybrid Qwen3.5-2B, where FastKV cannot prune DeltaNet layers.

Asif RazzaqWebsite | + postsBio ⮌

Asif RazzaqNVIDIA Releases Nemotron 3.5 Bolt E4B 4GB: An Open Small Language Model for Local Voice Agents, Game Characters and Robots

Asif RazzaqMoonshot AI Releases Kimi K3.1 Pro: A 2.8T Total/104B Active Params (MoE) Model with 1M Context

$uv pip install omni-parse-bench✓text✓tables✓layout→ scored per pageStar on GitHub→Sponsored

RELATED ARTICLESMORE FROM AUTHOR

How to Run Underdog Saluki 27B: 2-Bit Qwen3.8-27B Tool-Calling Guide

Meet Anthropic’s sandbox-runtime: OS-Level Sandboxing for AI Agents, No Containers