TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning
AuthorsLiana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia, Carlos Guestrin
AffiliationsStanford University · UC Berkeley https://github.com/tokengate-models/TokenGate
TokenGate teaches language models to skip less relevant tokens during prefill, aiming to cut latency and memory without sacrificing accuracy.
Key results
Fold faster than full-context prefill at no accuracy degradation.
Fold lower peak prefill memory than full-context prefill at no accuracy degradation.
Fold faster than the next-best token-pruning baseline at matched accuracy.
Fold lower peak prefill memory than the next-best token-pruning baseline at matched accuracy.
Fold lower time-to-first-token than full-context prefill with no accuracy degradation.
What the paper found
TokenGate adds a learned pruning gate to pretrained language models, letting them discard less relevant context tokens partway through prefill so later layers process fewer tokens and build a smaller KV cache. Its scorer selects the top-K tokens, while token-gated attention uses their importance scores to provide a gradient path through the otherwise non-differentiable selection step. The model and gate are adapted together with LoRA, using ordinary next-token training and sampled token budgets. On tasks including 2WikiMultihopQA, MultiNews, RepoBench-P, BioDEX, and needle-in-a-haystack, TokenGate reached up to 11.3-fold faster time-to-first-token and 3.4-fold lower peak prefill memory than full-context prefill without accuracy loss; against the strongest prior pruning methods, it was up to 6.4-fold faster and used up to 2-fold less memory at matched accuracy. On the hybrid Qwen3.5-2B model, it achieved 3.8-fold lower time-to-first-token with no accuracy degradation. The results show that learned, task-adaptive token selection can improve prefill efficiency without replacing a model’s existing architecture, including models such as Qwen3 and Qwen3.5.
Original abstract
Expert human readers process large inputs (e.g., code repositories and long documents) by skimming to closely read only what matters. On the other hand, language models process all input tokens through each model layer, resulting in significant processing costs for long-context inputs. We introduce TokenGate, a small, drop-in component that extends existing trained models to enable end-to-end learnable token pruning over long-context inputs for fast and memory-efficient prefill. To enable end-to-end learning despite the non-differentiable top-K operator for token pruning, we introduce token-gated attention, a mechanism inspired by Mixture-of-Experts that re-weights attention logits using token importance scores. Across diverse tasks, TokenGate achieves $1.4$-$11.3\times$ faster time-to-first token (TTFT) and $1.5$-$3.4\times$ lower peak memory at no accuracy degradation for prefill in transformers. Compared to the best-performing prior methods for token-pruning, TokenGate gives up to $6.4\times$ faster prefill and $2\times$ lower peak memory. TokenGate natively extends both standard transformers (e.g., Qwen3) and hybrid-attention models (e.g., Qwen3.5) with simple lightweight adaptation comparable to standard parameter-efficient fine-tuning. We make our code available at: https://github.com/tokengate-models/TokenGate.
Read the original paperMore in Efficient AI
Browse all 70 papers →AB-Attention: Accurate Block Attention for Accelerating Video Diffusion Transformers
Bin Xu, Mo Zhang, Zhengfeiyang Zhang, Xinfeng Yang, Yixing Ma, Wei Wang, HeXin Wang, Pengfei Zhou, Wangbo Zhao
AB-Attention speeds up video generation by skipping most attention computations while more accurately estimating the attention it omits.
Extreme Compression of Per-Layer Embeddings via Cross-Layer Latent Vector Quantization
Jiacheng Liu, Yaxin Luo, Xiaohan Zhao, Peng Sun, Xinyi Shang, Jiacheng Cui, Yonina C. Eldar, Zhiqiang Shen
A new way to compress Transformer embeddings dramatically reduces memory use while largely preserving model performance.
Rethinking Sequence Compression with Test-time Regression
Sihan Yang, Jie Huang, Hanxuan Li, Yifan Guo, Weigao Sun, Yiran Zhong, Weiyang Liu
A new regression-based compressor aims to keep language and vision models efficient while preserving more useful information in fewer tokens.