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

TL;DR

LOSA introduces a locality-aware sparse attention mechanism that reuses cached attention results for stable tokens, significantly improving efficiency and accuracy in block-wise diffusion language models.

Contribution

It proposes a novel sparse attention method that addresses KV Inflation by caching stable token representations, enhancing speed and accuracy in long-context DLMs.

Findings

Achieves up to +9 points in accuracy at high sparsity levels

Provides up to 4.14x attention speedup on GPUs

Maintains near-dense accuracy with reduced attention density

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.