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

Haocheng Xi; Harman Singh; Yuezhou Hu; Coleman Hooper; Rishabh Tiwari; Aditya Tomar; Minjae Lee; Wonjun Kang

cs.CLApr 13, 2026

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

Authors: Haocheng Xi, Harman Singh, Yuezhou Hu, Coleman Hooper, Rishabh Tiwari, Aditya Tomar, Minjae Lee, Wonjun Kang,

Organizations: University of California, Berkeley · FuriosaAI · ICSI · LBNL ·

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.

Figures & tables

Figure 1: Illustration of the KV-union effect: in block diffusion, each query selects a small set of prefix KV positions, but the effective cost is determined by the union across the block, inflating KV access. Our method computes sparse attention only for selected (active) tokens during denoising and reuses cached results for the others (static tokens), substantially reducing the size of the KV union and the latency.

Panel of Figure 2 — Figure 2: Visualizing representation change locality across denoising steps. Left and middle plots show the L2 norm of query vectors at timesteps $t-1$ and $t$ respectively. The right plot shows the per-token MSE between query vectors across the two timesteps. Only a small fraction of tokens exhibit large changes, motivating the reuse of cached-prefix attention for stable tokens.

Figure 3: KV-cache load reduction with locality-aware sparse prefix attention. X-axis: number of active query tokens in the block; Y-axis: percentage of the full KV cache that must be loaded. Evaluated on Trado-8B-Instruct on TriviaQA with block size 16 and 64K context length. LoSA only loads prefix KV positions in the union

\mathcal{I}=\bigcup_{i\in\mathcal{A}}\mathcal{I}_{i}

for the active tokens set

\mathcal{A}

, rather than the union over all

B

queries. The dashed orange line shows the KV load under QUEST, which must load the union across all queries.

Figure 4: Visualizing locality across denoising steps. We plot the distribution of per-token MSE changes in queries between steps $t\!-\!1$ and $t$ . For each layer, we sort the tokens in decreasing order of MSE changes from top to bottom in this heatmap. We observe that only a small fraction of tokens exhibit large changes (we show a threshold line which corresponds to 50% of the total MSE change is exhibited by tokens above that line), motivating the reuse of cached-prefix attention for the stable tokens.

Figure 5: Overview of the locality-aware sparse prefix attention workflow. We (i) measure per-token query changes and select the changed-token set $\mathcal{A}$ , (ii) run QUEST to obtain per-token prefix indices and take their union $\mathcal{I}$ , (iii) load prefix KV for $\mathcal{I}$ and compute updated prefix statistics for queries in $\mathcal{A}$ while reusing cached statistics for stable tokens, (iv) compute dense within-block attention, and (v) merge prefix and block contributions via online softmax and cache the updated prefix statistics for the next denoising step.

KV Cache		LongBench Accuracy (Trado-8B-Instruct)
Per-Query Budget	Method	HotPotQA	TriviaQA	NarrativeQA	Qasper	MultiFieldQA	Average ( $\uparrow$ )
-	Dense	49.45%	84.79%	19.04%	17.75%	53.29%	44.86%
128	QUEST	29.17%	69.21%	7.46%	13.64%	38.21%	31.54%
128	LoSA	48.27%	83.79%	17.19%	15.58%	45.00%	41.97% (+10.43)
256	QUEST	32.95%	75.23%	8.50%	14.04%	42.46%	34.64%
256	LoSA	44.53%	81.82%	19.42%	17.11%	47.81%	42.14% (+7.50)

Table 1: Accuracy (%) on LongBench under different retrieval budgets for Trado-8B-Instruct . LoSA consistently outperforms QUEST across all budget settings, with particularly large gains at low budgets (e.g., +10.43 at budget 128). The

\Delta

shown is the improvement over QUEST. Full comparison including SparseD is in Table 6 .

KV Cache		LongBench KV-Cache Density
Per-Query Budget	Method	HotPotQA	TriviaQA	NarrativeQA	Qasper	MultiFieldQA	Average ( $\downarrow$ )
128	QUEST	2.93%	3.80%	1.90%	6.72%	5.87%	4.24%
128	LoSA	1.71%	2.29%	0.98%	4.15%	3.77%	2.58% ( $1.64\times$ )
256	QUEST	5.84%	7.07%	3.89%	12.47%	11.14%	8.08%
256	LoSA	3.37%	4.47%	2.05%	7.90%	7.43%	5.04% ( $1.60\times$ )
512	QUEST	10.74%	13.06%	7.24%	23.13%	19.71%	14.78%

