TokenGate speeds up LLM prefill by up to 11.3x with learnable token pruning

Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia, and Carlos Guestrin have introduced TokenGate — a component that makes prefill in large language models significantly faster and more memory-efficient. Instead of processing every input token through each transformer layer, TokenGate lets models selectively skip unimportant context, much like a person skimming a long document: quickly scanning the content and reading only the key parts carefully.

The problem TokenGate addresses is acute for practical LLM deployment: processing long contexts demands substantial compute and memory. When a model receives thousands or tens of thousands of input tokens (say, an entire repository's source code or a long document), it's forced to push every token through all attention layers. TokenGate introduces a learnable token pruning mechanism that determines which tokens actually matter for the task at hand.

The token-gated attention mechanism

At the core of TokenGate's innovation is token-gated attention, a mechanism inspired by Mixture-of-Experts architecture. The key engineering insight solves the non-differentiability problem: instead of a hard top-K operator for pruning tokens (which makes end-to-end training difficult), the system uses a soft re-scoring approach. Each token is assigned an importance score that re-weights its contribution to the attention computation. This allows the entire pruning mechanism to be trained jointly with the rest of the model in a single process, without special tricks or extra stages.

Results and performance

In experiments across diverse tasks, TokenGate delivers impressive results:

• Time-to-first token (TTFT) speedup: 1.4–11.3x depending on the task and context length • Peak memory reduction: 1.5–3.4x without loss of accuracy • Comparison with the best competitors: up to 6.4x faster prefill and half the peak memory versus existing token-pruning methods

Compatibility and integration

TokenGate is designed as a drop-in component that extends already-trained models without requiring retraining. The system works with both standard transformers (e.g., Qwen3) and next-generation hybrid attention models (e.g., Qwen3.5). Adaptation requires only simple parameter-efficient techniques, comparable in complexity to standard parameter-efficient fine-tuning (PEFT).

The authors are sharing their results with the community: full code is available in the repository https://github.com/tokengate-models/TokenGate. This lets other researchers and engineers immediately integrate TokenGate into their own systems.

TokenGate addresses a real problem in deploying LLMs in production: the need to process long contexts quickly without excessive resource use. The learnable token pruning approach with token-gated attention opens a path toward more efficient use of modern large language models' capabilities.