LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2026
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866914470654115840 |
|---|---|
| author | Xi, Haocheng Singh, Harman Hu, Yuezhou Hooper, Coleman Tiwari, Rishabh Tomar, Aditya Lee, Minjae Kang, Wonjun Mahoney, Michael Xu, Chenfeng Keutzer, Kurt Gholami, Amir |
| author_facet | Xi, Haocheng Singh, Harman Hu, Yuezhou Hooper, Coleman Tiwari, Rishabh Tomar, Aditya Lee, Minjae Kang, Wonjun Mahoney, Michael Xu, Chenfeng Keutzer, Kurt Gholami, Amir |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2604_12056 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models Xi, Haocheng Singh, Harman Hu, Yuezhou Hooper, Coleman Tiwari, Rishabh Tomar, Aditya Lee, Minjae Kang, Wonjun Mahoney, Michael Xu, Chenfeng Keutzer, Kurt Gholami, Amir Computation and Language Machine Learning 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. |
| title | LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models |
| topic | Computation and Language Machine Learning |
| url | https://arxiv.org/abs/2604.12056 |