Block diffusion language models keep a massive key‑value (KV) cache during generation and attend to the entire cache at every denoising step, which limits memory usage and generation speed. Reducing these costs requires deciding, for the current block, which past tokens to use for denoising (selection) and which tokens to retain for future blocks (eviction).
We introduce MaskAhead, a training‑free approach that solves both tasks with a single mask‑query ranking mechanism. The current‑block mask scores KV entries by their estimated contribution to the attention output, enabling selection; probes of upcoming masked blocks similarly rank KV entries to determine eviction.
A quantized variant, Q‑MaskAhead, computes selection and attention directly from low‑bit KV while preserving most 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 only a 1.2‑point mean F1 drop, while Q‑MaskAhead pushes 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.
Review