LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models

• URL: http://arxiv.org/abs/2604.12056v1
• Date: Mon, 13 Apr 2026 20:53:51 GMT
• Title: LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models
• Abstract summary: Block-wise diffusion language models (DLMs) generate multiple tokens in any order, offering a promising alternative to the autoregressive decoding pipeline.<n>Naive sparse attention fails on DLMs due to a KV Inflation problem, where different queries select different prefix positions.<n>We propose LOSA (Locality-aware Sparse Attention), which reuses cached prefix-attention results for stable tokens and applies sparse attention only to active tokens.
• Score: 49.93891888238178
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Stability Implies Redundancy: Delta Attention Selective Halting for Efficient Long-Context Prefilling [14.757648323881883]
  Prefilling computational costs pose a significant bottleneck for Large Language Models.<n>Delta Attention Selective Halting (DASH) monitors the layer-wise update dynamics of the self-attention mechanism to selectively halt stabilized tokens.
  arXiv  Detail & Related papers  (2026-04-20T11:20:03Z)
• STDec: Spatio-Temporal Stability Guided Decoding for dLLMs [49.55447757907809]
  Large Diffusion Language Models (dLLMs) have achieved rapid progress, viewed as a promising alternative to the autoregressive paradigm.<n>Most dLLM decoders still adopt a global confidence threshold, and do not model local context from neighboring decoded states or temporal consistency of predicted token IDs.<n>We observe strong stability neartemporal stability in dLLM decoding, and our approach includes spatial-aware decoding and temporal-aware decoding.
  arXiv  Detail & Related papers  (2026-04-07T18:13:31Z)
• DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention [2.7422645382944935]
  We present DyLLM, a training-free inference framework that accelerates decoding by selectively computing only salient tokens.<n>Across diverse reasoning and code-generation benchmarks, DyLLM achieves up to 9.6x higher throughput.
  arXiv  Detail & Related papers  (2026-03-09T07:02:01Z)
• Stateful Token Reduction for Long-Video Hybrid VLMs [69.6930118088911]
  We study query-conditioned token reduction for hybrid video vision-language models (VLMs)<n>We propose a low-to-high progressive reduction schedule and a unified language-aware scoring mechanism for both attention and Mamba blocks.<n>Under an aggressive compression setting, our approach delivers substantial prefilling speedups with near-baseline accuracy at test time.
  arXiv  Detail & Related papers  (2026-02-27T08:11:06Z)
• MAGE: All-[MASK] Block Already Knows Where to Look in Diffusion LLM [9.69241599043101]
  Block diffusion LLMs are emerging as a promising next paradigm for language generation, but their use of KV caching makes memory access a dominant bottleneck in long-context settings.<n>This work identifies a key opportunity unique to block diffusion: attention at the first All-[MASK] denoising step reliably predicts important KV entries and budget requirements.<n>MAGE achieves near-lossless accuracy with a fraction of the KV budget while delivering up to 3-4x end-to-end speedup.<n>A lightweight fine-tuning strategy further strengthens [MASK]-guided patterns with minimal cost, requiring only a few hours of training on
  arXiv  Detail & Related papers  (2026-02-15T16:07:51Z)
• Attention Is All You Need for KV Cache in Diffusion LLMs [36.94369617373333]
  Elastic-Cache performs adaptive, layer-aware cache updates for diffusion large language models.<n>Our method achieves significantly higher throughput ($6.8times$ on GSM8K) than existing confidence-based approaches.
  arXiv  Detail & Related papers  (2025-10-16T17:59:48Z)
• NOSA: Native and Offloadable Sparse Attention [27.551376861663556]
  We present NOSA, a trainable sparse attention framework designed to support KV cache offloading.<n>We show that NOSA preserves near-lossless performance while achieving up to a 2.3x improvement in decoding throughput.
  arXiv  Detail & Related papers  (2025-10-15T14:33:16Z)
• Sparse-dLLM: Accelerating Diffusion LLMs with Dynamic Cache Eviction [72.27673320976933]
  Diffusion Large Language Models (dLLMs) enable breakthroughs in reasoning and parallel decoding.<n>Current caching techniques accelerate decoding by storing full-layer states, yet impose substantial memory usage.<n>We propose Sparse-dLLM, the first training-free framework integrating dynamic cache eviction with sparse attention.
  arXiv  Detail & Related papers  (2025-08-04T16:14:03Z)
• dKV-Cache: The Cache for Diffusion Language Models [53.85291644298835]
  Diffusion Language Models (DLMs) have been seen as a promising competitor for autoregressive language models.<n>We propose a KV-cache-like mechanism, delayed KV-Cache, for the denoising process of DLMs.<n>Our approach is motivated by the observation that different tokens have distinct representation dynamics throughout the diffusion process.
  arXiv  Detail & Related papers  (2025-05-21T17:32:10Z)
• Dialogue Without Limits: Constant-Sized KV Caches for Extended Responses in LLMs [6.222287867011644]
  We propose MorphKV, an inference-time technique that maintains a constant-sized KV cache while preserving accuracy.<n>Unlike retention or lossy compression, MorphKV iteratively refines the KV cache via lightweight updates guided by attention patterns of recent tokens.<n>Our studies show 52.9$%$ memory savings and 18.2$%$ higher accuracy on average compared to state-of-the-art prior works.
  arXiv  Detail & Related papers  (2025-03-02T18:12:50Z)
• Efficient Streaming Language Models with Attention Sinks [72.20260088848987]
  StreamingLLM is an efficient framework that enables Large Language Models to generalize to infinite sequence lengths without any fine-tuning. We show that StreamingLLM can enable Llama-2, MPT, Falcon, and Pythia to perform stable and efficient language modeling with up to 4 million tokens and more.
  arXiv  Detail & Related papers  (2023-09-29T17:59:56Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.