KV 缓存对于扩散 LLM 来说，注意力就是一切

这项工作研究了如何自适应地重新计算扩散大型语言模型（DLM）的关键值（KV）缓存，以最大化预测准确性同时最小化解码延迟。先前的方法的解码器在每个去噪步骤和层中重新计算所有 token 的 QKV，尽管 KV 状态在大多数步骤中变化不大，尤其是在浅层中，这导致大量冗余。我们有三个观察：（1）远处的 {bf MASK} token 主要起到长度偏置的作用，并且可以超过活动预测窗口进行块级缓存；（2）KV 动态随深度增加，表明从较深层开始选择性刷新就足够了；（3）最受关注的 token 表现出最小的 KV 漂移，为其他 token 提供了保守的缓存变化下限。基于这些观察，我们提出了 {bf Elastic-Cache}，一种无需训练、架构无关的策略，它同时决定 {何时} 刷新（通过对最受关注的 token 进行注意力感知漂移测试）和 {在哪里} 刷新（通过从选定层开始的深度感知调度，同时重用浅层缓存和窗口外的 MASK 缓存）。与固定周期方案不同，Elastic-Cache 对扩散 LLM 执行自适应的、层感知的缓存更新，减少冗余计算并加速解码，同时对生成质量的损失微乎其微。在 LLaDA-Instruct、LLaDA-1.5 和 LLaDA-V 上对数学推理和代码生成任务的实验表明，一致的加速效果：GSM8K（256 token）加速 8.7 倍，长序列加速 45.1 倍，HumanEval 加速 4.8 倍，同时始终保持比基线更高的准确性。我们的方法在 GSM8K 上实现的吞吐量（6.8 倍）显著高于现有的基于置信度的方法，同时保持了生成质量，从而能够实际部署扩散 LLM。