TokenGate adds a trainable token-pruning module to pretrained LLMs, delivering faster prefill with lower memory usage while preserving or potentially improving accuracy, making it a viable alternative to existing pruning methods across model scales and contexts.

Researchers from Stanford and UC Berkeley demonstrate TokenGate as a token-pruning module that accelerates long-context prefill without sacrificing accuracy.

Deployments require Python 3.10+ and CUDA PyTorch 2.9+, with optional FlashAttention-2 key-bias kernel to boost performance; there are checkpoints for multiple tasks on Qwen models and at least one NIAH gate for Qwen3-1.7B and Qwen3.5-2B.

Benchmark results on Qwen3-1.7B show 1.4x to 11.3x faster time-to-first-token and 1.5x to 3.4x lower peak memory across five tasks, with the best results reaching 11.3x TTFT and 3.42x memory savings at 16k context caps.

Training uses a differentiable-like path via a Mixture-of-Experts router, blending gate scores with attention logits to enable gradient flow for end-to-end improvement.

Code, checkpoints, and usage details are available on the GitHub repository and Hugging Face page under the Apache-2.0 license, with benchmarks referenced (RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX, NIAH).

Compared with FastKV and other baselines, TokenGate achieves up to 6.4x faster TTFT on certain tasks and generally outperforms competing pruning methods in these tests.

The approach uses a small MLP gate and LoRA (r=16) to score tokens and selectively keep the top-K tokens at a chosen gate layer, enabling pruned forward passes above the gate.

TokenGate also supports hybrid attention models like Qwen3.5-2B, where gating across early layers yields substantial speedups while maintaining accuracy, unlike FastKV which prunes only attention layers.

Performance scales with context length, with TTFT gains up to 13.0x and memory savings around 3.43x at 32k, 64k, and 128k context lengths.