MAGE: All-[MASK] Block Already Knows Where to Look in Block Diffusion LLM
Not on the map. Why?
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.
Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 25 of its 26 references that are in the map. First posted on arXiv 15 Feb 2026.
Builds on
25 of 26 references.
References · 25
- 2020Longformer: The Long-Document Transformercited by 3,362
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,115
- 2020Big Bird: Transformers for Longer Sequencescited by 1,673
- 2023LongBench: A Bilingual, Multitask Benchmark for Long Context Understandingcited by 1,378
- 2025Large Language Diffusion Modelscited by 872
- 2024SnapKV: LLM Knows What You are Looking for Before Generationcited by 848
- 2023H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Modelscited by 805
- 2024Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inferencecited by 487
- 2023Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Timecited by 428
- 2024PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funnelingcited by 420
- 2025Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decodingcited by 404
- 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Servingcited by 248
13 more
- 2025dKV-Cache: The Cache for Diffusion Language Modelscited by 140
- 2025Fast-dLLM v2: Efficient Block-Diffusion LLMcited by 128
- 2023SparQ Attention: Bandwidth-Efficient LLM Inferencecited by 119
- 2025SDAR: A Synergistic Diffusion-AutoRegression Paradigm for Scalable Sequence Generationcited by 118
- 2025Diffusion LLMs Can Do Faster-Than-AR Inference via Discrete Diffusion Forcingcited by 111
- 2024Loki: Low-rank Keys for Efficient Sparse Attentioncited by 87
- 2025Twilight: Adaptive Attention Sparsity with Hierarchical Top-p Pruningcited by 48
- 2025LongLLaDA: Unlocking Long Context Capabilities in Diffusion LLMscited by 38
- 2025Tactic: Adaptive Sparse Attention with Clustering and Distribution Fitting for Long-Context LLMscited by 24
- 2025CDLM: Consistency Diffusion Language Models For Faster Samplingcited by 20
- 2025Mask Tokens as Prophet: Fine-Grained Cache Eviction for Efficient dLLM Inferencecited by 17
- 2025SparseD: Sparse Attention for Diffusion Language Modelscited by 16
- 2025Beyond Next-Token Prediction: A Performance Characterization of Diffusion versus Autoregressive Language Modelscited by 13
Built on by
Citing papers per year: 2026: 1.
1 paper.