学会“扫读”的大模型：TokenGate 把长上下文预填充提速最高 11.3 倍

引言

长上下文是智能体和 RAG 的标配：每次调用都要重新塞进系统提示、工具定义、聊天历史和检索文档。真正贵的往往不是“生成”，而是生成之前的预填充（prefill）。斯坦福与 UC Berkeley 的研究者（Liana Patel、Rishabh Ranjan、Ion Stoica、Matei Zaharia、Carlos Guestrin）发布的 TokenGate，让模型像人一样“扫读”，只精读重要的部分。

核心原理：在中间层“闸门”筛掉大部分 token

Transformer 让每个 token 通过每一层，注意力计算随长度平方增长。TokenGate 在第 L 层插入一个小型 MLP 门控：

1. L 层以下照常处理全部 N 个 token；
2. 门控给每个 token 打分，top-K 保留得分最高的 K 个；
3. L 层以上只处理这 K 个 token，被剪掉的不再产生 KV 缓存。

只有提示中的“文档”部分可被剪枝，问题本身不会被剪。

技术难点：top-K 不可导

论文借用 MoE 路由器的技巧，提出 token-gated attention：把保留 token 的分数作为偏置加到上层注意力 logits 上，给打分器一条梯度通路，并配合可学习的每头增益和温度。训练用普通下一 token 损失，每步平均“全上下文损失”与“随机预算损失”，同一个门控可适配不同保留比例。整体是“小门控 + LoRA（r=16）”，无需从头预训练。

效果与边界

• 相比完整预填充，TTFT 加速 1.4–11.3 倍，峰值内存降低 1.5–3.4 倍，报告称无精度损失；
• 相比次优剪枝基线，TTFT 最高快 6.4 倍；
• 2WikiMultihopQA：达到全上下文 95% 精度时快 3.0 倍、省 2.7 倍内存；
• 16k 训练的门控外推到 32k 为 12.8 倍，64k 与 128k 为 13.0 倍；
• 在混合注意力的 Qwen3.5-2B 上，NIAH 32k 精度持平、快 3.8 倍；FastKV 只能剪注意力层。

需冷静看待：实验用的是 Qwen3-1.7B、Qwen3.5-2B 等小模型，上下文最长 128k，单张 A100、batch size 为 1；最好的 11.3 倍来自“大海捞针”任务，部分任务只有 1.4 倍。这是学术结果，不是生产数据，且每个任务需单独训练门控。

为什么重要

对反复重发长提示的智能体，预填充是成本和延迟大头。TokenGate 表明“选择性计算”不只在 MoE 的专家维度，也可以发生在序列维度。若能在 vLLM、SGLang 等真实服务栈和更大模型上复现，“长上下文”与“检索”的成本平衡点会后移。代码与检查点以 Apache-2.0 开源。

结语

模型变大之外，“少算”正成为另一条主线。下一步该看的，是它在大模型和真实推理引擎中的表现。

参考来源

• 论文：TokenGate（arXiv 2610.26754）
• MarkTechPost 解读
• AISeng Tech 日报
• GitHub：TokenGate