斯坦福伯克利TokenGate 预填充提速11.3倍

时间：2026年10月25日，美国加州斯坦福大学与旧金山加州大学伯克利分校；论文本身于2026年10月19日挂网 arXiv，代码与检查点同期在 GitHub 公开。 地点：斯坦福大学与加州大学伯克利分校（Stanford University and UC Berkeley），开源仓库位于 GitHub 的 tokengate-models/TokenGate，采用 Apache-2.0 许可证。 人物：论文作者为 Liana Patel、Rishabh Ranjan、Ion Stoica、Matei Zaharia、Carlos Guestrin。科技媒体 MarkTechPost 记者 Asif Razzaq 于10月25日首发报道。 事件详情：斯坦福大学与加州大学伯克利分校的研究团队发布了一套名为 TokenGate 的即插即用模块，专门用来加速大模型的超长上下文预填充阶段。预填充是模型在吐出第一个字之前必须先把整段输入读完的阶段，注意力计算量随序列长度大致呈平方增长，因此在动辄几十万到上百万 token 的上下文场景里，"首字延迟"和"预填充峰值显存"直接决定了线上服务的成本结构。 TokenGate 的思路是让模型学会"略读"。它是一个可以直接挂到已训练模型上的小型门控组件，引入可训练的 token 剪枝：前向传播被切成三段。门控层以下的各层照常处理全部 N 个 token；到门控层时，一个小型多层感知器根据每个 token 的注意力后残差给打分，top-K 选择器只保留得分最高的那些 token；门控层以上各层只处理这 K 个 token，被剪掉的 token 不会再生成对应的 KV 缓存条目。论文强调只有提示词中的文档部分可被剪枝，问题部分永远不会被剪掉。 真正的技术难点在于 top-K 选择本身不可导，常规做法会卡住端到端训练。TokenGate 的解法叫"token 门控注意力"，借鉴的是混合专家模型的路由技巧：把每个被保留 token 的得分转化为门控层之上注意力 logits 上的一个偏置项，从而给打分器留出一条梯度通路；再配合一个可学习的每注意力头增益和一个可学习温度，后者用来修正每个 query 所能看到的 key 数量减少带来的影响。训练使用普通的下一 token 损失，每一步把全上下文损失与一个随机预算损失取平均，因此一个训练好的门控可以跨不同预算复用。 实验在 Qwen3-1.7B 上进行，覆盖 RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX 和大海捞针（NIAH）五项任务，对比的基线包括 FastKV、SlimInfer、DASH、GemFilter、FlexPrefill 和 MInference，且都使用同一套 LoRA 骨干。论文计时用单张 A100、批大小为 1，训练用了 4 张 A100；实现方式是给 Qwen3-1.7B 或 Qwen3.5-2B 加一个多层感知器门控加 LoRA（秩 16）。 效果上，相比完整预填充，TokenGate 的首字延迟加速 1.4 倍至 11.3 倍，峰值预填充显存下降 1.5 倍至 3.4 倍，且没有报告到精度下降。对比此前表现最好的 token 剪枝方法，预填充最多快 6.4 倍，峰值显存最多低 2 倍。最佳案例出现在 NIAH 16k 上，首字延迟快 11.3 倍、显存低 3.42 倍；上下文拉到 128k 时首字延迟收益达到 13.0 倍，最差情况在部分任务上只有 1.4 倍。在 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 任务中，把 TokenGate 放在第 3 层可以达到与完整预填充相当的准确率，速度快 3.8 倍、比 FastKV 快 2.1 倍；放在第 1 或第 2 层则能达到 7.8 倍，但准确率会落在 6% 以内。研究还指出 FastKV 只能剪枝注意力层，这是它的结构性限制。代码仓库要求 Python 3.10 以上与 CUDA PyTorch 2.9 以上环境。 背景：推理侧的成本结构正在被重新审视。随着智能体工作流和百万级上下文窗口普及，prefill 已经取代 decode 成为主要的延迟与显存来源。GPT-6 更新了提示词缓存这类服务侧优化，PrfaaS 等预填充卸载系统也在解同一个问题，但这些路径大多停留在服务架构层面。TokenGate 选择了一条纯模型侧的路线：不换模型、不改架构，只用一次轻量微调就让已训练好的模型变成"会预算"的预填充系统。 这也解释了它与既有的免训练剪枝方法的分野。免训练剪枝便宜，但在高剪枝率下掉点明显；改造架构则通常需要从头预训练，还要换一套算子栈。TokenGate 保留预训练骨干，克隆一条打分分支，用 LoRA 做适配，改造成本接近常规的参数高效微调。团队用"专家略读长代码库"的类比来解释动机，但也明确指出这个类比不构成任意删除的许可。 影响： - 首字延迟和峰值显存这两项直接对应线上推理的算力账单，1.4 倍至 11.3 倍的加速区间意味着长文档与长代码库场景的部署成本可能出现结构性下降。 - 作为即插即用的组件而非新架构，它不要求团队更换模型或重训基座，只需一次 LoRA 级别的轻量适配，迁移成本远低于换模型。 - 相比此前最好的 token 剪枝方法最高 6.4 倍的加速与 2 倍显存收益，TokenGate 把"可选剪枝"的水位整体抬高了一档，方法论上的意义可能大于单点数字。 - 11.3 倍与 1.4 倍之间的巨大落差说明收益高度依赖任务形态，剪枝率选错就会吃掉大部分收益，落地前必须逐任务标定而非套用宣传数字。 总结：TokenGate 表面上是一次推理优化，实际上是一次关于"上下文该由谁来决定"的重新分工——把原本由人工设定的注意力预算，交回给模型自己学出来的那道门。这个思路最值得借鉴的地方不在 11.3 这个数字，而在于它用 token 门控注意力给不可导的 top-K 硬接了一条梯度通路，让"选择"这件事第一次可以被端到端地训练。团队还发现，在 TokenGate 骨干上随机 K 的表现都比原骨干上的启发式规则更好，这意味着剪枝方法的效果高度依赖骨干是否被协同适配过——对打算自建剪枝方案的团队来说，这条经验可能比论文本身更有实用价值。 参考来源： https://arxiv.org/abs/2610.26754 https://github.com/tokengate-models/TokenGate https://ainewsfeed.tech/articles/stanford-and-uc-berkeley-researchers-release-tokengate-for-faster-long-context-llm-prefill-36294-en-us-v1 https://www.dogely.com/ai-agent/9620.html https://megaoneai.com/benchmarks/tokengate-stanford-berkeley-llm-prefill-speedup https://newstodai.vercel.app/en/news/1a4147d6e8f37270d8c7f734fa4fc58e9e1e12a0 https://thetechbriefs.com/researchers-from-stanford-and-uc-berkeley-introduce-tokengate-token-gated-attention-makes-long-context-prefill-up-to-11-3x-faster