Research map / Architectures and efficiency / Attention and transformers

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

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes MAGE, a training-free sparse attention method for block diffusion LLMs that runs one exact attention pass on the all-[MASK] block at step 1 and reuses its per-block top-k KV indices for all remaining denoising steps, with selection overlapped on a second CUDA stream.

Problem

Sparse KV attention for block diffusion LLMs is harder because all B tokens in a block must share one per-block KV subset, and existing estimators (Quest, SparseD) lose up to 25% recall under this constraint, making long-context decoding memory-bound.

Approach

Block-diffusion training aligns the block-average query across denoising steps, so the all-[MASK] block at the first step reveals the per-block oracle top-k KV subset for the whole trajectory; MAGE computes this oracle once with exact attention (selection on an async CUDA stream) and reuses it for all later steps.

Key results

• At k=512, MAGE averages drops of 0.98/0.87/0.28 points vs exact on Fast-dLLM v2 7B / SDAR-8B-Chat / LLaDA 2.0 mini, versus Quest's 6.85/5.67/4.50 and SparseD's 2.27/2.80/1.13.
• At 128K context with k=1024 and 32 steps/block, MAGE achieves up to 6.82x end-to-end speedup over exact attention on SDAR-8B-Chat, 5.86x on Fast-dLLM v2 7B, and 3.67x on LLaDA 2.0 mini.
• MAGE runs up to 3.35x faster than Quest and 2.28x faster than SparseD end-to-end; first-step overhead is +1.9% to +6.4% over exact.

Surprising findings

• Block-diffusion training raises top-k recall of the all-[MASK] oracle over the AR pre-training backbone (e.g., 0.769 to 0.851 for LLaDA2.0-mini vs Ling-mini-base-2.0), suggesting the anchoring is training-induced.

Limitations

• Benefit diminishes in short-context regimes where FFN and non-attention computation dominate latency.
• Reduces KV cache access time but not the KV cache memory footprint.
• Evaluated on only three block-diffusion model families and LongBench/NIAH.

Experimental setup

Datasets: LongBench Needle-in-a-Haystack

Benchmarks: LongBench Needle-in-a-Haystack

Metrics: F1 ROUGE-L EditSim retrieval score top-k recall TPOB speedup

Baselines: Exact Attention (FlashInfer) Quest SparseD

Builds on: Quest SparseD FlashInfer NTK-based RoPE extrapolation

Closest papers by meaning

• From Next-Token to Next-Block: A Principled Adaptation Path for Diffusion LLMs (NeurIPS 2026 · similarity 0.91)
• UltraLLaDA: Scaling the Context Length to 128K for Diffusion Large Language Models (ICLR 2026 · similarity 0.89)
• Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models (NeurIPS 2026 · similarity 0.89)
• Polestar: Drift-Aware Cache Calibration and Token Commitment for Efficient Inference of Diffusion LLMs (NeurIPS 2026 · similarity 0.89)
• dKV-Cache: The Cache for Diffusion Language Models (NeurIPS 2025 · similarity 0.89)
• Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models (ICLR 2025 · similarity 0.89)