Table 2: KV-cache density (percentage of prefix tokens selected) on LongBench under different budgets. On all datasets and budget configurations, LoSA achieves an average of

1.54\times

lower KV-cache density than the baseline.

KV Cache		LongBench Accuracy (SDAR-8B-Chat)
Per-Query Budget	Method	HotPotQA	TriviaQA	NarrativeQA	Qasper	MultiFieldQA	Average ( $\uparrow$ )
-	Dense	49.35%	85.72%	19.06%	18.25%	49.49%	44.37%
128	QUEST	27.31%	70.40%	6.44%	17.29%	41.40%	32.57%
128	LoSA	43.36%	80.32%	18.69%	15.63%	46.74%	40.95% (+8.38)
256	QUEST	32.57%	78.36%	8.73%	18.21%	43.20%	36.21%
256	LoSA	45.68%	83.16%	15.65%	13.17%	48.33%	41.20% (+4.99)

Table 3: Accuracy (%) on LongBench under different retrieval budgets for SDAR-8B-Chat . Results are consistent with Trado-8B (Table 1 ): LoSA consistently outperforms QUEST, with the largest gains at low budgets (+8.38 at budget 128). Full comparison including SparseD is in Table 6 .

		Block Size = 16						Block Size = 32
Budget	Method	HotPotQA	TriviaQA	NarrativeQA	Qasper	MFieldQA	Avg. ( $\uparrow$ )	HotPotQA	TriviaQA	NarrativeQA	Qasper	MFieldQA	Avg. ( $\uparrow$ )
-	Dense	38.16%	75.61%	18.80%	29.57%	35.42%	39.51%	36.22%	75.23%	19.58%	13.40%	31.52%	35.19%
128	QUEST	30.01%	69.14%	9.83%	26.05%	30.79%	33.16%	25.55%	70.83%	9.53%	11.20%	27.30%	28.88%
128	LoSA	33.65%	66.46%	14.78%	24.71%	30.07%	33.93%	28.62%	69.50%	12.15%	13.78%	28.25%	30.46%
256	QUEST	33.06%	73.49%	12.34%	29.25%	34.67%	36.56%	31.61%	72.10%	10.87%	12.85%	29.10%	31.31%
256	LoSA	36.33%	72.74%	15.99%	28.11%	34.27%	37.49%	32.14%	72.20%	15.99%	14.80%	28.68%	32.76%

Table 4: LongBench accuracy (%) for Trado-4B-Instruct under two block sizes. LoSA consistently outperforms QUEST across model sizes and block-size configurations, confirming that the locality-aware approach generalizes beyond the 8B scale.

Budget	Method	HellaSwag	WinoGrande	BoolQ	Avg.
–	Dense	69.60	63.60	72.92	68.71
128	QUEST	67.60	66.80	73.96	69.45
128	LoSA	70.80	66.00	72.92	69.91
256	QUEST	73.60	63.20	75.00	70.60
256	LoSA	70.40	65.20	73.96	69.85

Table 5: Accuracy (%) on three commonsense reasoning benchmarks for Trado-8B-Instruct. LoSA performs comparably to QUEST; the slight gap at budget 256 is within statistical variance on these short-context tasks (

<

1K tokens) where the KV Inflation problem is less severe.

Figure 6: Latency breakdown comparison on prefix attention using Trado-8B-Instruct on TriviaQA with RTX A6000. Left: 64K context length with block size 16. Right: 32K context length with block size 32. LoSA and QUEST (Adapted) achieve

4.14\times

and

3.26\times

speedup over Dense Attention, respectively, by reducing memory-bound attention operations through sparsity. LoSA further improves over QUEST through locality-aware reuse.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

KV Cache		LongBench Accuracy – Comparison with SparseD
Per-Query Budget	Method	HotPotQA	TriviaQA	NarrativeQA	Qasper	MultiFieldQA	Average ( $\uparrow$ )
Trado-8B-Instruct
-	Dense	49.45%	84.79%	19.04%	17.75%	53.29%	44.86%
128	QUEST	29.17%	69.21%	7.46%	13.64%	38.21%	31.54%
	SparseD	38.97%	64.64%	13.70%	7.95%	39.55%	32.96%
	LoSA	48.27%	83.79%	17.19%	15.58%	45.00%	41.97% (+9.01)

Appendix

