arXiv:2510.14973cs.CLcs.AI2025-10被引 31

提出弹性缓存机制,让扩散型大模型推理更快更准。

Attention Is All You Need for KV Cache in Diffusion LLMs

  • 根据注意力动态判断何时刷新缓存,从深层开始逐层更新
  • 最长序列提速45.1倍,数学推理任务准确率更高
  • 无需训练,适配多种模型,适合实际部署

本文研究如何自适应地重计算扩散型大语言模型(DLMs)中的键值(KV)缓存,以在保持预测精度的同时最小化解码延迟。以往方法在每个去噪步骤和层级中均重算所有令牌的QKV,尽管大多数步骤中尤其是浅层的KV状态变化极小,造成大量冗余。我们提出三点观察:(1) 远距离的${f MASK}$令牌主要起长度偏差作用,可块级缓存于活跃预测窗口之外;(2) KV动态随深度增加,仅从深层开始选择性刷新即可;(3) 被最多关注的令牌具有最小的KV漂移,可作为其他令牌缓存更新的保守下界。基于此,我们提出${f Elastic-Cache}$,一种无需训练、与架构无关的策略,联合决定${when}$(通过最被关注令牌的注意力感知漂移检测)和${where}$(通过深度感知调度,从选定层级起逐层重算,复用浅层缓存及窗口外的MASK缓存)。相比固定周期方案,Elastic-Cache实现自适应、层级感知的缓存更新,显著减少冗余计算,加速解码且生成质量损失可忽略。在LLaDA-Instruct、LLaDA-1.5和LLaDA-V上进行的实验表明,在数学推理与代码生成任务中持续提速:在GSM8K(256 tokens)上达8.7倍,长序列达45.1倍,同时始终优于基线准确率。本方法吞吐量比现有置信度方法高6.8倍,且保持生成质量,支持扩散型大模型的实际部署。

原文摘要 · Abstract (English)

This work studies how to adaptively recompute key-value (KV) caches for diffusion large language models (DLMs) to maximize prediction accuracy while minimizing decoding latency. Prior methods' decoders recompute QKV for all tokens at every denoising step and layer, despite KV states changing little across most steps, especially in shallow layers, leading to substantial redundancy. We make three observations: (1) distant ${\bf MASK}$ tokens primarily act as a length-bias and can be cached block-wise beyond the active prediction window; (2) KV dynamics increase with depth, suggesting that selective refresh starting from deeper layers is sufficient; and (3) the most-attended token exhibits the smallest KV drift, providing a conservative lower bound on cache change for other tokens. Building on these, we propose ${\bf Elastic-Cache}$, a training-free, architecture-agnostic strategy that jointly decides ${when}$ to refresh (via an attention-aware drift test on the most-attended token) and ${where}$ to refresh (via a depth-aware schedule that recomputes from a chosen layer onward while reusing shallow-layer caches and off-window MASK caches). Unlike fixed-period schemes, Elastic-Cache performs adaptive, layer-aware cache updates for diffusion LLMs, reducing redundant computation and accelerating decoding with negligible loss in generation quality. Experiments on LLaDA-Instruct, LLaDA-1.5, and LLaDA-V across mathematical reasoning and code generation tasks demonstrate consistent speedups: $8.7\times$ on GSM8K (256 tokens), and $45.1\times$ on longer sequences, while consistently maintaining higher accuracy than the baseline. Our method achieves significantly higher throughput ($6.8\times$ on GSM8K) than existing confidence-based approaches while preserving generation quality, enabling practical deployment of diffusion LLMs.

扩散模型缓存优化推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。