Mask-Guided KV Cache Eviction in Block Diffusion Language Models
Abstract Overview
This paper studies inference-time key-value (KV) cache management for block diffusion language models, where long prefixes are both stored in memory and repeatedly read during denoising. The authors argue that two distinct decisions must be made over the prefix: a temporary read selection for the current block and a persistent retention decision for future blocks. They propose MaskAhead, a training-free method that uses current-block mask queries to choose sparse reads and future-block mask probes to rank which KV entries should remain in the persistent cache. The method uses a value-aware scoring rule and physically compacts the cache under a fixed budget, and the paper also evaluates a 4-bit storage variant, Q-MaskAhead.
Novelty
The distinctive idea is to separate read selection from irreversible retention and to use different query horizons for the two decisions: current masks for immediate reads and future-mask probes for eviction. The work also combines this recurrent eviction scheme with value-aware ranking and, in Q-MaskAhead, performs both ranking and attention from compact 4-bit stored KV representations without a bf16 shadow cache.
Results
Across the six main long-prompt QA settings, MaskAhead achieves the highest F1 among the evicting methods, including one tie, while reducing KV memory by 9.5× on average with a mean 1.2-point F1 drop relative to dense inference. Q-MaskAhead further lowers memory usage, reaching an average 20.1× KV reduction on the QA settings with a mean 2.3-point F1 loss relative to dense inference. In a batch-profile pilot, MaskAhead and Q-MaskAhead also reduced KV memory by 85.1% and 89.9% at batch 32 and improved recorded decode-stage time relative to dense bf16.
Key Points
- MaskAhead treats sparse reading and cache retention as separate inference decisions and uses different mask-query horizons for each.
- The method performs training-free, value-aware KV eviction with physical cache compaction, rather than only reducing reads or offloading the full cache.
- Experiments across three block diffusion language models show strong memory savings with comparatively small quality degradation, and the 4-bit Q-MaskAhead variant pushes memory lower at additional quality cost.