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

2 hours ago AI闪点

一个不改动主干模型的门控模块

斯坦福大学与加州大学伯克利分校团队在 10 月 25 日公开了 TokenGate，一个可以插进现有大模型的 token 剪枝模块。思路很直接：预填充阶段的注意力开销随上下文长度平方增长，如果能提前判断哪些 token 无关紧要，只让重要的 token 走完整计算，长上下文的等待时间就能压缩。模块本体很小，一个小型 MLP 门控加一组秩为 16 的 LoRA，挂到 Qwen3-1.7B 或 Qwen3.5-2B 上，上下文最长测到 128k。

前向过程切成三段：门控层以下照常处理全部 token；门控层用 MLP 给每个 token 打分，取 top-K 保留；门控层以上只处理这 K 个 token，被剪掉的连 KV 都不生成。实际使用时只有文档部分可剪枝，问题本身永远保留。难点在 top-K 选择不可导，没法端到端训练，团队借 MoE 路由器的做法，把保留 token 的得分转成上层注意力 logits 上的偏置，配上可学习的每头增益和温度参数，让梯度回到打分器。

11.3 倍加速怎么测出来

论文在 Qwen3-1.7B 上跑了 RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX 和大海捞针五项任务，对比 FastKV、SlimInfer、DASH、GemFilter 等六个基线，都在同一套 LoRA 主干上跑。对比完整预填充，TTFT 快 1.4 到 11.3 倍，峰值预填充内存降 1.5 到 3.4 倍；只跟次优基线比仍快 1.2 到 6.4 倍。最亮眼的是 NIAH 16k，TTFT 快 11.3 倍、内存低 3.42 倍。上下文继续拉长还能再涨，16k 上训练的门控迁到 32k 时加速 12.8 倍，64k 和 128k 上都是 13.0 倍。混合注意力模型上，Qwen3.5-2B 在 NIAH 32k 上于第 3 层打门控就追平完整预填充准确率，比完整预填充快 3.8 倍、比 FastKV 快 2.1 倍。

算子层的冗余与集群层的冗余

这类改造的落脚点还是推理成本。我们此前拆解过 llm-d 1.0 的分布式推理路径，它靠 17000 节点把首 token 延迟砍掉 85%，处理的是集群侧吞吐；TokenGate 处理单请求内的冗余计算，一个在系统层，一个在算子层，恰好可叠加。成本压力还体现在词元市场上，公开服务 TPS 中位数只有 78.23，单价却涨到 13.46 元每百万 token，压缩预填充就是直接削掉这部分支出。

代价也要说清楚。剪枝层选在第 113 层，相当靠后，前 112 层的完整计算一分不少省不掉；1.4 倍是它自己报的最差情形，部分任务收益明显缩水。它还需要训练，一个门控加 LoRA 大约是普通 LoRA 的 1.3 倍耗时，平均训练时间 1.8 个 A100 小时对 2.2 个。落到生产上，还得看业务的上下文分布值不值得这次训练。论文计时用单张 A100、batch size 为 1，训练用了 4 张 A100，代码与 checkpoint 以 Apache-2.0 协议开源。

延伸阅读

AI 自学笔记 每天更新 AI 前沿动态，订阅 RSS，或按 Ctrl/⌘+D 收藏本站，明天见。

继续读

同一批更新里这几篇值得一看：

Related News