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

2026-04-13

AI Summary

This work addresses the memory bottleneck of dense attention mechanisms in block diffusion language models under long-context scenarios, where conventional sparse attention fails due to KV cache inflation. The authors propose Locality-aware Sparse Attention (LoSA), which exploits the observation that during denoising, the hidden states of most tokens change minimally. LoSA dynamically identifies active tokens and applies sparse attention only to them, while reusing cached prefix attention results for stable tokens. This approach substantially reduces the number of KV indices that must be loaded, achieving accuracy close to dense attention across multiple block diffusion models and benchmarks—improving average accuracy by up to 9 points—while reducing attention density by 1.54× and delivering up to 4.14× speedup on an RTX A6000 GPU.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• block-wise diffusion language models
• long-context
• memory-bound attention
• sparse attention
• KV inflation

Innovation

• Locality-aware Sparse Attention
• Block-wise Diffusion Language Models
• KV Inflation
• Attention Sparsity
• Token Stability