TokenGate Speeds Up LLM Prefill by Up to 11.3x with Learned Token Pruning

A team of researchers has published TokenGate, a small plug-in component that lets already-trained language models skip low-value tokens in long inputs. According to the authors, time to first token (TTFT) falls by 1.4–11.3x, peak prefill memory drops by 1.5–3.4x, and accuracy does not degrade across a range of tasks. Against the best existing token-pruning methods, the gain reaches up to 6.4x in prefill speed and up to 2x in memory.

Anyone who has run models on large code repositories or hefty documents will know the problem. A transformer pushes every input token through every layer, whether or not that token is needed for the answer. A human expert works differently: they skim the text and read closely only the parts that matter. The authors build on that analogy and try to teach the model to do the same.

How to get around non-differentiable top-K

Token pruning itself is not new. The difficulty is in training: selecting the top-K tokens is a non-differentiable operation, no gradient passes through it, and the importance scorer cannot be trained end-to-end together with the model. Many earlier approaches therefore rely on heuristics or separate training stages.

TokenGate solves this with a mechanism called token-gated attention. It is inspired by Mixture-of-Experts: token importance scores reweight the attention logits. As a result, the scorer receives a gradient signal during ordinary training, and the whole chain, including the decision about what to keep, is trained end-to-end.

One important practical detail: this is not a new architecture built from scratch. TokenGate extends existing models. The authors show it working with both standard transformers (for example, Qwen3) and hybrid attention models (for example, Qwen3.5). By their description, the adaptation is lightweight and comparable to standard parameter-efficient fine-tuning.

What the measurements showed

The figures from the paper are summarized below.

• TTFT: 1.4–11.3x faster
• Peak prefill memory: 1.5–3.4x lower
• Accuracy: no degradation on the tasks tested
• Versus the best earlier pruning methods: up to 6.4x faster prefill and up to 2x lower peak memory
• Compatibility: standard transformers (Qwen3) and hybrid attention (Qwen3.5)
• Code: open, repository tokengate-models/TokenGate on GitHub

The ranges are wide, and that is expected: the gain depends on the task, the input length, and the share of tokens that can be safely discarded. The abstract does not say which configurations produce the extreme values, so it is better to rely on your own measurements than on the 11.3x maximum.

Where it could help and what to check

The method promises the biggest benefit where prefill dominates latency: coding agents that load entire repositories into context, RAG systems with long excerpts, and document and log analysis. Lower peak memory also means that a longer context may fit on the same GPU, or that more requests can be processed in parallel.

Before adopting it, check a few things:

1. Run TokenGate on your own data. "No accuracy degradation" in the paper refers to the authors' set of tasks, and yours may be stricter. 2. Test separately the scenarios where precise retrieval of a single detail from a long text matters: that is where aggressive token pruning risks throwing out the one thing you need. 3. Estimate the cost of fine-tuning: it is claimed to be comparable to PEFT, but it is worth calculating for your model. 4. Check how the component fits into your serving stack and whether your inference engine supports it.

The limitation is obvious: this is a preprint, the numbers come from the authors, and there is no independent reproduction yet. But the code is open and the method applies to the popular Qwen families, so the claims can be checked quickly.