KV Cache Quantization for Diffusion Language Models

Yan Tai; Yunan Ding; Guangzhao Li; Siqi Luo; Chaoyang Zhang; Guangtao Zhai; Guodong Guo; Xiaohong Liu

cs.CLOct 29, 2026

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

Abstract

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.

Prefix cache
Method	GSM8K	HumanEval	MATH	MBPP	Avg.
LLaDA full attention	78.62	42.07	33.36	29.80	45.96
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.

Prefix cache
Method	GSM8K	HumanEval	MATH	MBPP	Avg.
Dream full attention	74.15	39.63	34.18	54.40	50.59
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 .

Mode	L2	KL	Top-1	Top- $K$	Eff. bits
	Win rate (%)
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

Panel of Figure 3 — 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 .

KV-Cache Quantization

Key-value cache quantization compresses transformer attention state to reduce memory use and bandwidth during long-context inference.

See all 104 papers

EntropyCache: Decoded Token Entropy Guided KV Caching for Diffusion Language Models
Mar 19, 2026Minsoo Cheong, Donghyun Son, Woosang Lim +1Diffusion Language Model Inference Diffusion Model Caching
Seoul National University
Mask-Guided KV Cache Eviction in Block Diffusion Language Models
Oct 4, 2026Gleb Molodtsov, Ekaterina Alimaskina, Evgeny Uskov +2KV-Cache Eviction Diffusion Language Model Inference
BRAIn Lab · Yandex Research · Yandex School of Data Analysis
SoloQ: Calibration-Free Quantization for Diffusion Language Models
Oct 5, 2026Donghyun Lee, Arkapravo Ghosh, Varun Manjunath +5Diffusion Model Quantization KV-Cache Quantization
University of Southern California · Yale University · Korea University

Mar 19, 2026J/K move · Enter open · S save

EntropyCache: Decoded Token Entropy Guided KV Caching for Diffusion Language Models

Minsoo Cheong, Donghyun Son, Woosang Lim, Sungjoo Yoo

Seoul National University

Diffusion Language Model Inference Diffusion Model Caching Diffusion Language Models

Open paper AskPDF arXiv

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\times

26.4\times

speedup on standard benchmarks and

22.4\times

24.1\times

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.

Abstract

KV-Cache Quantization

Explore similar work

EntropyCache: Decoded Token Entropy Guided KV Caching for Diffusion Language Models

Mask-Guided KV Cache Eviction in Block Diffusion Language Models

SoloQ: Calibration-Free Quantization for Diffusion Language Models