Prefix Sliding for efficient test-time scaling
Niklas Muennighoff, Zhengyang Wang, Zeyi Chen, Weijia Shi, Binyuan Hui, John Yang, Dapeng Jiang, Mika Senghaas, Fares Obeid, Johannes Hagemann, Sami Jaghouar, Ludwig Schmidt, Percy Liang, Jason Wei, Andrew Y. Ng, Luke Zettlemoyer, Yejin Choi, Mike Lewis
cs.CL, cs.AI, cs.LG
2026-08-27
在 Qwen3-1.7B 上，Prefix Sliding 只保留提示前缀和最近几千个 token，免训练即可把长推理加速约 3 倍；窗口 8192 时 AIME25 达到 35.8，略高于全注意力的 34.2。
让模型在测试时多想一会儿，是现在最常用的 test-time scaling。卡点在全注意力：整段推理轨迹都要进 KV cache，每多吐一个 token，注意力代价就线性往上爬。难问题动辄几万 token，显存和延迟很快顶死。更长的上下文还会把模型带偏，旧 token 干扰、循环复读、前面写过的结论被淹没。
一个具体的画面：算「((42 + 84) × 4) - 5」时，42 + 84 一旦算完，那段推演过程就没价值了，下一步只用得上中间结果。这篇在 Qwen3-1.7B 的 AIME25 轨迹上画了注意力分布：系统指令、题目、<think> 分隔符，再加上最近大约 1000 个 token，吃掉了大部分注意力。中间那段推理几乎没人看。
Prefix Sliding 生成时只留两类 token。
中间那些已经过期的推理直接丢掉。前缀 100 token、窗口 4096 时，内存里最多 4196 个 token。模型已经想了百万还是十亿步，每一步的代价都一样。
位置编码默认 Continue PE：被滑出去的位置编号不重排，KV cache 里已经乘过 RoPE 的向量直接复用。Reset PE 要把同一批 token 重新赋位置再算一遍，更贵。附录在 AIME25 上两者差距不大，实现就选了便宜那条。
训练侧用截断反传。滑动窗跨层的理论感受野是 W×L，实践里大约只有 1.5×W。一条 10 万 token 的轨迹、窗口 2048 时，训练器只拿最后 8192 个 token 当上下文，损失只算最后 2048 个。7B 的异步 RL 上核过：窗口 8192、乘数 4，AIME24 能跟全量反传的全注意力打平。FlashAttention 核做了两级过滤：tile 内部对不合法的 (q, k) 做 mask，tile 之间直接跳过前缀和窗口之外的块。
主实验是 Qwen3-1.7B，avg@64，温度 0.6。免训练、窗口 4096 时，AIME25 33.9、GPQA 37.0、MATH500 91.5；全注意力是 34.2、37.6、91.7。窗口开到 8192，AIME25 反超到 35.8，GPQA 到 38.0。吞吐在 128K 序列上，窗口 4096 为 5224 tok/s，全注意力只剩 448；32K 上是 5479 对 1477，大约 3.7 倍。摘要里写的「3 倍」对应图 1 的墙钟思考时间：同一时间内能多吐 token，不是每个 token 本身更聪明。
单卡 H100 上，Prefix Sliding 的自定义核和普通滑动窗核几乎同速，过了预热期后稳定在大约 5000 tok/s。全注意力会一直变慢。
等内存预算的 RL 两边都卡在 8192。全注意力最长就是 8K，Prefix Sliding 能滚到 10 万 token 以上，reward 跟着上去。消融里，纯滑动窗会忘掉题目；last-k 和摘要压缩都要重复处理 token，内存呈锯齿状，还多一组超参。AIME25、最大生成 262144、局部窗 4096 的设置下，Prefix Sliding 的性能-效率曲线最好，而且只多一个超参：窗口大小。
这是给现成推理模型加的一层注意力掩码，不用重训就能把长思考的边际成本压成常数。竞赛数学、需要想很久的 agent 任务，可以先拿窗口 4096 或 8192 试。训练侧则让「超长 generation 直接截断丢掉」这件事变得没那么必要。
它不是新架构。Mamba、RWKV 那些线性模型要重训；Prefix Sliding 现有 Transformer 权重可以直接套。StreamingLLM 只留开头 4 个 sink token。工具定义可能几千步都没人看，后面一旦调用就必须还在，所以整段前缀都要钉住。H2O 按历史注意力留 heavy hitter，是往后看的；Prefix Sliding 是往前看的。两者可以叠，这篇没做成公平对照，因为 H2O 接不进 vLLM 和 FlashAttention。
短任务别指望它。HealthBench 平均只要 2086 token，窗口 2048 时几乎不滑，速度和全注意力差不多。
LiveCodeBench 上窗口至少要 16384 才能追上全注意力。模型会先写出函数骨架，再在注释里想几千 token，等写回代码时开头已经滑出窗口。这是免训练评估。RL 可能会教会模型少在注释里神游，这篇没验证。
工具输出和多轮是真坑。网页或文件一次性灌进来，窗口比内容短就会读不全，重要指令也可能被冲掉。后续用户消息该钉进前缀，还是交给窗口自然淘汰，论文只把问题摊开，没有方案。前缀本身如果极长，prefill 的 KV 并没有被压缩。
主数字来自 1.7B，截断反传只在 7B 上做过短实验。和 RNN、SSM、混合滑动窗架构没有同算力对照。图 9 里 last-k 和摘要的超参是在另一套小上下文上扫的，搬到 262K 生成时未必最优。