TokenGate：可学习的 token 剪枝让长上下文预填充更快更省显存
TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning
一篇很实用的推理优化论文：给 Qwen3 加个小模块做 token 剪枝，TTFT 最快快 11 倍，代码开源可以直接试。
TokenGate 是一个即插即用的小组件，可以在已训练好的模型上做端到端可学习的 token 剪枝。为了让不可导的 top-K 操作能参与训练，作者提出 token-gated attention，借鉴 Mixture-of-Experts 的思路用 token 重要性分数重新加权注意力 logits。在多种任务上，TokenGate 把首 token 生成时间（TTFT）加快 1.4-11.3 倍，峰值显存降低 1.5-3.4 倍，且不损失精度。相比此前最好的 token 剪枝方法，预填充最快提升 6.4 倍、峰值显存降低 2 倍。该方法兼容 Qwen3 这类标准 transformer 和 Qwen3.5 这类混合注意力模型，适配成本接近常规参数高效微调。
TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning
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.