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× faster time-to-first token (TTFT) and 1.5-3.4× 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× faster prefill and 2× 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.
Figures & tables
Figure 1: Overview of TokenGate . (a) Layers up to L prefill all N context tokens as usual; at layer L the gate module keeps K of them, and every layer above L processes only those K tokens. (b) The gate module scores each hidden state with a small scoring network, and a top- K selector forwards the kept hidden states together with their importance scores. (c) Token-gated attention multiplies each attention weight by the gate value gj of its key, which enables end-to-end training. (d) On 2WikiMultihopQA question answering, with Qwen3-1.7B, TokenGate pushes the Pareto frontier of accuracy versus time-to-first-token and peak prefill memory; at 95% of the full-context accuracy it prefills 3.0× faster and with 2.7× less memory, outperforming the best training-free token pruning methods, like FastKV ( Jo et al., 2025 ) and SlimInfer ( Long et al., 2025 ) .
Figure 2: Visualization of TokenGate applied on sample inputs from the 2WikiMultihopQA dataset (left) and MultiNews dataset (right), with K=256 kept tokens, highlighted in red (darker denotes higher token importance score by the gate). Detailed visualizations of real samples from all datasets are in Appendix E .
Figure 3: Accuracy vs. time-to-first-token (top) and accuracy vs. peak GPU prefill memory (bottom) with Qwen3-1.7B. Points that incur significant overhead over standard prefill at full context (grey cross), or significantly poor accuracy, are excluded; see Appendix D.2 for full plots.
Figure 4
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
Value
Base model
Qwen3-1.7B, bf16
LoRA
r=16 , α=32 on query and value projections, all layers
Optimizer
AdamW; Muon ( Jordan et al., 2024 ) for the gate matrices ( TokenGate only)
Learning rate
10−4 , linear warmup over 1000 steps, then constant
Weight decay
0.1
Gradient clipping
global norm 1.0
Appendix
Table 1: Training hyperparameters for TokenGate and the SFT baseline, shared across all five datasets.
Full train time
Per-step time
Dataset
SFT
Ours
Ratio ↓
SFT
Ours
Ratio ↓
Repo
1.7h
1.1h
0.6×
0.50s
0.64s
1.3×
2Wiki
3.8h
4.0h
1.1×
0.34s
0.43s
1.3×
MNews
1.3h
1.8h
1.4×
0.21s
0.27s
1.3×
BioDEX
1.2h
1.0h
0.8×
0.23s
0.30s
1.3×
NIAH
2.8h
1.0h
0.4×
0.99s
1.24s
1.3×
Appendix
Table 2: Training time for SFT baseline and model trained with TokenGate for Qwen3-1.7B on one A100. Ratio shows Ours / SFT.
Figure 6: Main results (Figure 3 ) without x or y limits: accuracy vs. time-to-first-token (top) and accuracy vs. peak GPU prefill memory (bottom) with Qwen3-1.7B, all points shown.
NIAH 16k
NIAH 32k
NIAH 64k
NIAH 128k
TTFT speedup
11.3×
12.8×
13.0×
13.0×
Peak memory reduction
3.42×
3.43×
3.43×
3.43×
Appendix
Table 3: TokenGate TTFT speedup and peak prefill memory reduction at no accuracy degradation over full context for the NIAH-16k, 32k, 64k and 128k test sets with Qwen3-1.7B.
Figure 7: Left: token-gated forward and training ablations on BioDEX (Qwen3-1.7B, gate at layer 1): F1 vs. TTFT (batch 1, A100) for the full TokenGate recipe against three single-change variants: no full-context loss (no full-context term in Eq. 10 ), no temperature scaling (Eq. 7 removed), and no head-wise scaling (raw bias, with no normalization or gain; Eq. 3 removed). Right: gate-layer ablation on BioDEX (Qwen3-1.7B): F1 vs. TTFT (batch 1, A100) with the TokenGate placed at layer 0, 1, or 2.
Dataset
What is kept
Failure mode
RepoBench-P
code syntax rather than identifiers: commas and empty argument lists “()” dominate, line breaks are avoided, and kept tokens are isolated rather than spans
wrong line, as with the full context
2WikiMultihopQA
entities and dates: names and years are over-kept and function words avoided; on the shown prompt, the answer phrase (“Film Award for …Director”)
a different award at small budgets, exact at the shown one; multi-hop answers need the largest budgets to match the full context
MultiNews
mostly function words and scattered event words (“charged”, “president”); names and places are under-kept
right events, but wrong name, occupation and place, consistent with the under-kept names
BioDEX
terms of the diagnosis and outcome (“septic shock”, “failure”) and the connectives of lists (“and”, “or”); numbers are almost never kept
misses reactions, but so does the full-context model, which bounds the score
NIAH 16k
the needle sentence as one contiguous span, plus rare subwords elsewhere (“-haired”, “Lisp”, “.html”)
none: exact at every budget
Appendix
Table 4: Summary of what TokenGate keeps and how its answers fail. Examples refer to the prompts of Figures 8 – 12 at their shown budgets; the patterns hold over the first eight prompts of each test draw.
Figure 8: Context tokens kept by TokenGate on one RepoBench-P test prompt (Appendix E ), with the reference answer and the model’s greedy answers at the shown budget and with the full context. Kept tokens are shaded by gate score (darker is higher) and pruned tokens are grey.
Figure 9: Context tokens kept by TokenGate on one 2WikiMultihopQA test prompt (Appendix E ), with the reference answer and the model’s greedy answers at the shown budget and with the full context. Kept tokens are shaded by gate score (darker is higher) and pruned tokens are grey.
Figure 10: Context tokens kept by TokenGate on one MultiNews test prompt (Appendix E ), with the reference answer and the model’s greedy answers at the shown budget and with the full context. Kept tokens are shaded by gate score (darker is higher) and pruned tokens are grey.
Figure 11: Context tokens kept by TokenGate on one BioDEX test prompt (Appendix E ), with the reference answer and the model’s greedy answers at the shown budget and with the full context. Kept tokens are shaded by gate score (darker is higher) and pruned tokens are grey.
Figure 12: Context tokens kept by TokenGate on one NIAH 16k test prompt (Appendix E ), with the reference answer and the model’s greedy answers at the shown budget and with the full context. Kept tokens are shaded by gate score (darker is higher) and pruned tokens are grey.
Prefilling computational costs pose a significant bottleneck for Large Language Models (LLMs) and Large Multimodal Models (LMMs) in long-context settings. While token pruning reduces sequence length, prior methods rely on heuristics that break compatibility with hardware-efficient kernels like FlashAttention. In this work, we observe that tokens evolve toward semantic fixing points, making further processing redundant. To this end, we introduce Delta Attention Selective Halting (DASH), a training-free policy that monitors the layer-wise update dynamics of the self-attention mechanism to selectively halt stabilized tokens. Extensive evaluation confirms that DASH generalizes across language and vision benchmarks, delivering significant prefill speedups while preserving model accuracy and hardware efficiency. Code will be released at https://github.com/verach3n/DASH.git.
Yujie Chen, Tailai Chen, Yifeng Gao +4
Shanghai Jiao Tong University · University of California, San Diego · Carnegie Mellon University
The prefill stage in long-context LLM inference remains a computational bottleneck. Recent token-ranking heuristics accelerate inference by selectively processing a subset of semantically relevant tokens. However, existing methods suffer from unstable token importance estimation, often varying between layers. Evaluating token-ranking quality independently from heuristic-specific architectures is challenging. To address this, we introduce an Answer-Informed Oracle, which defines ground-truth token importance by measuring attention from generated answers back to the prompt. This oracle reveals that existing heuristics exhibit high variance across layers: rankings can degrade sharply at specific layers, a failure mode invisible to end-to-end benchmarks. The diagnosis suggests a simple fix: aggregate scores across layers rather than relying on any single one. We implement this as Cross-Layer Attention Aggregation (CLAA), which closes the gap to the oracle upper bound and reduces Time-to-First-Token (TTFT) by up to 39% compared to the Full KV Cache baseline.
Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.