확산 LLM에서 KV 캐시를 위한 주의력이 전부다

초록

본 연구는 디퓨전 대형 언어 모델(DLMs)의 키-값(KV) 캐시를 적응적으로 재계산하여 예측 정확도를 극대화하고 디코딩 지연 시간을 최소화하는 방법을 탐구합니다. 기존 방법들의 디코더는 모든 디노이징 단계와 레이어에서 모든 토큰에 대해 QKV를 재계산하지만, 특히 얕은 레이어에서는 KV 상태가 대부분의 단계에서 거의 변하지 않아 상당한 중복이 발생합니다. 우리는 세 가지 관찰을 통해 다음과 같은 사실을 발견했습니다: (1) 멀리 떨어진 {bf MASK} 토큰은 주로 길이 편향으로 작용하며, 활성 예측 창을 넘어 블록 단위로 캐싱할 수 있다; (2) KV 동적성은 깊이에 따라 증가하므로, 더 깊은 레이어부터 선택적으로 갱신하는 것이 충분하다; (3) 가장 많이 주목받는 토큰은 가장 작은 KV 드리프트를 보이며, 이는 다른 토큰에 대한 캐시 변경의 보수적인 하한을 제공한다. 이를 바탕으로, 우리는 {bf Elastic-Cache}를 제안합니다. 이는 학습이 필요 없고 아키텍처에 구애받지 않는 전략으로, {언제} 갱신할지(가장 많이 주목받는 토큰에 대한 주의 기반 드리프트 테스트를 통해)와 {어디서} 갱신할지(선택된 레이어부터 재계산하면서 얕은 레이어 캐시와 창 밖의 MASK 캐시를 재사용하는 깊이 인식 스케줄을 통해)를 공동으로 결정합니다. 고정 주기 방식과 달리, Elastic-Cache는 디퓨전 LLMs를 위해 적응적이고 레이어 인식 캐시 업데이트를 수행하여 중복 계산을 줄이고 디코딩을 가속화하며 생성 품질의 손실을 무시할 수준으로 유지합니다. LLaDA-Instruct, LLaDA-1.5, LLaDA-V에서 수학적 추론 및 코드 생성 작업에 대한 실험은 일관된 속도 향상을 보여줍니다: GSM8K(256 토큰)에서 8.7배, 더 긴 시퀀스에서 45.1배, HumanEval에서 4.8배의 속도 향상을 달성하면서도 기준선보다 더 높은 정확도를 유지합니다. 우리의 방법은 기존의 신뢰도 기반 접근법보다 훨씬 높은 처리량(GSM8K에서 6.8배)을 달성하면서도 생성 품질을 보존하여, 디퓨전 LLMs의 실용적인 배포를 가능하게 합니다.

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.7times on GSM8K (256 tokens), 45.1times on longer sequences, and 4.8times on HumanEval, while consistently maintaining higher accuracy than the baseline. Our method achieves significantly higher throughput (6.8times on GSM8K) than existing confidence-based approaches while preserving generation quality, enabling practical deployment of diffusion LLMs.

확산 LLM에서 KV 캐시를 위한 주의력이 전부다

Attention Is All You Need for KV Cache in Diffusion LLMs

초록

Support