TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning

AI summary

Overview

Research area: Machine learning systems and LLM inference efficiency — specifically long-context prefill optimization via learned token pruning.

Technical level: Advanced. The paper combines transformer internals (attention logits, residual streams, KV cache), Mixture-of-Experts-style gating, and custom attention kernels; the experimental sections are readable with intermediate background, but the method section assumes familiarity with attention mechanics.

Scope: The paper introduces TokenGate, a trainable drop-in gate module plus a "token-gated attention" mechanism that lets existing pretrained transformers skip most input tokens in their upper layers, and it reports latency, memory, and accuracy results across five task types and two model families.

What This Paper Is About

Transformers push every input token through every layer, so long prompts (code repositories, multi-document inputs, medical papers) incur attention cost that grows quadratically with input length and MLP compute plus KV-cache memory that grow linearly. Prior fixes split into training-free token pruning, which is cheap to deploy but degrades accuracy before delivering large speedups, and architectural changes, which are expensive because they require pretraining from scratch. The paper's goal is to endow already trained models with learnable token pruning so that prefill becomes faster and lighter without losing accuracy, at an adaptation cost comparable to standard parameter-efficient fine-tuning.

Key Contributions

1. TokenGate, a drop-in gating component. A small module inserted at a chosen transformer layer L: a scoring network (an MLP over each token's normalized post-attention residual, normalized with RMSNorm) scores prunable tokens, and a top-K selector keeps K of them plus all non-prunable tokens. Layers above L then process only the kept tokens and never create KV-cache entries for pruned ones.
2. Token-gated attention. Because top-K selection is non-differentiable, the paper borrows the Mixture-of-Experts routing trick but applies it along a different axis: instead of weighting expert subnetworks to sparsify parameters, it adds a per-key gate bias to the attention logits to sparsify processed tokens, giving the scoring network a gradient path through the task loss.
3. Two calibration mechanisms. A learnable per-layer, per-head gain (passed through softplus, initialized so its scale matches the layer's logit standard deviation) and a learnable temperature scaled by the log-ratio of a query's original position to its rank among kept tokens, intended to make one trained model work across pruning ratios and input lengths.
4. An efficient training and serving recipe. Training uses standard next-token cross-entropy with LoRA on the backbone plus the gate parameters, averaging loss at full budget (K = all prunable tokens) and a log-uniform sampled budget between 1 and |P|/2; serving uses a patched FlashAttention-2 kernel accepting a per-key, per-head bias, plus CUDA-graph capture of the pruned forward pass.

Main Findings

• Prefill speed and memory at full accuracy: Across the evaluated tasks, TokenGate reduces time-to-first-token by 1.4–11.3× and peak prefill memory by 1.5–3.4× at no accuracy degradation relative to the fine-tuned base model run at full context.
• Comparison to prior pruning methods: Against the best-performing prior token-pruning baselines, TokenGate gives up to 6.4× faster prefill and 2× lower peak memory in the abstract and conclusion; Section 5.1 reports up to 2.1× lower peak memory at matched accuracy.
• Pareto-optimality: The accuracy-versus-TTFT and accuracy-versus-peak-memory plots place TokenGate on the Pareto frontier on all evaluated datasets, including against FastKV measured both at its default layer and at layer 1 (the same layer index TokenGate uses).
• Illustrative example (2WikiMultihopQA, Qwen3-1.7B, K=256): at 95% of full-context accuracy, TokenGate prefills 3.0× faster with 2.7× less memory, outperforming FastKV and SlimInfer.
• Hybrid attention support: On Qwen3.5-2B (full-attention layers every fourth layer, Gated-DeltaNet elsewhere), TokenGate at layer 3 matches full-context accuracy at 3.8× lower TTFT and beats the best FastKV run by 2.1× lower TTFT at the same accuracy. Gating at the lower GatedDeltaNet layers (1, 2) reaches up to 7.8× faster prefill while staying within 6% of standard prefill accuracy. FastKV can only prune at full-attention layers.
• The learned selector matters: On 2WikiMultihopQA, the trained TokenGate selector beats both a random-K selector and the FastKV selector applied at the same layer 1, across all token budgets.
• Co-adaptation matters independently: Even with random token selection, the TokenGate-trained model outperforms the LoRA-fine-tuned base model at every token budget, suggesting training redistributes information across tokens and makes the model more pruning-robust.
• Cheap adaptation: A TokenGate training step is 1.3× slower than the LoRA baseline because each step samples two token budgets, but runs converge in fewer steps on four of five tasks, so total training time is modestly reduced rather than increased.
• Kernel overhead is small: The patched FlashAttention-2 call with a per-key bias costs 1.0–1.1× a plain fused call; it supports head dimensions 128 (Qwen3) and 256 (Qwen3.5).

Methodology in Plain English

The model is split into three phases. Layers below a chosen gate layer run normally over the entire prompt. At the gate layer, a small network gives each "prunable" token an importance score; instructions, questions, or designated answer tokens can be marked non-prunable and are always kept. The top-scoring K prunable tokens, plus all non-prunable tokens, continue upward; everything else stops there. Because the hard keep/drop decision is not differentiable, the importance scores are also fed into the attention computation of every upper layer as an extra additive term on the attention logits — the same idea as a Mixture-of-Experts router weighting its selected experts, but re-purposed to weight tokens instead of parameters. Two learnable scalars per layer and head (a gain and a temperature) keep this term well-scaled and keep the softened attention comparable to full-context attention even when many keys have been removed.

For training, the authors fine-tune existing checkpoints with LoRA plus the gate parameters using ordinary next-token prediction — no auxiliary loss. Each step computes one full-budget forward (so every token receives a gradient and full-context accuracy is anchored) and one forward at a randomly sampled smaller budget drawn log-uniformly between 1 and half the prunable tokens, then averages the two losses; the two share the pre-gate compute graph, which is why the extra cost is limited to the pruned forward. For serving, they modified FlashAttention-2 to accept a per-key, per-head bias added before the online softmax, gathered the lower layers' KV cache to the kept tokens on a separate stream, and captured the pruned forward pass as CUDA graphs per padded length.

Evaluation covers five task families: RepoBench-P next-line code completion (500 samples, mean 11.0k tokens, up to 39k, edit similarity), 2WikiMultihopQA multi-hop QA (200 samples, mean 7.5k tokens, token-level F1), MultiNews multi-document summarization (200 samples, mean 2.7k tokens, up to 14k, ROUGE-L), BioDEX drug-reaction extraction (500 samples between 2k and 39k tokens, mean 4.5k, set F1), and needle-in-a-haystack (NIAH 32K). Training sets were disjoint from evaluation (e.g., 14.9k RepoBench v1.1 Python/Java cross-file samples, 20k rebuilt 2WikiMultihopQA questions, 44.8k MultiNews samples). Baselines are FastKV, SlimInfer, DASH, GemFilter, FlexPrefill, and MInference, all run on LoRA-tuned base models (r=16) trained to convergence on 4 A100 GPUs, matching TokenGate's LoRA configuration. Absolute accuracy values per dataset appear only in figures, not in the text.

Why This Matters

The paper targets prefill — the phase that determines how long a user waits before the first generated token and how much GPU memory a request occupies, which in turn caps how many concurrent requests a server can hold. Its distinguishing claim is that the accuracy/efficiency trade-off of token pruning can be bought with a LoRA-sized training budget rather than pretraining, so an organization can adapt a checkpoint it already has instead of adopting a new architecture.

Real-world applications the paper motivates or evaluates:

• Biomedical literature extraction — pulling reported adverse drug reactions from full-text papers (BioDEX), where prompt tokens run from 2k to 39k.
• Coding agents and repository-level completion — RepoBench-P prompts of up to 39k tokens containing cross-file context.
• Deep research and document assistants — multi-hop QA and multi-document summarization, where the documents arrive at request time and their KV cache cannot be precomputed.
• Retrieval and agent systems feeding long, request-time context — the paper frames this as the motivating setting: an instruction plus documents supplied or retrieved for that request.

Industry relevance is direct: lower TTFT and smaller peak prefill memory translate into serving cost and concurrency, and the "drop-in, LoRA-scale adaptation" framing is aimed at teams that cannot retrain foundation models but can afford parameter-efficient fine-tuning. The demonstrated extension to a hybrid-attention model (Qwen3.5-2B) matters because these architectures are increasingly common and, per the paper, attention-score-based pruning heuristics like FastKV can only be applied to their full-attention layers.

Future Directions

• Multiple gate modules. The paper states the design generalizes to several gates at different layers, each selecting from what lower layers kept and contributing to upper-layer token-gated attention, but the reported experiments use a single gate; how gates interact across depth is left open.
• Faster kernels for hybrid layers. For Qwen3.5's Gated DeltaNet layers the authors train via FlexAttention with the bias added as a score modification and note "more optimized kernels are possible" — closing that gap is an explicit item.
• Tighter accuracy trade-offs at aggressive pruning. The result that gating at layers 1–2 yields up to 7.8× faster prefill while staying within 6% of full-context accuracy raises the question of how much of that remaining gap can be recovered.
• Composition with decode-side optimizations. The paper positions KV-cache eviction and merging methods (and early exit / per-token routing for decoded tokens) as complementary rather than competing, which suggests combining prefill pruning like TokenGate with decode-time cache compression, and optionally pruning the pre-gate layers' cache as the appendix mentions.

Target Audience

Researchers and engineers working on LLM inference efficiency and long-context serving; systems builders who maintain model-serving stacks and care about TTFT, peak memory, and concurrency; practitioners who want to adapt an existing checkpoint with parameter-efficient fine-tuning rather than pretrain a new architecture; and readers interested in cross-pollinating Mixture-of-Experts routing ideas into token-level sparsity. Prerequisite comfort with attention math and KV-cache mechanics is needed for the method sections.

Authors’ 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.