Table 6: Full LongBench accuracy comparison including SparseD. The

\Delta

shown is the improvement of LoSA over the best competing method at each budget. At budget 1024 on Trado-8B and budget 512 on SDAR-8B, SparseD slightly outperforms LoSA in average accuracy.

Panel of Figure 7 — Figure 7: Visualizing representation change locality across denoising steps for query, keys and values of the token in the block that is being decoded. (Top) L2 norm of vectors at step $t\!-\!1$ . (Middle) L2 norm of vectors at step $t$ . (Bottom) Per-token change in representations, measured using MSE between steps $t\!-\!1$ and $t$ . Only a small fraction of tokens exhibit large changes, motivating reuse of cached-prefix attention for stable tokens.

Figure 8: Combined QKV Normalized MSE Heatmap with Threshold Selection ( $\tau=0.5$ ) for Data Sample 0. The red line indicates the boundary where cumulative normalized MSE reaches 50% of the layer total. Tokens above the line are selected; tokens below are dimmed.

Figure 9: Combined QKV Normalized MSE Heatmap with Threshold Selection ( $\tau=0.5$ ) for Data Sample 100.

Figure 10: Combined QKV Normalized MSE Heatmap with Threshold Selection ( $\tau=0.5$ ) for Data Sample 200.

Figure 11: Combined QKV Normalized MSE Heatmap with Threshold Selection ( $\tau=0.5$ ) for Data Sample 300.

Diffusion Language Model Inference

Diffusion language model inference studies how to generate, guide, and serve text efficiently through iterative denoising and parallel token updates.

See all 177 papers

Efficient Long-Context Modeling in Diffusion Language Models via Block Approximate Sparse Attention
May 19, 2026Wenhu Zhang, Yiming Wu, Huanyu Wang +6Self-Attention Long-Context Language Modeling
The Hong Kong University of Science and Technology · The University of Hong Kong · Zhejiang University +1
Retrofitting Linear Attention into Diffusion Language Models
Aug 6, 2026Jinha Kim, Younghun Roh, Jaeyeon KimBlockwise Diffusion Language Models Diffusion Language Model Inference
Apple · *Work done at MIT · Google DeepMind +1
PulseCol: Periodically Refreshed Column-Sparse Attention for Accelerating Diffusion Language Models
May 20, 2026Yanyi Lyu, Letian Chen, Futing Sun +3Self-Attention Diffusion Language Model Inference
Harbin Institute of Technology (Shenzhen)

May 19, 2026J/K move · Enter open · S save

Efficient Long-Context Modeling in Diffusion Language Models via Block Approximate Sparse Attention

Wenhu Zhang, Yiming Wu, Huanyu Wang, Yaoyang Liu+5

The Hong Kong University of Science and Technology · The University of Hong Kong · Zhejiang University+1

Self-Attention Long-Context Language Modeling Sparse Attention Diffusion Language Models Diffusion Model Inference Acceleration Block-Sparse Attention

Open paper AskPDF arXiv

Diffusion Language Models (DLMs) enable globally coherent, bidirectional, and controllable text generation, offering advantages over traditional autoregressive LLMs, while scaling to ultra-long sequences remains costly. Many existing block-sparse attention methods select blocks by fixed sampling patterns over the high-resolution attention space, such as tail regions or anti-diagonal stripes. Such prior-driven sampling can miss salient tokens and introduce instability under distribution shifts. In this paper, we propose the Block Approximate Sparse Attention framework (BA-Att) with block-wise pre-downsampled operation, which identifies informative regions within a compact downsampled space, avoiding reliance on brittle positional priors. To analyze its theoretical behavior, we define an oracle post-downsample attention map and formalize the approximation error between pre- and post-downsample schemes. Based on this insight, we introduce a lightweight norm-sorting module and a covariance-compensated correction that approximates full covariance using diagonal QK variances, reducing computational complexity. Extensive experiments show that our operator achieves up to 6.95x acceleration over FlashAttention in attention computation, and maintains near full-attention performance at 50% sparsity across language models, multimodal language models, and video generation models, demonstrating strong efficiency and generalization.

Abstract

Diffusion Language Model Inference

Explore similar work

Efficient Long-Context Modeling in Diffusion Language Models via Block Approximate Sparse Attention

Retrofitting Linear Attention into Diffusion Language Models

PulseCol: Periodically Refreshed Column-Sparse Attention for Accelerating Diffusion Language Models