斯坦福与 UC Berkeley 提出 TokenGate：token 门控注意力让长上下文预填充最高快 11.3 倍
Researchers from Stanford and UC Berkeley Introduce TokenGate: Token-Gated Attention Makes Long-Context Prefill Up to 11.3x Faster - MarkTechPost
斯坦福和 UC Berkeley 研究者发布 TokenGate，一个面向长上下文预填充的可学习 token 剪枝模块，团队报告首 token 时延（TTFT）提升 1.4 倍至 11.3 倍且无精度损失，峰值预填充显存降低 1.5 倍至 3.4 倍。方法是在 Qwen3-1.7B 或 Qwen3.5-2B 上加一个小型 MLP 门控和 LoRA（r=16），用 MoE 路由思路让 top-K 选择可端到端训练，只有文档部分可剪枝、问题部分不剪。代码与 checkpoint 以 Apache-2.0 在 GitHub 开源，测试上下文最长 128k token，NIAH 16k 上达到 11.3 倍加速和 3.42 倍显存下降。
正文 · AI 翻译
在 Google 上添加为首选来源
斯坦福和 UC Berkeley 的研究者发布了 TokenGate，一个用于长上下文预填充的模块。TokenGate 的可学习 token 剪枝让现有 LLM 能够略读其输入。只有重要的 token 才会得到完整处理。研究团队报告称，首 token 时延（TTFT）提升 1.4 倍至 11.3 倍，且无精度损失。峰值预填充显存下降 1.5 倍至 3.4 倍。代码和 checkpoint 以 Apache-2.0 协议在 GitHub 上开源。
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 是一个即插即用组件，为预训练 LLM 添加可训练的 token 剪枝。Transformer 会将每个 token 送入每一层。注意力成本随长度呈二次增长。免训练剪枝成本低，但在高剪枝率下会损失精度。架构改动通常需要从头预训练。
TEXTTABLELAYOUT你的 PDF 解析器真的解析对了吗？omni_parse_bench · 对文本、表格和布局评分 · 开源在 GitHub 上加星→赞助
TokenGate 的可学习 token 剪枝如何工作？
前向传播分为 3 个阶段：
标准前向：门控层 L 以下的层照常处理全部 N 个 token。
门控模块：一个 MLP 根据第 L 层注意力后的残差为每个 token 打分。top-K 选择器保留得分最高的 token。
剪枝前向：L 以上的层只处理这 K 个 token，且从不为被剪枝的 token 创建 KV 条目。
只有提示中的文档部分可被剪枝。问题部分永远不会被剪枝。
什么是 token 门控注意力？
top-K 选择不可微，这会阻碍端到端训练。TokenGate 借用了混合专家（MoE）路由器的技巧。每个保留 token 的得分会成为门控之上注意力 logits 的偏置。该偏置为梯度提供了一条回到打分器的路径。
可学习的逐头增益和可学习的温度对此进行细化。温度用于校正每个查询看到的键更少的问题。训练使用普通的下一 token 损失。每一步对全上下文损失和随机预算损失取平均。因此 1 个训练好的门控可适用于多种预算。
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 的启发式方法。即使在 TokenGate 骨干上，随机 K 选择得分也更高。研究团队将其归功于骨干协同适应。
它在混合注意力模型上有效吗？
有效，Qwen3.5-2B 混合了 Gated DeltaNet 和全注意力层。在 NIAH 32k 上，第 3 层的 TokenGate 达到了与完整预填充相同的精度。它比完整预填充快 3.8 倍，比 FastKV 快 2.1 倍。在第 1 或第 2 层门控达到 7.8 倍，精度差距在 6% 以内。FastKV 只能剪枝注意力层。
特性TokenGateFastKVSlimInferGemFilterDASH
Token 选择可学习 MLP 门控基于最后查询的注意力块注意力、多阶段过滤层上的最后查询注意力隐状态稳定性
是否需要训练是（门控 + LoRA r=16）否否否否
TokenGate 的 Qwen3-1.7B 运行中的剪枝层113（默认）8、17、251211
是否在 Gated DeltaNet 层剪枝是否（仅注意力层）未披露未披露未披露
最佳自报预填充增益（自有设置）11.3 倍 TTFT、3.4 倍显存1.82 倍预填充2.53 倍 TTFT（LLaMA3.1-8B，RTX 4090）相比 SOTA 提速 2.4 倍、GPU 显存减少 30%未披露
代码GitHubGitHubGitHubGitHub已公布
自报增益使用不同设置，因此不能直接比较。剪枝层来自 TokenGate 的附录 C.2。
如何部署 TokenGate？
GitHub 仓库需要 Python 3.10+ 和 CUDA PyTorch 2.9+。可选的 FlashAttention-2 键偏置内核可加速稠密 Qwen3。其开销为普通调用的 1.0 倍至 1.1 倍。若不使用，推理会回退到 FlexAttention。
{state} 占位符标记可剪枝的文档。generate 接受 max_kept_prefill_tokens 预算。Checkpoint 覆盖 Qwen3-1.7B 上的全部 5 个任务，外加一个 Qwen3.5-2B NIAH 门控。示例：liana313/tokengate-niah-qwen3-1.7b。
训练仿照 TRL/PEFT 配方。每一步比普通 LoRA 慢 1.3 倍。在 5 个任务中的 4 个上运行收敛更快。平均训练时间为 1.8 对 2.2 A100 小时。
关键要点
TokenGate 通过一个小型门控和 LoRA 为现有 LLM 添加可学习 token 剪枝。
预填充快 1.4 倍至 11.3 倍，峰值显存低 1.5 倍至 3.4 倍。
Token 门控注意力使 top-K 选择可训练，采用 MoE 路由器风格。
在 TTFT 上比次优剪枝基线最多快 6.4 倍。
它在混合 Qwen3.5-2B 上有效，而 FastKV 无法剪枝 DeltaNet 层。
Asif Razzaq网站 | + 文章作者简介 ⮌
Asif RazzaqNVIDIA 发布 Nemotron 3.5 Bolt E4B 4GB：一个面向本地语音智能体、游戏角色和机器人的开放小型语言模型
Asif RazzaqMoonshot AI 发布 Kimi K3.1 Pro：一个 2.8T 总参数/104B 激活参数（MoE）模型，支持 1M 上下文
$uv pip install omni-parse-bench✓文本✓表格✓布局→ 按页评分在 GitHub 上加星→赞助
相关文章更多来自作者
如何运行 Underdog Saluki 27B：2 比特 Qwen3.8-27B 工具调用指南
认识 Anthropic 的 sandbox-runtime：面向 AI 智能体的操作系统级沙箱，无需容器