Zhenrui Wang

Author directory

2026

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.