💬 NLP

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

The paper proposes LoSA, a locality-aware sparse attention mechanism that accelerates block-wise diffusion language models by reusing cached attention results for stable tokens and applying sparse attention only to active tokens, thereby overcoming KV inflation to achieve significant speedups and improved accuracy in long-context scenarios.

Original authors: Haocheng Xi, Harman Singh, Yuezhou Hu, Coleman Hooper, Rishabh Tiwari, Aditya Tomar, Minjae Lee, Wonjun Kang, Michael Mahoney, Chenfeng Xu, Kurt Keutzer, Amir Gholami

Published 2026-04-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Haocheng Xi, Harman Singh, Yuezhou Hu, Coleman Hooper, Rishabh Tiwari, Aditya Tomar, Minjae Lee, Wonjun Kang, Michael Mahoney, Chenfeng Xu, Kurt Keutzer, Amir Gholami

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a long story, but instead of writing one word at a time (like a normal person), you are writing in blocks of 16 words at once. You can rearrange those 16 words in any order you want before finalizing them. This is how Block-wise Diffusion Language Models (DLMs) work. They are powerful because they can "think" about a whole chunk of text simultaneously, which is great for complex reasoning.

However, there's a big problem: Memory Traffic.

The Problem: The "Crowded Bus" Effect

To write these blocks, the model needs to remember everything it wrote previously (the "prefix"). In computer terms, this is called the KV Cache.

When the model tries to write a new block, it has to look back at all those previous words to decide what comes next.

The Naïve Approach: Imagine you have a bus with 16 passengers (the 16 words in your new block). Each passenger wants to look at a specific seat in the back of the bus (the previous text) to decide what to say.
The Inflation Problem: If Passenger 1 looks at Seat 5, Passenger 2 looks at Seat 100, and Passenger 3 looks at Seat 500, the bus driver has to stop and open the doors for all those different seats to let everyone check. Even though each person only looked at one seat, the bus had to load 16 different seats just to satisfy the group. This is called KV Inflation. It's slow and wastes a lot of energy.

The Insight: The "Quiet Classroom"

The researchers noticed something interesting about how these models work. When the model updates a block of words, most of the words don't actually change much.

Think of it like a classroom:

Active Tokens: A few students are raising their hands, shouting, and changing their answers. They are "active."
Stable Tokens: The other 14 students are sitting quietly. Their answers haven't changed since the last time you checked. They are "stable."

If you ask the quiet students, "What do you think about the story so far?" they will give you the exact same answer they gave you a second ago. There is no need to ask them again!

The Solution: LoSA (Locality-Aware Sparse Attention)

The paper introduces a method called LoSA (Locality-Aware Sparse Attention). It's like a smart teacher who manages the classroom efficiently.

Here is how LoSA works, step-by-step:

Spot the Active Students: First, LoSA quickly checks which students (tokens) are actually changing their minds. It ignores the quiet ones.
Reuse the Quiet Answers: For the "Stable Tokens" (the quiet students), LoSA doesn't ask them to look at the back of the bus again. It just reuses the answer they gave last time. It's like saying, "You didn't change your mind, so I'll just write down what you said before."
Focus on the Active Ones: LoSA only sends the "Active Tokens" (the ones changing) to look at the back of the bus.
The Result: Because only a few students are looking, the bus driver only has to open the doors for a small group of seats.

Why This is a Game Changer

Speed: By not loading the whole bus of seats for every single word, the model runs 4.14 times faster on standard graphics cards.
Smarter: Because the "quiet" students still get to use their full memory of the story (instead of being forced to guess based on a tiny slice of memory), the model actually writes better stories (higher accuracy) than other methods that try to be too sparse.
Efficiency: It uses 1.5 times less memory traffic, meaning the computer doesn't have to work as hard to move data around.

The Analogy Summary

Old Way (Naïve Sparse Attention): Asking 16 people to look at 16 different random spots in a library. The librarian has to run to 16 different shelves to get the books. It's chaotic and slow.
LoSA: Realizing that 14 of those people are just reading the same page they were reading a second ago. You tell them, "Keep reading what you have." You only send the 2 people who are turning pages to the librarian. The librarian only has to run to 2 shelves.

In short: LoSA makes AI models faster and smarter by realizing that most of the time, the AI is just repeating what it already knows, so it shouldn't waste energy checking the same things over and over again.

1. Problem Statement

Block-wise Diffusion Language Models (DLMs) generate text by updating blocks of tokens in any order via iterative denoising, offering advantages for complex reasoning over autoregressive models. However, they face a critical bottleneck in long-context scenarios: Memory-Bound Attention.

