TokenGate teaches small LLMs to skim long documents, cutting prefill time up to 11x in the paper's tests

Long prompts are expensive before the first token appears. Every document token in a RAG context goes through every layer and leaves an entry in the KV cache, even when most of it is irrelevant to the question. TokenGate, a paper posted on 19 October by Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia and Carlos Guestrin, trains a model to skim instead. Code is on GitHub under Apache-2.0, and six trained checkpoints are on Hugging Face.

How it works

TokenGate adds a small gate module inside one early layer of an existing model. Layers up to the gate read the whole context as usual. The gate then scores every document token and keeps only the top K. Layers above it process just those tokens, so pruned tokens never get a KV cache entry in the upper layers. That is where the savings in time and memory come from.

Picking the top K is not differentiable, so the authors borrow from mixture-of-experts routing. In what they call token-gated attention, each kept token’s importance score scales its attention weight in the layers above. That gives the gate a gradient, so it learns which tokens matter for the task. The gate is trained together with a LoRA adapter (rank 16), and the paper puts the extra training cost at about 1.3 times a plain LoRA step. In use, you mark the prunable document in the prompt and set a keep budget, max_kept_prefill_tokens.

What the paper reports

The tests used Qwen3-1.7B and the hybrid Qwen3.5-2B, on multi-hop question answering, code completion, summarization, biomedical extraction and needle-in-a-haystack retrieval. Timings are at batch size 1 on one A100. The authors report:

• 1.4 to 11.3 times faster time to first token and 1.5 to 3.4 times lower peak prefill memory than full-context prefill, with no loss of accuracy.
• Up to 6.4 times faster prefill and 2 times lower peak memory than the best training-free pruning methods they tested, including FastKV.
• On 2WikiMultihopQA, prefill 3.0 times faster with 2.7 times less memory while keeping 95% of full-context accuracy.
• A needle-in-a-haystack model trained at 16K tokens held its accuracy at 32K, 64K and 128K, with time to first token about 13 times faster at the longer lengths.
• On Qwen3.5-2B, whose Gated DeltaNet layers FastKV cannot prune, 3.8 times lower time to first token at no accuracy loss.

The catches

Each result comes from a model fine-tuned on that task’s training set, so TokenGate is a replacement for task fine-tuning, not a switch you flip on a general chat model. Every result is on models of 2 billion parameters or less. The fast path also depends on a custom FlashAttention-2 kernel that you build with nvcc; without it, the code falls back to FlexAttention. The released checkpoints are tied to the paper’s tasks, one per dataset.

The idea fits a real gap. Retrieval already narrows which documents reach the model, and TokenGate narrows which tokens inside them get read in full. Whether it holds up on larger models and mixed workloads is the open question.