Zhenrui Wang
Author directory2026
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
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Zhenrui Wang | Haodong Chang | Yang Sui | Rongjian Liang | Bo Yuan | Jiang Hu
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
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.