New KV Cache Quantization Methods for Diffusion Language Models: TRACE and dLLM-KVQ

Why KV Cache Quantization Is Trickier for Diffusion Language Models

KV cache quantization has been thoroughly studied in autoregressive language models. Under causal attention, the cached prefix states are exact, and quantization mainly addresses the problem of cache reconstruction. But diffusion language models (dLLMs) break this premise.

Fast-dLLM inference reuses KV states across denoising steps to accelerate generation, yet the bidirectional token representations change as masked tokens are progressively revealed. This means that, before being quantized, the reused cache may already deviate from the result of full recomputation. This research focuses precisely on the KV cache quantization problem under this "cache reuse error."

KV Cache and Autoregressive Decoding: In standard autoregressive (AR) language models, generation proceeds token by token serially. Each time a new token is generated, its corresponding Key/Value tensors are appended to the cache, and subsequent attention computations read these caches directly without recomputing historical positions. Causal attention guarantees that the cached prefix states will not be rewritten by subsequent tokens, so once written, the cache is fully consistent with the result of full recomputation. The goal of KV cache quantization is to compress these exact states to low bit-width while controlling reconstruction error. Acceleration schemes for dLLMs such as Fast-dLLM borrow this idea and cache and reuse states that "appear unchanged" across denoising steps, but bidirectional attention causes these states to change as masked tokens are revealed—this is the root cause of the article's statement that "the reused cache may already deviate from the recomputed result."

Diffusion Language Models and Iterative Masked Decoding: dLLMs (such as LLaDA and Dream) no longer generate from left to right; instead, they start from a fully masked sequence and progressively "reveal" masked tokens over several denoising steps. At each step, the model performs bidirectional attention over the entire sequence, and the representations of already revealed tokens are updated due to newly revealed tokens. This is fundamentally different from the AR model's assumption that "once the prefix is determined, it no longer changes." Therefore, KV states reused across steps are only approximations, and their error compounds with quantization error—this is the core reason why KV cache quantization for dLLMs requires separate design.

From Autoregressive to Diffusion: The Failure of Migrating Quantization Assumptions

Quantization methods for autoregressive models are built on the assumption that "the cache is exact," but the iterative masked decoding process of dLLMs invalidates this assumption. At each denoising step, some tokens are revealed, and bidirectional attention requires cached states to be updated based on newly revealed tokens. If quantization schemes designed for autoregressive models are applied directly, the error in the reused cache compounds with quantization error, leading to degraded output quality.

The authors adapt KIVI and KVQuant to Fast-dLLM-style prefix cache and dual cache reuse scenarios, and propose two new methods: TRACE and dLLM-KVQ.

KIVI and KVQuant: Both are representative KV cache quantization methods for autoregressive LLMs. KIVI quantizes the Key cache per-channel and the Value cache per-token, exploiting the different distribution characteristics of K and V; KVQuant introduces sensitivity-weighted non-uniform quantization, outlier handling, and calibration data to reduce accuracy loss at low bit-width. Their common implicit premise is that the cache states used for calibration and quantization are identical to those fully recomputed at inference time. After migrating these two types of methods to Fast-dLLM's prefix cache and dual cache reuse scenarios, the authors found that reuse error breaks this premise, thus necessitating calibration targeted at reveal trajectories such as TRACE, and strategies such as dLLM-KVQ that allocate bits per layer to compensate.

TRACE: Constructing Calibration Trajectories for Diffusion Decoding

The core idea of TRACE is to reconstruct feasible dLLM reveal trajectories from prompt-answer pairs. These trajectories provide calibration states that match the iterative masked decoding process.

Unlike autoregressive models, which use fixed-prefix calibration, calibration for dLLMs needs to cover the dynamic process of tokens being progressively revealed. TRACE constructs these trajectories so that quantization calibration more closely matches the cache state distribution during actual inference, thereby reducing the impact of reuse error.

dLLM-KVQ: Selecting High-Precision K Cache Entries per Layer

dLLM-KVQ uses the trajectories generated by TRACE to select, based on output-level and attention-level behavioral metrics, the K cache entries in each layer that need to be kept at high precision. That said, standard low-bit quantization is applied to the value cache.

The intuition behind this asymmetric strategy is that the K cache is more sensitive to the computation of attention weights, while the V cache has relatively higher quantization tolerance. By retaining high precision only for some K entries, the method strikes a balance between compression rate and model quality.

Asymmetric Sensitivity of K/V Caches: In attention computation, Keys directly participate in softmax scoring, and tiny perturbations to them can change the distribution of attention weights, which is then amplified in the output; Values are weighted sums computed after the weights are determined, so errors are more easily averaged out. Existing AR quantization research has also observed that K is more sensitive than V, often using finer granularity or higher bit-width for K. dLLM-KVQ combines this intuition with the dynamics of diffusion decoding: with the help of TRACE trajectories, it identifies the most critical K entries in each layer according to output-level and attention-level behavioral metrics and keeps them at high precision, while the remaining K entries and all V entries undergo standard low-bit quantization, thereby achieving a better balance between compression rate and quality.

Experimental Results: The Highest Average Score Among Quantization Methods

Evaluated on the LLaDA and Dream models across four benchmarks—GSM8K, HumanEval, MATH, and MBPP—dLLM-KVQ achieves the highest average score among all quantization methods in the reported 4-bit and 2-bit prefix/dual cache settings.

This result shows that a quantization scheme designed for dLLM cache reuse error is more effective than directly migrating autoregressive quantization methods. The research also suggests that cache management for dLLM inference acceleration requires methods specifically targeted at the dynamics of bidirectional representations, rather than simple engineering adaptation.

Implications for dLLM Inference Deployment

The value of this work lies in clarifying an easily overlooked issue: KV cache quantization for dLLMs is not a direct extension of autoregressive quantization. Cache reuse error is an error source unique to dLLMs and must be explicitly handled in quantization calibration and bit allocation.

For practical deployment, the trajectory calibration idea provided by TRACE and the per-layer mixed-precision strategy of dLLM-KVQ offer reference directions for running diffusion language models under limited GPU memory. As dLLMs are increasingly applied to tasks such as code generation and mathematical reasoning, the importance of such inference optimization work will become even more evident.