d$^2$Cache: Accelerating Diffusion-Based LLMs via Dual Adaptive Caching

TL;DR

This work tackles the inefficiency of diffusion-based LLMs by introducing d$^2$Cache, a training-free, fine-grained approximate KV cache that selectively updates a small subset of token KV states at each decoding step while caching the rest. Leveraging a two-stage strategy—certainty-prior based selection for masked tokens and attention-rollout informed selection for prompt/decoded tokens—d$^2$Cache enables quasi left-to-right decoding and substantial speedups without degrading generation quality. The approach is validated on LLaDA and Dream across multiple benchmarks, delivering 3.2×–4.0× speedups and competitive or improved accuracy relative to strong baselines. Overall, d$^2$Cache offers a practical and effective pathway to accelerate dLLMs, with open-source code available for adoption.

Abstract

Diffusion-based large language models (dLLMs), despite their promising performance, still suffer from inferior inference efficiency. This is because dLLMs rely on bidirectional attention and cannot directly benefit from the standard key-value (KV) cache as autoregressive models (ARMs) do. To tackle this issue, we introduce \textit{Dual aDaptive Cache} (d$^2$Cache), which is a training-free approximate KV cache framework for accelerating dLLM inference. d$^2$Cache features a two-stage fine-grained selection strategy to identify tokens and adaptively update their KV states at each decoding step, while caching the KV states of the remaining tokens for reuse. Furthermore, d$^2$Cache naturally offers a more reliable decoding alternative, which can enable quasi left-to-right generation and mitigate premature overconfidence in tokens at the end of the sequence. Extensive experimental results on two representative dLLMs (\ie, LLaDA and Dream) demonstrate that d$^2$Cache not only achieves substantial inference speedups, but also yields consistent improvements in generation quality. The code is available at https://github.com/Kamichanw/d2Cache.

Figures (9)

• Figure 1: (a) In ARMs, causal attention requires each token to interact only with its preceding tokens. (b) In dLLMs, bidirectional attention requires each token to attend to both its preceding and subsequent tokens, such that any modification in the subsequent tokens necessitates recomputation of the entire sequence. (c) The proposed d$^2$Cache adaptively selects a small subset of tokens in dLLMs and updates their KV states through a two-stage fine-grained process. The KV states of the remaining tokens can be approximately cached for reuse in subsequent decoding step.
• Figure 2: (a) PCA of 77th masked token's trajectory on LLaMA-8B-Instruct with GSM8K ($L$=328, $n$=256, and $T$=256). (b) Sequential distances between token pairs decoded in adjacent steps.
• Figure 3: Attention rollout analysis over sequence, where the example and setting are the same as in \ref{['fig:kv-dynamics']}. (a) Attention rollout visualization at step 126, showing the sum of rollout values over all key positions (top) and the pairwise rollout values across different positions (bottom). (b) The total absolute differences in rollout values between each two adjacent decoding steps.
• Figure 4: Visualization of the decoding order using certainty prior with different $\sigma$ and NAR decoding. Each dot at $(i, t)$ indicates that the token at position $i$ is decoded at step $t$.
• Figure 5: Comparisons of different update strategies, including updating tokens only during the rapid-change phase (Phase 2 Only) and updating tokens during both the gradual-change and rapid-change phases (Phases 1 & 2).