投机解码(speculative decoding)已经成了推理侧的标准配置:小模型先猜一串 token,大模型一次验证多个,输出分布不变,把多步解码压进一次前向。但 reasoning 和 agent 应用把输入、输出同时拉长之后,这套体系里一个不起眼的部件开始顶到天花板——草稿模型的 KV cache。

10 月 22 日提交到 arXiv 的论文《RIFT: KV Cache Compression for Draft Models with Complementary Compressibility Training》(arXiv:2610.31098)正面处理了这个问题。十位作者来自 NVIDIA,其中四位也是论文所用的 SPEED-Bench 基准的共同作者。

先看一个反直觉结果

论文的第一组实验就颠覆了一个默认认知。同样 8K 的 KV 预算,让 Qwen3.5-4B 给 Qwen3.5-122B 当草稿、跑 32K 输入的 SPEED-Bench:带重要性打分的 SnapKV,输给了"无脑只留最近 token"的滑动窗口(SWA)。随着生成推进,SnapKV 的平均接受长度(AL)下降 36.7%,滑动窗口反而提升 41%。

机制层面的证据更直接:全缓存注意力中落在保留集合上的比例,SnapKV 从 67.0% 一路掉到 25.4%,滑动窗口从 43.5% 涨到 68.7%,两条曲线在中途交叉。作者把这种现象命名为 decode staleness(解码陈旧化):prefill 阶段选出来的"重要"KV,到长输出阶段已经不再重要。

为什么会陈旧

论文给出两条归因。其一,投机解码里最近的解码前缀是目标模型验证过的输出,草稿模型盯住它才能保持与目标对齐;其二,reasoning 场景下最近 token 本身就是一路滚动的"思维链摘要"。消融实验很说明问题:关闭 reasoning 后,staleness gap 从 2.11 掉到 0.60 AL——推理链越长,陈旧化越狠。

还有一个架构维度的发现:在 Nemotron3(hybrid-SSM)和 GPT-OSS(硬编码 SWA 交错全注意力)这类混合注意力模型上,SWA 压过 SnapKV;只有在全注意力的 Llama3.1 上 SnapKV 才保持优势。而 KV 压缩文献的评测恰恰大多建立在全注意力加短输出的设定上。

RIFT 的解法:三档策略加一次训练

  • RIFT-stat:最简单的静态版。把预算 B 按比例 ρ 切成"最近 token 窗口 + 重要性 prompt 位置"两块。ρ=0.5 时 AL 达 5.614,比 SnapKV 的 4.073 高出 37.8%,也高于纯 SWA 的 5.230。
  • RIFT-dyn:解码过程中周期性重打分、刷新选择,质量更高但要付额外延迟。
  • RIFT-ada:自适应刷新,比 dyn 少 51.1% 的刷新次数,AL 反而还高 0.122。
  • 压缩感知训练:用稀疏注意力掩码 + LoRA + 蒸馏(CE+KL)做一轮 post-training,让草稿模型天生适应压缩缓存。Qwen3.5-4B 最高拿到 +16.19%,Nemotron3.5-Lightning 最高 +16.42%。

合成成绩:SPEED-Bench 32K 输入、4× 压缩比下,Qwen3.5-4B 保留 96% 的全缓存 AL,Nemotron3.5-Lightning 保留 97%(给 Nemotron3-Ultra 当草稿);SnapKV 在两组设定里只有 63-67%。结合压缩感知训练后,32K 输入、4K 输出、8K 预算下接受长度保留超过 98%。对比 StreamingLLM、H2O、TriAttention 三个解码时基线,RIFT-dyn 和 RIFT-ada 在 32K/100K 两档上下文、两种温度下全部占优。

别急着上线,看清两个边界

论文的 Limitations 写得诚实。第一,在纯 GPU 部署里,外部草稿模型的端到端加速本来就不如 MTP、EAGLE、DFlash 这类 tail-augmented 方案——RIFT 的真实靶场是异构推理:GPU 负责打分,SRAM 类低内存加速器只跑草稿,论文引用了 NVIDIA LPX 低延迟加速器技术博客作为这一方向的早期信号。第二,解码中重打分的开销在 GPU 上约占草稿解码延迟的 6.0%,搬到 SRAM 设备上这个比例还要重新量。

对行业的信号比单篇结果更重要:混合注意力正在成为旗舰模型的默认架构,而 KV 压缩的整套评测基准大多还是全注意力时代的遗产。方法选型之前,先问两个问题——你的输出长度分布是什么,你的模型是哪种注意力结构。这两个问题的答案,可能比压缩方法本身的榜单名次更值钱。