The KV Inflation Problem: In autoregressive models, sparse attention reduces latency by limiting the Key-Value (KV) cache accessed per query. However, in block-wise DLMs, a block contains $B$ queries processed simultaneously. If a naive sparse attention method (like QUEST) is applied, each of the $B$ queries selects a different subset of $k$ KV positions.
The Consequence: The hardware must load the union of all selected KV positions across the entire block. This union size ( $L'$ ) can grow up to $B \times k$ , effectively negating the memory savings intended by sparsity. This phenomenon, termed KV Inflation, leads to high memory traffic and low speedup, even with aggressive sparsity settings.

2. Key Observation: Locality of Representation Changes

The authors identify a unique structural property in block-wise diffusion that is absent in autoregressive decoding: Locality of Representation Changes.

Active vs. Stable Tokens: Between consecutive denoising steps ( $t-1$ and $t$ ), only a small fraction of tokens in the current block undergo significant changes in their hidden states (transitioning from [MASK] to actual tokens). These are termed Active Tokens.
Stability: The majority of tokens (Stable Tokens) exhibit negligible changes in their query vectors ( $q$ ) and hidden states. Consequently, their attention patterns and outputs remain approximately constant across steps.
Implication: Since stable tokens do not change significantly, their attention results from the previous step can be reused, eliminating the need to recompute sparse attention for them.

3. Methodology: LoSA (Locality-Aware Sparse Attention)

LoSA leverages the observation above to mitigate KV Inflation by reducing the number of queries that participate in the sparse selection process.

Core Workflow

Locality Pruning:
- At each denoising step, the model calculates the Mean Squared Error (MSE) between the query vectors of the current step ( $t$ ) and the previous step ( $t-1$ ) for every token in the block.
- Tokens are ranked by this "locality score." The top- $k$ tokens with the largest changes are identified as the Active Set ( $A$ ). The remaining tokens are Stable Tokens.
Sparse Attention for Active Tokens:
- A sparse selector (e.g., QUEST) is applied only to the active tokens in set $A$ .
- The union of KV indices selected by these active tokens ( $I = \bigcup_{i \in A} I_i$ ) is computed. Since $|A| \ll B$ , the size of this union is significantly smaller than the union of all $B$ queries.
- The model loads the KV cache for indices in $I$ and computes the prefix attention for active tokens.
Cache Reuse for Stable Tokens:
- For stable tokens (those not in $A$ ), the model reuses the cached prefix attention outputs ( $o_p$ ) and log-normalizers ( $L_p$ ) from the previous step.
- No new KV cache loading or attention computation is performed for these tokens.
Merging:
- The newly computed attention for active tokens and the reused attention for stable tokens are merged with the within-block (suffix) attention using Online Softmax to produce the final output.
Initialization:
- For the first denoising step of a block, dense prefix attention is used to initialize the cache, ensuring high accuracy before the reuse mechanism begins.

Efficiency Analysis

Overhead: The cost of calculating locality scores and ranking tokens is $O(B \cdot d + B \log B)$ , which is negligible compared to the attention computation.
Memory: The additional memory required to cache attention outputs is minimal ( $\sim 2$ KB per head for typical settings) compared to the full KV cache.
KV Load Reduction: By shrinking the query set from $B$ to $|A|$ , LoSA significantly reduces the union size of KV indices loaded, directly lowering memory bandwidth usage.

4. Key Contributions

Identification of KV Inflation: The paper formally identifies and analyzes the "KV Inflation" problem, explaining why standard sparse attention fails in block-wise DLMs due to the union effect of multiple queries.
Discovery of Representation Locality: The authors demonstrate that hidden state changes in block-wise diffusion are highly localized, with only a small fraction of tokens changing significantly between steps.
Proposal of LoSA: A novel sparse attention mechanism that dynamically separates tokens into active and stable groups, reusing cached results for stable tokens to drastically reduce KV cache traffic.
Comprehensive Evaluation: Extensive experiments across multiple models (Trado-8B/4B, SDAR-8B) and benchmarks (LongBench, reasoning tasks) validating the method's effectiveness.

5. Experimental Results

The authors evaluated LoSA against baselines like QUEST (naive sparse attention), SparseD, and Sparse-dLLM.

Accuracy:
- LoSA preserves near-dense accuracy even at aggressive sparsity levels.
- On LongBench (Trado-8B), LoSA achieved +10.43% average accuracy improvement over QUEST at a low budget (128).
- At budget 256, LoSA maintained a +7.50% advantage.
- LoSA consistently outperformed baselines across different block sizes (16 and 32) and model scales (4B and 8B).
Efficiency (KV Density):
- LoSA achieved 1.54× lower attention density (percentage of KV cache loaded) compared to QUEST on average.
Speedup:
- On NVIDIA RTX A6000, LoSA achieved up to 4.14× speedup over dense attention, compared to 3.26× for QUEST.
- On newer RTX 5090 hardware, the speedup trend held at 3.67×, confirming the method is memory-bound and scalable.
Short-Context Performance: On short-context benchmarks (HellaSwag, WinoGrande), LoSA performed comparably to QUEST, as the KV Inflation problem is less severe in short contexts.

6. Significance

LoSA represents a significant advancement in the efficient inference of Diffusion Language Models.

Solves a Fundamental Bottleneck: It directly addresses the memory-bound nature of block-wise DLMs, which previous sparse attention methods failed to solve due to the KV Inflation problem.
Accuracy-Efficiency Trade-off: Unlike many sparse methods that sacrifice accuracy for speed, LoSA improves accuracy at high sparsity levels by preserving full dense attention information for stable tokens while only approximating for the changing ones.
Hardware Agnostic: The method relies on reducing memory traffic, making it highly effective across different GPU generations (A6000 to RTX 5090) as memory bandwidth remains the primary constraint.

In summary, LoSA enables block-wise diffusion models to scale to long contexts efficiently by exploiting the temporal stability of token representations, offering a viable path toward high-performance, non-autoregressive text generation.