KV Cache Quantization for Diffusion Language Models
Authors: Yan Tai, Yunan Ding, Guangzhao Li, Siqi Luo, Chaoyang Zhang, Guangtao Zhai, Guodong Guo, Xiaohong Liu
Organizations: Shanghai Jiao Tong University · Ningbo Institute of Digital Twin, Eastern Institute of Technology, Ningbo, China · The Hong Kong Polytechnic University · Shanghai Innovation Institute
KV cache quantization is well studied for autoregressive language models, where cached prefix states are exact under causal attention and quantization is mainly a cache reconstruction problem. Diffusion language models (dLLMs) break this assumption. Fast dLLM inference reuses KV states across denoising steps, but bidirectional token representations can change as masked tokens are revealed, so the reused cache may already differ from full recomputation before quantization is applied. We study KV cache quantization under this cache-reuse error. We adapt KIVI and KVQuant to Fast-dLLM-style prefix and dual cache reuse, and introduce TRACE and dLLM-KVQ. TRACE reconstructs feasible dLLM reveal trajectories from prompt-answer pairs, providing calibration states that match iterative masked decoding. dLLM-KVQ uses these trajectories to choose per-layer K-cache entries that remain in high precision according to output- and attention-level behavior metrics, while applying standard low-bit value quantization. Across LLaDA and Dream evaluations on GSM8K, HumanEval, MATH, and MBPP, dLLM-KVQ achieves the highest average score among quantized methods across the reported 4-bit and 2-bit prefix/dual cache settings.
Figures & tables
Figure 1: Fast-dLLM cache reuse and cache-reuse error. (a) Dual cache reuses prefix and suffix KV states from a block-initial forward pass while recomputing only the active block. (b) The reused cache can differ from the corresponding full-recompute cache after masked tokens inside the block are revealed.
Figure 2: dLLM-KVQ pipeline. (a) TRACE reconstructs masked-denoising trajectories from prompt-answer pairs. (b) Offline diagnostics fit per-layer K-cache policies from behavior and attention metrics. (c) Online quantization applies the selected policy to block-initial reused cache states.
Method
GSM8K
HumanEval
MATH
MBPP
Avg.
LLaDA full attention
78.62
42.07
33.36
29.80
45.96
Prefix cache
Cache only
78.47
42.07
33.18
26.60
45.08
KIVI
77.33
42.68
32.40
25.60
44.50
KVQuant
77.48
42.07
32.68
24.60
44.21
dLLM-KVQ
77.94
43.90
32.92
26.60
45.34
Table 1: LLaDA results under 4-bit and 2-bit KV cache quantization. LLaDA full attention denotes no-cache original inference. Cache-only rows are unquantized cache references. For dLLM-KVQ, each row reports the selected configuration in the current result matrix. Average denotes the arithmetic mean over GSM8K, HumanEval, MATH, and MBPP.
Method
GSM8K
HumanEval
MATH
MBPP
Avg.
Dream full attention
74.15
39.63
34.18
54.40
50.59
Prefix cache
Cache only
74.60
43.29
34.20
55.20
51.82
KIVI
73.69
40.85
33.74
54.20
50.62
KVQuant
74.83
37.80
33.86
54.60
50.27
dLLM-KVQ
74.53
41.46
34.28
56.20
51.62
Table 2: Dream results under 4-bit and 2-bit KV cache quantization. Dream full attention denotes no-cache original inference. Cache-only rows are unquantized cache references. For dLLM-KVQ, each row reports the selected configuration in the current result matrix. Average denotes the arithmetic mean over GSM8K, HumanEval, MATH, and MBPP.
Method
K high-precision rule
MBPP
KIVI K/V
None; all low-bit
26.60
Fixed K rule
Attention, p=4
25.60
Fixed K rule
Attention, p=2
25.20
dLLM-KVQ
Selected per layer
26.60
Table 3: Task-level LLaDA MBPP policy controls under 4-bit K/V cache quantization. Scores are pass@1. The p value denotes how many K tokens per token group are kept in high precision. Fixed K-rule rows use the same rule in every layer, while dLLM-KVQ selects this rule separately for each layer using trajectory calibration. Dual rows use the same generation length, block length, and cache mode as Table 1 .
Win rate (%)
Mode
L2
KL
Top-1
Top- K
Eff. bits
Prefix
96.4
88.9
72.9
85.9
4.77
Dual
94.9
90.0
71.7
82.5
4.60
Table 4: Trajectory-level win rates for the compiled K-side policy on LLaDA 4-bit calibration traces. Each win rate is the percentage of paired records improved over unprotected K quantization; L2 and KL are converted to improvements by lower distance/divergence, while Top-1 and Top- K use higher agreement/recall.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Representative K/V cache distributions used to choose quantization granularity.
Figure 4: K-cache absolute-value surfaces across denoising steps in the same active block. The top row shows LLaDA and the bottom row shows Dream; columns show early, middle, and late denoising steps for block [2016,2048) .
Figure 5: V-cache absolute-value surfaces for the same layer, head, block, and denoising steps as Figure 4 .
Diffusion-based large language models (dLLMs) rely on bidirectional attention, which prevents lossless KV caching and requires a full forward pass at every denoising step. Existing approximate KV caching methods reduce this cost by selectively updating cached states, but their decision overhead scales with context length or model depth. We propose EntropyCache, a training-free KV caching method that uses the maximum entropy of newly decoded token distributions as a constant-cost signal for deciding when to recompute. Our design is grounded in two empirical observations: (1) decoded token entropy correlates with KV cache drift, providing a cheap proxy for cache staleness, and (2) feature volatility of decoded tokens persists for multiple steps after unmasking, motivating recomputation of the k most recently decoded tokens. The skip-or-recompute decision requires only O(V) computation per step, independent of context length and model scale. Experiments on LLaDA-8B-Instruct and Dream-7B-Instruct show that EntropyCache achieves 15.2×-26.4× speedup on standard benchmarks and 22.4×-24.1× on chain-of-thought benchmarks, with competitive accuracy and decision overhead accounting for only 0.5% of inference time. Code is available at https://github.com/mscheong01/EntropyCache.
Block diffusion language models keep a large key-value (KV) cache throughout generation and attend to it at every denoising step, limiting both memory capacity and generation speed. Reducing these costs requires deciding which past tokens to use for denoising the current block (selection) and which to keep in memory for future blocks (eviction). We propose MaskAhead, a training-free method that solves both tasks with a single mask-query-based ranking mechanism. Current-block masks guide selection, while probes of upcoming masked blocks guide eviction. Both rank KV entries by their estimated contribution to the attention output. Our quantized variant, Q-MaskAhead, computes selection and attention directly from low-bit KV, largely preserving the selected entries. Experiments on Fast-dLLM-v2, DreamReasoner, and LLaDA2.0-mini cover long-generation reasoning, long-prompt question answering, and needle-in-a-haystack retrieval. On long-prompt QA, MaskAhead reduces KV memory by 9.5× on average with a 1.2-point mean F1 loss relative to dense inference. Q-MaskAhead increases the reduction to 20.1× with a 2.3-point mean F1 loss. In a batch-32 systems profile, MaskAhead achieves 1.23× end-to-end and 1.68× decode-stage speedups over dense inference.
Diffusion large language models dLLMs) have emerged as a promising alternative to autoregressive language models through bidirectional diffusion-based token generation. However, their growing model sizes and high inference costs make efficient deployment challenging: full-sequence denoising repeatedly invokes compute-intensive forward passes, while block-diffusion models additionally introduce a memory-intensive KV-cache. Low-bit weight-activation quantization is therefore attractive, yet existing dLLM post-training quantization methods rely on calibration data despite activation distributions shifting across masking states and denoising steps. We present SoloQ, a calibration-free quantization framework that maps weights and activations into a normalized rotated basis with a predictable marginal distribution, enabling data-independent quantization. SoloQ combines a structured K-RPBH rotation with a lightweight rescaling correction for calibration-free quantization. Its predictable post-rotation distribution supports both distribution-matched codebooks and hardware-native NVFP4. For block-diffusion models, SoloQ further applies commit-time KV-cache quantization to compress persistent states without perturbing the actively denoised block. Across full-sequence dLLMs (LLaDA and Dream) and block-diffusion dLLMs(Fast-dLLM v2 and Nemotron-Labs-Diffusion), SoloQ retains accuracy under 4-bit quantization and outperforms calibration-based baselines on knowledge- and reasoning-intensive benchmarks. With NVFP4, SoloQ reduces peak memory by up to 2.61X and accelerates end-to-end inference by up to 2.24X.
Donghyun Lee, Arkapravo Ghosh, Varun Manjunath +5
University of Southern California · Yale University · Korea University