LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models
On the map: Diffusion models → Discrete and text diffusion
Abstract (arXiv)
Block-wise diffusion language models (DLMs) generate multiple tokens in any order, offering a promising alternative to the autoregressive decoding pipeline. However, they still remain bottlenecked by memory-bound attention in long-context scenarios. Naive sparse attention fails on DLMs due to a KV Inflation problem, where different queries select different prefix positions, making the union of accessed KV pages large. To address this, we observe that between consecutive denoising steps, only a small fraction of active tokens exhibit significant hidden-state changes, while the majority of stable tokens remain nearly constant. Based on this insight, we propose LOSA (Locality-aware Sparse Attention), which reuses cached prefix-attention results for stable tokens and applies sparse attention only to active tokens. This substantially shrinks the number of KV indices that must be loaded, yielding both higher speedup and higher accuracy. Across multiple block-wise DLMs and benchmarks, LOSA preserves near-dense accuracy while significantly improving efficiency, achieving up to +9 points in average accuracy at aggressive sparsity levels while maintaining 1.54x lower attention density. It also achieves up to 4.14x attention speedup on RTX A6000 GPUs, demonstrating the effectiveness of the proposed method.
Key facts. Cited by 2 papers in the map (data through 14 Oct 2026). Placed in Discrete and text diffusion, one of 14 directions in Diffusion models. Builds on 39 of its 44 references that are in the map. First posted on arXiv 13 Apr 2026.
Builds on
39 of 44 references.
Closest ancestors · 7
Same direction, newest first
- 2025SDAR: A Synergistic Diffusion-AutoRegression Paradigm for Scalable Sequence Generation7 Oct
- 2025SparseD: Sparse Attention for Diffusion Language Models28 Sep
- 2025Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models8 Sep
- 2025Dream 7B: Diffusion Large Language Models21 Aug
- 2025Large Language Diffusion Models14 Feb
- 2024Simple and Effective Masked Diffusion Language Models11 Jun
- 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution25 Oct
Other references · 32
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
- 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,100
- 2018HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answeringcited by 3,754
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2020Longformer: The Long-Document Transformercited by 3,413
- 2017TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehensioncited by 2,955
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,186
- 2019BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questionscited by 2,183
- 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 2,022
- 2021Structured Denoising Diffusion Models in Discrete State-Spacescited by 1,719
- 2020Big Bird: Transformers for Longer Sequencescited by 1,684
20 more
- 2020Reformer: The Efficient Transformercited by 1,668
- 2020Rethinking Attention with Performerscited by 1,601
- 2019Generating Long Sequences with Sparse Transformerscited by 1,546
- 2023LongBench: A Bilingual, Multitask Benchmark for Long Context Understandingcited by 1,417
- 2024SnapKV: LLM Knows What You are Looking for Before Generationcited by 880
- 2023H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Modelscited by 832
- 2017The NarrativeQA Reading Comprehension Challengecited by 597
- 2024Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inferencecited by 501
- 2021A Dataset of Information-Seeking Questions and Answers Anchored in Research Paperscited by 486
- 2023Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Timecited by 439
- 2024PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funnelingcited by 432
- 2025Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decodingcited by 420
- 2024MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attentioncited by 383
- 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Servingcited by 258
- 2024SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Accelerationcited by 160
- 2024SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantizationcited by 114
- 2025SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Trainingcited by 66
- 2025Sparse-dLLM: Accelerating Diffusion LLMs with Dynamic Cache Evictioncited by 65
- 2025ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMscited by 43
- 2025SageAttention2++: A More Efficient Implementation of SageAttention2cited by 22
Built on by
Citing papers per year: 2026: 2.
2 papers.
Similar, not cited
Close in wording, no citation either way
- 2026Efficient Long-Context Modeling in Diffusion Language Models via Block Approximate Sparse Attention
- 2026Retrofitting Linear Attention into Diffusion Language Models
- 2026AB-Sparse: Sparse Attention with Adaptive Block Size for Accurate and Efficient Long-Context Inference
- 2025Progressive Sparse Attention: Algorithm and System Co-design for Efficient Attention in LLM Serving
- 2025Blockwise SFT for Diffusion Language Models: Reconciling Bidirectional Attention and Autoregressive Decoding
- 2026Mask-Guided KV Cache Eviction in Block Diffusion Language Models
- 2026Fixed State, Long Reach: What a Constant-Size Cache Buys Block Diffusion at Scale
- 2025Rectified Sparse Attention