Внимание — это всё, что нужно для кэша ключей и значений в диффузионных языковых моделях

Attention Is All You Need for KV Cache in Diffusion LLMs

October 16, 2025

Авторы: Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen

cs.AI

Аннотация

В данной работе исследуется, как адаптивно пересчитывать кэши ключей и значений (KV) для диффузионных больших языковых моделей (DLM), чтобы максимизировать точность предсказаний при минимизации задержки декодирования. В существующих методах декодеры пересчитывают QKV для всех токенов на каждом шаге шумоподавления и на каждом слое, несмотря на то, что состояния KV изменяются незначительно на большинстве шагов, особенно на поверхностных слоях, что приводит к значительной избыточности. Мы делаем три наблюдения: (1) удалённые токены {bf MASK} в основном служат как смещение по длине и могут быть кэшированы блоками за пределами активного окна предсказания; (2) динамика KV увеличивается с глубиной, что указывает на достаточность выборочного обновления, начиная с более глубоких слоёв; и (3) токен, к которому обращаются чаще всего, демонстрирует наименьшее смещение KV, что даёт консервативную нижнюю границу изменения кэша для других токенов. На основе этого мы предлагаем {bf Elastic-Cache}, стратегию, не требующую обучения и независимую от архитектуры, которая совместно определяет, {когда} обновлять (с помощью теста на смещение, учитывающего внимание, для наиболее часто используемого токена) и {где} обновлять (с помощью расписания, учитывающего глубину, которое пересчитывает кэши, начиная с выбранного слоя, повторно используя кэши поверхностных слоёв и кэши MASK за пределами окна). В отличие от схем с фиксированным периодом, Elastic-Cache выполняет адаптивные, учитывающие слои обновления кэшей для диффузионных LLM, сокращая избыточные вычисления и ускоряя декодирование с незначительной потерей качества генерации. Эксперименты на LLaDA-Instruct, LLaDA-1.5 и LLaDA-V в задачах математического рассуждения и генерации кода демонстрируют стабильное ускорение: в 8.7 раз на GSM8K (256 токенов), в 45.1 раз на более длинных последовательностях и в 4.8 раз на HumanEval, при этом сохраняя более высокую точность по сравнению с базовым методом. Наш метод обеспечивает значительно более высокую пропускную способность (в 6.8 раз на GSM8K) по сравнению с существующими подходами, основанными на уверенности, сохраняя качество генерации, что делает возможным практическое использование диффузионных LLM.

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.