SampKV: Saliency-Aware Mixed-Precision KV Cache for Efficient Diffusion Language Model Inference

Zhenrui Wang, Haodong Chang, Yang Sui, Rongjian Liang, Bo Yuan, Jiang Hu


Abstract
Diffusion language models (dLLMs) have recently shown strong potential in various tasks with competitive decoding efficiency, but their bidirectional attention incurs quadratic compute and memory complexity. Many prior works rely on KV cache optimizations to accelerate inference. However, the full-sequence BF16 KV cache remains the dominant memory bottleneck. We propose SampKV, a Saliency-aware Mixed-Precision KV cache for dLLMs that stores most history in 2 bits while promoting a small budget of attention-salient tokens to high precision. A customized Triton kernel operates directly on the mixed-precision cache and avoids materializing full-precision KV tensors, alleviating the storage and traffic bottleneck. Experiments on LLaDA and Dream models show that the memory bottleneck is substantially alleviated: up to 54% peak memory reduction at matched workload, or up to 2.06x throughput improvement at matched peak memory compared to the BF16 baseline, while maintaining the accuracy.
Anthology ID:
2026.emnlp-main.1968
Volume:
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Month:
October
Year:
2026
Address:
Budapest, Hungary
Editors:
Hung-Yi Lee, Sunipa Dev, Jessy Li, Desmond Elliott
Venue:
EMNLP
SIG:
SIGDAT
Publisher:
Association for Computational Linguistics
Note:
Pages:
44808–44821
Language:
URL:
https://aclanthology.org/2026.emnlp-main.1968/
DOI:
Bibkey:
Cite (ACL):
Zhenrui Wang, Haodong Chang, Yang Sui, Rongjian Liang, Bo Yuan, and Jiang Hu. 2026. SampKV: Saliency-Aware Mixed-Precision KV Cache for Efficient Diffusion Language Model Inference. In Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing, pages 44808–44821, Budapest, Hungary. Association for Computational Linguistics.
Cite (Informal):
SampKV: Saliency-Aware Mixed-Precision KV Cache for Efficient Diffusion Language Model Inference (Wang et al., EMNLP 2026)
Copy Citation:
PDF:
https://aclanthology.org/2026.emnlp-main.1968.pdf