斯坦福大学与加州大学伯克利分校的研究人员发布了 TokenGate，这是一个面向长上下文预填充（prefill）的模块。TokenGate 通过可学习的 token 剪枝，让现有的大语言模型能够对输入进行“快速浏览”：只有真正重要的 token 会获得完整的逐层处理。研究团队报告称，在不损失准确率的情况下，首个 token 生成时间（TTFT）可提升 1.4 倍至 11.3 倍，预填充阶段的峰值内存下降 1.5 倍至 3.4 倍。代码与检查点已在 GitHub 上以 Apache-2.0 许可证发布。
TL;DR（要点速览）
MLP 门控：在 Qwen3-1.7B 或 Qwen3.5-2B 上加入一个小型 MLP 门控模块和 LoRA（r=16），测试上下文长度最高达 128k token。
运行环境：论文中的计时实验使用 1 块 A100、批大小为 1；训练使用 4 块 A100。
性能表现：在 5 项任务上相对 6 个基线方法取得帕累托最优。
最佳情况：NIAH 16k 任务上 TTFT 快 11.3 倍、内存降低 3.42 倍；在 128k 上下文中 TTFT 增益可达 13.0 倍。
最差情况：在部分任务上加速比降至 1.4 倍。
TokenGate 是什么？
TokenGate 是一个可直接嵌入（drop-in）的组件，为预训练大语言模型增加可训练的 token 剪枝能力。标准 Transformer 会把每一个 token 送入每一层，而注意力开销随序列长度呈二次增长。免训练的剪枝方法成本低，但在高剪枝率下会损失准确率；而修改模型架构通常需要从头预训练。
TokenGate 的可学习 token 剪枝如何工作？
前向传播被拆分为三个阶段：
标准前向：门控层 L 以下的层照常处理全部 N 个 token。
门控模块：一个 MLP 根据第 L 层注意力后的残差为每个 token 打分，top-K 选择器保留得分最高的 token。
剪枝后前向：L 层以上的层只处理这 K 个 token，且不会为被剪掉的 token 创建 KV 条目。
只有提示词中的文档部分可被剪枝，问题本身永远不会被剪枝。
什么是 token 门控注意力？
Top-K 选择不可微，这会阻断端到端训练。TokenGate 借用了混合专家（MoE）路由器的技巧：每个被保留 token 的得分被转化为门控层以上注意力 logits 的偏置，从而为评分器提供梯度回传路径。一个可学习的逐头增益和可学习温度进一步对此进行细化，温度用于校正每个查询所看到的键变少的问题。训练使用普通的下一 token 预测损失，每一步对全上下文损失和随机预算损失取平均，因此一个训练好的门控可在多种预算下通用。
TokenGate 在基准测试上表现如何？
研究团队在 RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX 以及大海捞针（NIAH）任务上测试了 Qwen3-1.7B，基线包括 FastKV、SlimInfer、DASH、GemFilter、FlexPrefill 和 MInference，均在相同的 LoRA 骨干上进行对比。
相对完整预填充：TTFT 快 1.4 倍至 11.3 倍，内存降低 1.5 倍至 3.4 倍。
相对次优基线：TTFT 快 1.2 倍至 6.4 倍，内存最多降低 2.1 倍。
2WikiMultihopQA：在达到完整上下文准确率 95% 的情况下，预填充快 3.0 倍、内存减少 2.7 倍。
更长上下文：在 16k 上训练的门控在 32k 上达到 12.8 倍，在 64k 和 128k 上达到 13.0 倍；内存节省维持在约 3.43 倍。
在相同层上，可学习选择器优于 FastKV 的启发式方法；即便随机 K 选择在 TokenGate 骨干上得分也更高，研究团队将此归功于骨干协同适配。
它能在混合注意力模型上工作吗？
可以。Qwen3.5-2B 混合了 Gated DeltaNet 层与全注意力层。在 NIAH 32k 上，第 3 层使用 TokenGate 达到与完整预填充相当的准确率，运行速度比完整预填充快 3.8 倍，比 FastKV 快 2.1 倍。在第 1 或第 2 层进行门控可达到 7.8 倍加速，准确率差距在 6% 以内。FastKV 只能剪枝注意力层。
TokenGate 与 FastKV、SlimInfer、GemFilter、DASH 的对比
特性对比：
Token 选择：TokenGate 使用可学习 MLP 门控；FastKV 基于最后查询的注意力；SlimInfer 采用块注意力、多阶段；GemFilter 在过滤层使用最后查询注意力；DASH 基于隐状态稳定性。
是否需要训练：TokenGate 需要（门控 + LoRA r=16）；其余四个均不需要。
TokenGate 的 Qwen3-1.7B 运行中的剪枝层：113（默认）；FastKV 为 8、17、25；GemFilter 为 12；DASH 为 11。
是否在 Gated DeltaNet 层剪枝：TokenGate 可以；FastKV 不可以（仅注意力层）；SlimInfer、GemFilter、DASH 未披露。
自报最佳预填充增益（各自设置下）：TokenGate 为 11.3 倍 TTFT、3.4 倍内存；FastKV 为 1.82 倍预填充；SlimInfer 为 2.53 倍 TTFT（LLaMA3.1-8B，RTX 4090）；GemFilter 相比 SOTA 快 2.4 倍、GPU 内存少 30%；DASH 未披露。
代码：TokenGate、FastKV、SlimInfer、GemFilter 均已在 GitHub 上，DASH 仅公布。
注意：各项自报增益使用不同实验设置，不能直接横向比较。剪枝层信息来自 TokenGate 论文附录 C.2。
如何部署 TokenGate？
GitHub 仓库需要 Python 3.10+ 和 CUDA PyTorch 2.9+。可选的 FlashAttention-2 key-bias 内核可加速稠密 Qwen3，开销为普通调用的 1.0 倍至 1.1 倍；不使用时推理回退到 FlexAttention。
使用 {state} 占位符标记可剪枝的文档部分。generate 接受 max_kept_prefill_tokens 预算参数。检查点覆盖 Qwen3-1.7B 上的全部 5 项任务，另有一个 Qwen3.5-2B 的 NIAH 门控。示例：liana313/tokengate-niah-qwen3-1.7b。
训练方式与 TRL/PEFT 配方一致，每一步比普通 LoRA 慢 1.3 倍；5 项任务中有 4 项收敛更快，平均训练时间为 1.8 个 A100 小时，而普通 LoRA 为 2.2 个 A100 小时。
关键要点
TokenGate 通过小型门控和 LoRA 为现有大语言模型加入可学习 token 剪枝。
预填充速度提升 1.4 倍至 11.3 倍，峰值内存降低 1.5 倍至 3.4 倍。
token 门控注意力以 MoE 路由器的方式使 top-K 选择可训练。
在 TTFT 上比次优剪枝基线最高快 6.4 倍。
它可在混合架构的 Qwen3.5-2B 上工作，而 FastKV 无法剪枝 DeltaNet 层。
可查阅论文、GitHub 仓库和 Hugging Face 论文页面。