# MAGE: All-[MASK] Block Already Knows Where to Look in Block Diffusion LLM

arXiv: https://arxiv.org/abs/2602.14209 · HTML: https://mlsanity.com/paper/2602.14209/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 26 references that are in the map.
- First posted on arXiv 15 Feb 2026.

## Builds on (most cited first, 25 of 25)

- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference](https://mlsanity.com/paper/2406.10774/) (2024)
- [Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time](https://mlsanity.com/paper/2305.17118/) (2023)
- [PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling](https://mlsanity.com/paper/2406.02069/) (2024)
- [Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding](https://mlsanity.com/paper/2505.22618/) (2025)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [dKV-Cache: The Cache for Diffusion Language Models](https://mlsanity.com/paper/2505.15781/) (2025)
- [Fast-dLLM v2: Efficient Block-Diffusion LLM](https://mlsanity.com/paper/2509.26328/) (2025)
- [SDAR: A Synergistic Diffusion-AutoRegression Paradigm for Scalable Sequence Generation](https://mlsanity.com/paper/2510.06303/) (2025)
- [SparQ Attention: Bandwidth-Efficient LLM Inference](https://mlsanity.com/paper/2312.04985/) (2023)
- [Diffusion LLMs Can Do Faster-Than-AR Inference via Discrete Diffusion Forcing](https://mlsanity.com/paper/2508.09192/) (2025)
- [Loki: Low-rank Keys for Efficient Sparse Attention](https://mlsanity.com/paper/2406.02542/) (2024)
- [Twilight: Adaptive Attention Sparsity with Hierarchical Top-p Pruning](https://mlsanity.com/paper/2502.02770/) (2025)
- [LongLLaDA: Unlocking Long Context Capabilities in Diffusion LLMs](https://mlsanity.com/paper/2506.14429/) (2025)
- [Tactic: Adaptive Sparse Attention with Clustering and Distribution Fitting for Long-Context LLMs](https://mlsanity.com/paper/2502.12216/) (2025)
- [CDLM: Consistency Diffusion Language Models For Faster Sampling](https://mlsanity.com/paper/2511.19269/) (2025)
- [SparseD: Sparse Attention for Diffusion Language Models](https://mlsanity.com/paper/2509.24014/) (2025)
- [Mask Tokens as Prophet: Fine-Grained Cache Eviction for Efficient dLLM Inference](https://mlsanity.com/paper/2510.09309/) (2025)
- [Beyond Next-Token Prediction: A Performance Characterization of Diffusion versus Autoregressive Language Models](https://mlsanity.com/paper/2510.04146/) (2025)

## Built on by (most cited first, 2 of 2)

- [Hardware Acceleration of Block-Diffusion LLM for Edge Devices](https://mlsanity.com/paper/2609.01084/) (2026)
- [Mask-Guided KV Cache Eviction in Block Diffusion Language Models](https://mlsanity.com/paper/2610.06996/) (2026)

## Abstract (arXiv)

Block diffusion LLMs are an emerging paradigm for parallel language generation, but their KV caching makes memory access the dominant bottleneck in long-context inference. Sparse attention, which attends only to a small KV subset per query, can reduce this latency with minimal accuracy loss. In block diffusion, however, the B tokens of each block must share a single KV subset, and we show this per-block constraint degrades existing sparse KV estimators by up to 25% in recall. We address this challenge by exploiting a property that emerges from the block-diffusion training objective: it aligns the block-average query across denoising steps, so the All-[MASK] block at the first step already reveals the per-block KV subset for the entire trajectory. We exploit this in MAGE ([MASK]-Guided Sparse Attention), a training-free method that runs one exact attention pass at the first step and reuses its top-k index sets for all remaining steps within the block. Across three block-diffusion families on LongBench, MAGE matches Exact Attention at k=512 with near-lossless accuracy, achieves up to 6.82x end-to-end speedup at 128K context, and runs up to 3.35x and 2.28x faster than Quest and SparseD, designed for AR LLMs and fully bidirectional diffusion LLMs, respectively.