据MarkTechPost 10月25日报道，斯坦福大学和加州大学伯克利分校研究人员发布TokenGate，一种面向长上下文预填充的模块。研究团队称，该模块让现有大模型在输入中“略读”，只对重要Token进行完整处理，首Token时延（TTFT）提升1.4倍至11.3倍，且无精度损失，峰值预填充内存降低1.5倍至3.4倍。代码和检查点已在GitHub以Apache-2.0许可发布。
TokenGate在Qwen3-1.7B或Qwen3.5-2B上加入小型MLP门控和LoRA（r=16），测试上下文长度最高128k。其前向过程分为三阶段：门控层以下按常规处理全部N个Token；门控模块根据第L层注意力后的残差为每个Token打分，并由top-K选择器保留高分Token；门控层以上只处理被保留的K个Token，且不为被剪枝Token创建KV条目。提示词中只有文档部分可被剪枝，问题本身从不被剪枝。
由于top-K选择不可微，TokenGate借用混合专家（MoE）路由器的做法，把每个保留Token的分数作为门控层以上注意力logits的偏置，为评分器提供梯度路径，并加入可学习的每头增益和温度参数。训练使用常规下一Token损失，每步平均全上下文损失和随机预算损失，使一个训练好的门控可跨预算使用。论文计时使用单张A100、批大小为1，训练使用4张A100。
研究团队在RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX和NIAH上测试Qwen3-1.7B，对比FastKV、SlimInfer、DASH、GemFilter、FlexPrefill和MInference等基线。与完整预填充相比，TTFT快1.4倍至11.3倍，内存低1.5倍至3.4倍；与次优基线相比，TTFT快1.2倍至6.4倍，内存最多低2.1倍。最佳结果出现在NIAH 16k：TTFT快11.3倍，内存低3.42倍；在128k上下文，TTFT增益达到13.0倍，但部分任务的最差加速降至1.4倍。2WikiMultihopQA在保持完整上下文95%精度时，预填充快3.0倍，内存少2.7倍。
在更长上下文上，一个在16k训练的门控在32k达到12.8倍，在64k和128k达到13.0倍，内存节省保持在3.43倍附近。研究团队称，在相同层上，学习到的选择器优于FastKV的启发式方法。对混合注意力模型Qwen3.5-2B，在NIAH 32k上，第3层门控与完整预填充精度相当，运行速度快3.8倍，比FastKV快2.1倍；在第1或第2层门控达到7.8倍加速，精度差距在6%以内。FastKV只能剪枝注意力层。
部署方面，GitHub仓库要求Python 3.10+和CUDA PyTorch 2.9+；可选FlashAttention-2 key-bias内核加速密集Qwen3，开销为普通调用的1.0倍至1.1倍，否则回退到FlexAttention。{state}占位符用于标记可剪枝文档，generate接收max_kept_prefill_tokens预算。检查点覆盖Qwen3-1.7B上的5个任务，以及Qwen3.5-2B的NIAH门控。训练流程类似TRL/PEFT方案，每步比普通LoRA慢1.3倍，5个任务中有4个收敛更快，平均训练时间1.8个A100小时，对比2.2个A100小时。研究团队自报的加速结果使用不同设置，不能直接比较。