Mask-Guided KV Cache Eviction in Block Diffusion Language Models

2026-10-04

AI Summary

This study addresses the high KV cache memory consumption and slow generation speed inherent in block diffusion language models by proposing MaskAhead, a training-free method. MaskAhead introduces a unified masked query ranking mechanism that enables efficient selection and eviction of KV cache entries through a single-pass evaluation of attention output contributions. Furthermore, a quantized variant, Q-MaskAhead, is presented to support direct computation with low-bit-width KV representations. Experimental results on long-prompt question-answering tasks demonstrate that MaskAhead reduces KV memory usage by an average factor of 9.5× while incurring only a marginal 1.2-point drop in F1 score. Additionally, Q-MaskAhead achieves up to 20.1× memory reduction and accelerates decoding speed by 1.68×, highlighting its effectiveness for efficient inference.

Abstract

Block diffusion language models keep a large key-value (KV) cache throughout generation and attend to it at every denoising step, limiting both memory capacity and generation speed. Reducing these costs requires deciding which past tokens to use for denoising the current block (selection) and which to keep in memory for future blocks (eviction). We propose MaskAhead, a training-free method that solves both tasks with a single mask-query-based ranking mechanism. Current-block masks guide selection, while probes of upcoming masked blocks guide eviction. Both rank KV entries by their estimated contribution to the attention output. Our quantized variant, Q-MaskAhead, computes selection and attention directly from low-bit KV, largely preserving the selected entries. Experiments on Fast-dLLM-v2, DreamReasoner, and LLaDA2.0-mini cover long-generation reasoning, long-prompt question answering, and needle-in-a-haystack retrieval. On long-prompt QA, MaskAhead reduces KV memory by $9.5\times$ on average with a 1.2-point mean F1 loss relative to dense inference. Q-MaskAhead increases the reduction to $20.1\times$ with a 2.3-point mean F1 loss. In a batch-32 systems profile, MaskAhead achieves $1.23\times$ end-to-end and $1.68\times$ decode-stage speedups over dense inference.

Problem

• Block Diffusion Language Models
• KV Cache
• Memory Efficiency
• Generation Speed

Innovation

• KV Cache Eviction
• Block Diffusion Language Models
• Training-free
• Mask-Guided Selection
• Quantized Attention