Mask-Guided KV Cache Eviction in Block Diffusion Language Models
cs.LG, cs.AI, cs.CL
Submitted: 2026-10-04
Updated: 2026-10-04
Code: https://github.com/brain-lab-research/MaskAhead
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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.
Terminology
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.
Sources
- LookaheadKV: Fast and Accurate KV Cache Eviction by Glimpsing into the Future without Generation
- KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization
- Mask Tokens as Prophet: Fine-Grained Cache Eviction for Efficient dLLM Inference
- ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMs
- KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
- NPU Design for Diffusion Language Model Inference
- GroupKV: Hierarchical KV Cache Management for Long-Context Diffusion LLM Inference
- SparseD: Sparse Attention for Diffusion Language Models
- Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding
- DreamReasoner-8B: Block-Size Curriculum Learning for Diffusion Reasoning Models
- Efficient Streaming Language Models with Attention Sinks
- More Tokens, Lower Precision: Towards the Optimal Token-Precision Trade-off in KV Cache Compression
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks