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

Liana Patel; Rishabh Ranjan; Ion Stoica; Matei Zaharia; Carlos Guestrin

cs.CLOct 19, 2026

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

Authors: Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia, Carlos Guestrin

Organizations: Stanford University · UC Berkeley

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.

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

g_{j}

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\times

faster and with

2.7\times

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.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Base model	Qwen3-1.7B, bf16
LoRA	$r{=}16$ , $\alpha{=}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.

Dataset	SFT	Ours	Ratio $\downarrow$	SFT	Ours	Ratio $\downarrow$
	Full train time			Per-step time
Repo	1.7h	1.1h	$0.6\times$	0.50s	0.64s	$1.3\times$
2Wiki	3.8h	4.0h	$1.1\times$	0.34s	0.43s	$1.3\times$
MNews	1.3h	1.8h	$1.4\times$	0.21s	0.27s	$1.3\times$
BioDEX	1.2h	1.0h	$0.8\times$	0.23s	0.30s	$1.3\times$
NIAH	2.8h	1.0h	$0.4\times$	0.99s	1.24s	$1.3\times$

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\times$	$12.8\times$	$13.0\times$	$13.0\times$
Peak memory reduction	$3.42\times$	$3.43\times$	$3.43\times$	$3.43\times$

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.

Token Pruning

Token pruning removes redundant tokens from transformer inputs or intermediate representations to reduce computation, memory use, and inference time.

See all 54 papers

Stability Implies Redundancy: Delta Attention Selective Halting for Efficient Long-Context Prefilling
Apr 20, 2026Yujie Chen, Tailai Chen, Yifeng Gao +4Self-Attention Long-Context Language Model Inference
Shanghai Jiao Tong University · University of California, San Diego · Carnegie Mellon University
CLAA: Cross-Layer Attention Aggregation for Accelerating LLM Prefill
Feb 17, 2026Bradley McDanel, Steven Li, Harshit KhaitanCross-Layer Attention LLM Inference Acceleration
Franklin and Marshall College · Meta Reality Labs
SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
Jun 22, 2026Huzama Ahmad, Se-Young YunSparse Attention Block-Sparse Attention
KAIST

Apr 20, 2026J/K move · Enter open · S save

Stability Implies Redundancy: Delta Attention Selective Halting for Efficient Long-Context Prefilling

Yujie Chen, Tailai Chen, Yifeng Gao, Zoe Wanying He+3

Shanghai Jiao Tong University · University of California, San Diego · Carnegie Mellon University

Self-Attention Long-Context Language Model Inference Efficient VLM Inference LLM Inference Acceleration Attention Mechanisms

Open paper AskPDF arXiv

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.

Abstract

Token Pruning

Explore similar work

Stability Implies Redundancy: Delta Attention Selective Halting for Efficient Long-Context Prefilling

CLAA: Cross-Layer Attention Aggregation for Accelerating LLM Prefill

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers