TokenGate 长上下文预填充提速 11.3 倍，斯坦福伯克利开源门控剪枝

一句话结论，TokenGate 用一个可训练的轻量门控模块，在不改动主干权重结构的前提下把长上下文预填充的首字延迟压到原来的 1/11.3，同时峰值显存降到 1/3.4，精度基本不掉，是目前公开结果里把可学习剪枝推进到端到端联合训练的第一份工作。

长上下文这条赛道上，注意力开销随序列长度平方增长是所有人都绕不开的账单。预填充阶段要把整段输入过一遍模型的每一层，随后为每个 token 写出 KV 缓存。这两步共同决定了用户按下回车之后要等多久看到第一个字，也决定了同一张卡上能同时挂多少路请求。围绕这件事，过去两年产出过一批免训练的门控剪枝方法，思路大致相同，在模型中段用注意力打分挑出最重要的 token，其余的不再往上层传。它们的优点是部署便宜，缺点是启发式打分在浅层信号太弱，一旦剪得激进，准确率就先塌了。

斯坦福大学与加州大学伯克利分校的团队在 2026 年 10 月 19 日公开了 TokenGate，把这件事往前推了一步，门控的依据换成了可训练的打分网络，跟着主干一起训。论文作者包括 Liana Patel、Rishabh Ranjan、Ion Stoica、Matei Zaharia 与 Carlos Guestrin，arXiv 编号 2610.26754，同日以 Apache 2.0 许可在 GitHub 放出代码。这篇工作交付的是一层能插进现有模型的可训练组件，落在的是推理成本这个非常具体的痛点上。

门控到底插在模型的哪个位置

TokenGate 的结构可以用三段来描述。门控层以下的各层照常处理全部 token，不做任何删减。到了门控层，一个很小的多层感知机对每个 token 的归一化残差读数打一个标量分数，分数用 RMSNorm 沿序列归一化，然后按分数取前 K 个 token 保留。被剪掉的 token 的残差流就在这一层停住，不再往上走，也就不会为它们生成任何 KV 缓存条目。门控层以上只处理保留下来的那 K 个 token。

这套设计有三个实际后果。第一，剪枝发生在模型中段而不是开头，用户的问题、已经解码出来的 token 这些可以被标记为不可剪枝的部分始终完整参与计算。第二，被剪掉的 token 根本不进入上层的 MLP 与注意力，所以省下的是实打实的算力而不只是注意力权重。第三，门控层以下的部分完全不变，意味着主干可以是一个已经训好的模型，不需要重新预训练。

论文把输入拆成请求指令与待处理文档两部分，文档可以是代码仓库、论文、病历、财报，可以是用户上传的，也可以是检索系统或智能体在运行时拉进来的。这一点决定了它的适用面，因为这类输入的 KV 缓存没法提前算好，只能在请求时现场处理。

端到端训练的关键在 token 门控注意力

top-K 选择本身不可导，梯度传不回去，这是整个方法最核心的技术障碍。TokenGate 的解法叫 token 门控注意力，思路借鉴自混合专家模型的路由器。MoE 用路由分数给专家子网络加权，在激活参数这个维度上制造稀疏；TokenGate 把重要性分数直接加权到上层注意力的 logit 上，在处理的 token 维度上制造稀疏。

具体做法分成两处修改。一是按头缩放，每个注意力头都有一个可学习的增益系数，乘在重要性分数上，再做 softmax 得到门控值。二是温度缩放，因为剪枝之后每个 query 能看到的 key 数量从原来的位置缩到了它自己排在保留列表中的名次，softmax 相对完整上下文变得更尖锐，需要一个可学习的温度做重新校准。这个温度按保留数量与原始位置的日志比缩放，灵感来自 SSMax 的做法。两者合起来的注意力权重就是原始的 query-key 点积除以头维根号，再加上缩放后的重要性分数，一起送进 softmax。

训练方式上没有额外的辅助损失，用标准的下一 token 交叉熵就行。关键技巧是双预算采样。每个训练实例先用完整上下文跑一遍得到满预算损失，再随机采一个远小于满预算的裁剪预算跑第二遍，两个损失取平均。前者给所有 token 都留了梯度，包括那些在采样中被剪掉的，后者让门控见过从温和到激进的各个剪枝比例，这样训练出来的模型换到别的 token 预算上不用重训。总训练时间与标准 LoRA 微调接近，每步因为要跑两次前向而慢 1.3 倍，但在五个任务里有四个比基线更早收敛，整体时间反而略低。

加速数据在什么口径下测出来的

实验用 Qwen3-1.7B 这个标准 Transformer 作为主干，用 Qwen3.5-2B 这个混合注意力架构作为泛化验证，两者都在四张 A100 上训练，LoRA 秩为 16。任务覆盖五个方向，代码补全用 RepoBench-P，多跳问答用 2WikiMultihopQA，多文档摘要用 MultiNews，生物医学药物反应抽取用 BioDEX，还有大海捞针 NIAH。基线选了免训练剪枝里的 FastKV、SlimInfer、DASH、GemFilter，以及稀疏注意力路线的 FlexPrefill 与 MInference，所有基线都在同一套 LoRA 主干上跑到收敛。

在 2WikiMultihopQA 上，以 Qwen3-1.7B 达到完整上下文 95% 精度为前提，TokenGate 的预填充速度快 3.0 倍、峰值显存少 2.7 倍，这个点比单看峰值倍数更能说明问题，因为它把加速和不掉点绑在了一起。相比之下，FastKV 在同样的第 1 层打门控时全面落后，论文给的原因很直接，注意力打分这种启发式在浅层能提供的信号太弱。

消融实验把优势拆成了两块。第一块是门控本身可训练，把训练好的门控换成随机选 K 个 token，准确率在所有预算下都掉；把同一层位置的 FastKV 选择器拿来对比，差距更大。第二块是协同适应，同样用随机 K 的粗选择，跑在 TokenGate 模型上的准确率仍然高于跑在普通微调基线上的准确率。这说明主干与门控一起训过之后，模型本身对剪枝变得更鲁棒，信息在 token 之间的分布被重新摊开了。

论文的可视化还暴露了一个有意思的现象，不同任务门控学到的剪枝模式完全不一样。做问答时它保留的是答案关键词，做摘要时它保留的却是短语结尾的停用词。这种差异说明门控确实在按任务学，而不是套用一套固定的重要性排序。

混合注意力架构上的表现更值得看

Qwen3.5-2B 每四层一个全注意力层，其余层用 Gated-DeltaNet。FastKV 这类方法只能在全注意力层上做剪枝，而 TokenGate 两种层都能挂。在 NIAH 32k 上，第 3 层打门控就能追平完整预填充的准确率，首字延迟只有完整预填充的 1/3.8，比最好的 FastKV 跑法还快 2.1 倍。更有意思的是把门控放到低层的 Gated-DeltaNet 上，也就是第 1、2 层，预填充最多可以快 7.8 倍，精度落在完整预填充的 6% 以内。

这个结果的实践含义是，混合注意力架构上的剪枝位点比想象中更靠前。剪枝越靠前，省下的层数越多，这个道理其实直白，但过去受限于方法只能在特定层生效而难以兑现。

工程落地要补的几块拼图

论文把 TokenGate 定位为标准微调的替代品而非补充，训练模型时同时提升精度与预填充效率。这个说法有吸引力，但落到生产环境还有几处需要自己补。

第一个是它不是即插即用的推理优化。免训练方法部署时换个模型就能用，TokenGate 需要针对目标模型重跑一次微调，主干加门控一起训，成本是实打实的训练开销。第二个是实验规模仍在小模型区间，主干选的是 1.7B 与 2B，7B 以上以及生产级模型上的收益尚未有公开验证。第三个是内核依赖，官方推荐编译一个带 key 偏置的 FlashAttention-2 定制内核，没装的话推理会自动回退到 FlexAttention 并给出警告，训练则需要显式指定后端。第四个是社区验证尚早，仓库 10 月 19 日建仓，到 10 月 26 日星标数只有 30 个左右，与那些动辄上万星的开源项目不是一个量级，第三方复现结论还没有形成。

还有一点值得提醒。TokenGate 剪掉的是 token，不是句子，剪掉的内容在残差流上被截断，后续层看不到。这与注意力稀疏或者 KV 缓存淘汰在机制上不同，它更接近于结构性的信息丢弃。对长文档摘要、代码补全这类任务，这种丢弃恰好落在冗余部分；但如果任务本身要求覆盖全文的所有细节，剪枝比例就得上小心。

它放在整条推理优化链条的哪个位置

