Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

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

About

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.

Omin Kwon, Yeonjae Kim, Doyeon Kim, Minseo Kim, Yeonhong Park, Jae W. Lee• 2026

Related benchmarks

TaskDatasetResultRank
Needle-In-A-Haystack RetrievalNeedle-in-a-Haystack 32K context (test)
Accuracy70
30
Needle-In-A-Haystack RetrievalNeedle-in-a-Haystack 8K context (test)
Accuracy100
30
Peak decode throughputLLaDA mini 16B 8K context 2.0
Throughput (tokens/s)1.37e+3
4
Peak decode throughputLLaDA 16B 16K context 2.0-mini
Throughput (tokens/s)796
4
Peak decode throughputLLaDA mini 16B (32K context) 2.0
Throughput (tokens/s)390
4
Peak decode throughputSDAR-8B-Chat 8K context
Throughput (tokens/s)414
4
Peak decode throughputSDAR-8B-Chat 16K context
Throughput (tokens/s)238
4
Peak decode throughputSDAR-8B-Chat 32K context
Throughput (tokens/s)121
4
Showing 8 of 8 rows

Other info

Follow for update