把这篇工作放回大图里看更清楚。此前围绕推理成本的一大片工作其实在打解码阶段，办法是预填充完成后驱逐或合并 KV 缓存条目，或者做提前退出与逐 token 路由。这两类方法跟预填充优化是互补关系，管的是生成阶段的账。TokenGate 管的是生成之前的那一段，也就是用户按下回车到看到第一个字之间的等待。

这个时间段在真实产品里的权重并不低。聊天产品的首字延迟直接决定用户感知，代码助手要先把整个仓库读进去，深度研究工具要一次性吞下几十份材料。在这些场景里预填充就是主成本。相关工作里还有从更底层入手的路线，比如换 FlashAttention 或稀疏注意力机制，但那些路径依然要求每个 token 都过一遍每一层，TokenGate 的路子是把 token 数量本身降下来。

值得留意的是另一组对照数字。我们此前拆解过 llm-d 1.0 的分布式推理改造，那套做法在集群层面把首 token 延迟砍掉 85%、吞吐提升 2.9 倍；而 DeepSeek V4.1 Flash 把 KV Cache 缩小 437 倍，是在模型架构层面直接改缓存形态。TokenGate 走的是第三个维度，算子层的冗余。前者不动模型，后者重训模型，TokenGate 只在两者之间加一层可训练组件，改动面最小但需要重训。

从稀疏化路线看，它和 Recursive Flint 把 40B 模型压到只激活 2.5B 属于同一个大方向，都是让模型少算一部分。区别在于 Recursive Flint 稀疏的是参数维度，TokenGate 稀疏的是输入 token 维度，论文自己也把两者做了明确区分。混合专家模型是参数维度的稀疏，稀疏注意力是注意力内部结构的稀疏，TokenGate 则是直接决定哪些 token 不进上层。

同一个实验室此前的另一条路线是 斯坦福与 Together AI 推出的 NCM 文本生成范式，用下一概念建模把生成加速 31 倍，走的是训练目标层面的重构。TokenGate 没动训练目标，动的是前向过程里的数据流。三种做法可以叠在一起，这也是它作为 drop-in 组件的价值所在。

FAQ

TokenGate 适合什么样的场景

它针对的是长上下文预填充本身开销过高的场景，比如代码助手读整个仓库、深度研究工具一次吞几十份材料、研究型问答在长文档里做多跳检索。判断标准很简单，如果你的服务里用户等待首字的主要时间花在预填充上，并且能接受针对模型做一次专项微调，这个方法就值得试。反过来，如果瓶颈在解码阶段的生成速度，或者你无法为每个模型单独做训练，收益就会打折扣。

和 FastKV、SlimInfer 这类免训练方法怎么选

免训练方法的优势是零训练成本、换个模型就能用，适合快速验证；TokenGate 的优势是精度保持得住，在同等剪枝率下明显压过启发式方法，论文报告对次优基线最高还有 6.4 倍的首字延迟优势。实践中的常见路径是先用免训练方法把基线跑出来、确认这条业务确实值得投入，再考虑要不要付一次训练成本换更干净的曲线。

它训练出来的模型还能像原来一样用吗

可以。TokenGate 模型是一个普通的 Transformer，仍然走标准的 generate 接口，保留 2048 或者更高的上下文都可以，使用方不需要额外的运行时依赖。需要注意的是它保留了插桩能力，仓库提供了加载模型与套用对话模板的入口，预填充时可以通过参数指定保留的 token 预算，从完全不剪枝到只留一百个 token 都可以在同一个 checkpoint 上切换。

小结

TokenGate 做的事情不复杂，把一段输入里不重要的 token 在模型中段就停下来，并且让哪些 token 不重要这件事跟着任务一起学出来。11.3 倍的首字延迟与 3.4 倍的显存下降来自四个具体设计的叠加。门控放在中段而非开头，剪掉的 token 不生成 KV，重要性分数通过门控注意力进入梯度，训练时用双预算采样覆盖不同剪枝比例。

它没有解决所有问题。小模型验证、微调成本、内核依赖、社区复现不足这几项都摆在那里，生产部署前要自己算一遍账。但如果你的产品正被长上下文的预填充账单卡住，这套东西提供了一条比换模型、比改架构、比调集群都轻的路径，值得放进技术评估清单里跑一轮。

相关阅读

• DeepSeek V4.1 Flash 深度解读，KV Cache 缩小 437 倍的架构改动
• llm-d 1.0 发布，17000 节点吞吐提升 2.9 倍
• Recursive Flint 开源，40B 模型只激活 2.5B
• 斯坦福与 Together AI 推出 NCM 范式，31 倍加速