Source-linked AI summary
SparseD: Sparse Attention for Diffusion Language Models
Zeqing Wang, Gongfan Fang, Xinyin Ma, Xingyi Yang, Xinchao Wang
TL;DR
DLMs face high inference latency from quadratic attention, while AR-oriented sparse patterns do not match DLM-specific head and denoising behavior. SparseD pre-computes head-specific patterns once, reuses them later, and keeps full attention early; it preserves accuracy with up to 1.50× speedup over FlashAttention at 64k context and 1,024 steps.
Problem
DLM inference is bottlenecked by quadratic attention, while AR sparse-attention patterns fail to capture DLM head-specific behavior and early-step generation requirements.
Method
SparseD pre-computes head-specific sparse patterns once, reuses them across denoising steps, and applies full attention during early steps.
Results
SparseD preserves accuracy with up to 1.50× speedup over FlashAttention at 64k context and 1,024 diffusion steps.
Takeaways & Limitations
SparseD provides a practical sparse-attention solution for efficient long-context DLM deployment without substantial generation-quality loss.
Takeaways & Limitations
The work focuses on algorithmic design and leaves system-level optimization and integration with cache-based methods for future work.
Abstract
from arXiv · showhide
While diffusion language models (DLMs) offer a promising alternative to autoregressive models (ARs), existing open-source DLMs suffer from high inference latency. This bottleneck is mainly due to the attention's quadratic complexity with respect to context length in computing all query-key pairs. Intuitively, to reduce this complexity, a natural strategy is to restrict attention to sparse patterns that retain only the most relevant connections. Such approaches are well-established in ARs, where attention follows fixed and clearly defined sparse patterns. However, in DLMs, we observe distinct sparsity behaviors: (1) attention patterns vary across heads, (2) attention patterns in each head remain highly similar across denoising steps, and (3) early denoising steps are critical for generation. These findings render sparse attention methods designed for ARs largely incompatible with DLMs, as they fail to capture head-specific structures and risk degrading generation when applied in early denoising steps. To address these challenges, we propose SparseD, a novel sparse attention method for DLMs. Leveraging the observations, SparseD only requires pre-computing head-specific sparse patterns one time, and reuses them across all steps. This prevents recomputing sparse patterns at each denoising step. Meanwhile, SparseD uses full attention in the early steps, then switches to sparse attention later to maintain generation quality. Together, these establish SparseD as a practical and efficient solution for deploying DLMs in long-context applications. Experimental results demonstrate that SparseD achieves lossless acceleration, delivering up to $1.50\times$ speedup over FlashAttention at a 64k context length with 1,024 denoising steps.
1 INTRODUCTION
DLMs offer parallel generation but face quadratic attention costs that create high latency, especially for long contexts. SparseD addresses DLM-specific attention behavior by reusing head-specific patterns after early full-attention steps, achieving near-lossless acceleration.
- 1 INTRODUCTION: DLM generation latency grows with context length because bidirectional attention computes all query–key pairs, producing quadratic complexity.This bottleneck affects both prompt prefill and generated tokens.
- 1 INTRODUCTION: DLM attention varies across heads, remains consistent within each head across denoising steps, and makes early steps crucial for generation.These properties make fixed autoregressive sparse patterns poorly suited to DLMs.
- 1 INTRODUCTION: SparseD pre-computes head-specific block-wise query–key patterns, reuses them across denoising steps, and applies full attention early to protect generation quality.The design avoids recomputing sparse patterns at every step while respecting the importance of early denoising.
- 1 INTRODUCTION: SparseD separately selects important prefill and generation-token scores so early-step selection does not overfocus on prefill tokens.Generation-token attention scores are initially lower but increase later, motivating isolated selection.
- 1 INTRODUCTION: 1.50× speedup over FlashAttention at 64k context with 1,024 diffusion steps, while greatly preserving original accuracy.Experiments on recent DLMs report negligible accuracy loss alongside the speedup.
2 RELATED WORKS
Related work positions DLMs as parallel, bidirectional alternatives to autoregressive models while highlighting inference latency as a major unresolved challenge. Prior acceleration approaches include fixed sparse attention for ARs and cache-based methods for DLMs.
- 2 RELATED WORKS: DLMs denoise all tokens jointly in a bidirectional manner, unlike autoregressive models that generate tokens sequentially.This parallel generation supports strong language understanding and generation performance.
- 2 RELATED WORKS: Sparse attention reduces DLM attention cost by restricting computation to selected query–key pairs, extending a mature strategy from autoregressive models.AR methods commonly use fixed patterns such as sink attention and sliding windows.
- 2 RELATED WORKS: Cache-based DLM accelerators reduce redundant computation by reusing key–value states or token blocks, but can noticeably degrade accuracy in long-context settings.The cited approaches include dKV-Cache and Fast-dLLM.
3 METHOD
SparseD exploits DLM-specific attention behavior: head patterns differ, remain stable across denoising steps, and are especially sensitive early in generation. It therefore uses early full attention, then reuses head-specific block-sparse patterns selected separately for prefill and generation tokens.
- 3.1 PRELIMINARY: Sparse attention retains only top-ρ% query–key pairs, masking the remaining pairs before computing attention.The selected sparse pattern is applied through a mask added to the attention logits.
- 3.2 OBSERVATIONS: DLM attention varies across heads, remains highly consistent within each head across denoising steps, and is particularly sensitive to sparsity during early steps.These properties make fixed autoregressive patterns unsuitable while motivating reuse of head-specific patterns later in denoising.
- 3.3 SPARSED: SparseD applies full attention in early diffusion steps, then pre-computes and reuses head-specific sparse patterns in subsequent steps.This design avoids recomputing sparse patterns at every denoising step while protecting the steps most important for generation.
- 3.3 SPARSED: SparseD selects important attention scores separately for prefill and generation keys, using the same selection ratio for both groups.Isolated selection addresses DLM heads whose attention toward generation tokens changes across the key dimension.
- 3.3 SPARSED: For hardware-friendly selection, SparseD average-pools attention scores into blocks and processes query blocks sequentially to reduce memory use.The pooled score matrix has dimensions l//block size × l//block size, and smaller query blocks avoid materializing the full attention matrix.
4 EXPERIMENTS
SparseD is evaluated on two DLMs across benchmark accuracy, latency, and component configurations. It largely preserves accuracy while improving long-context latency, with reuse and isolated selection supporting the method’s efficiency and quality.
- Accuracy: SparseD preserves benchmark accuracy, with only a 0.04% average drop on Dream-7B-Instruct and a 0.29% improvement on LLaDA-1.5.On RULER-8k, dKV-Cache and Fast-dLLM instead reduce Dream-7B-Instruct accuracy by approximately 16%, while reducing LLaDA-1.5 accuracy by 5.3% and 14.6%.
- Latency: At 64k context with 128 steps, SparseD achieves 1.23× and 1.25× speedups over FlashAttention on Dream-7B-Instruct and LLaDA-1.5, respectively.It matches FlashAttention at 4k and 8k and shows clear advantages beyond 16k; Figure 4 covers varying sequence lengths.
- Latency: At 64k context and 1,024 diffusion steps, SparseD reaches 1.50× and 1.48× speedups over FlashAttention on Dream-7B-Instruct and LLaDA-1.5, respectively.At 128 steps, the corresponding speedups are 1.23× and 1.25×; Figure 5 reports increasing acceleration as diffusion steps increase.
- Ablations: Removing skipping sparse attention severely lowers accuracy, while recomputing patterns each step increases latency compared with reusing them.The ablation evaluates accuracy on RULER-4k and latency on a 64k RULER sample.
- Ablations: Isolated selection improves accuracy with negligible latency overhead, confirming its contribution alongside skipping and sparse-pattern reuse.The ablation table excludes each component individually relative to SparseD.
- Hyper-parameters: Accuracy improves with skipping ratios up to a moderate range, reaching above 90.89% for LLaDA-1.5 at skip = 20% and plateauing at 90.00% for Dream-Instruct beyond 30%.The experiments use skip = 20% for both models.
- Hyper-parameters: For selection ratio ρ, Dream-7B-Instruct accuracy rises from around 40% at ρ = 5% to nearly 89.76% at ρ = 30%, while latency increases steadily.Accuracy saturates after ρ = 30% for Dream-7B-Instruct and after ρ = 20% for LLaDA-1.5; long-context experiments use ρ = 30%.
5 CONCLUSIONS
SparseD adapts sparse attention to DLM-specific head and denoising behavior by reusing head-specific patterns while retaining full attention early. It preserves tested accuracy and reaches up to 1.50× speedup over FlashAttention at 64k context with 1,024 steps.
- 5 CONCLUSIONS: SparseD pre-computes one sparse attention pattern per head, reuses patterns across diffusion steps, and applies full attention early to preserve generation quality.The design follows observations of head-specific patterns, cross-step consistency, and the importance of early diffusion steps.
- 5 CONCLUSIONS: SparseD achieves lossless performance on all tested benchmarks and up to 1.50× speedup over FlashAttention at 64k context with 1,024 diffusion steps.
A.1 ATTENTION PATTERNS
The attention patterns of LLaDA-8B-Base and Dream-7B-Instruct show head-specific structure and strong consistency across denoising steps, with rare corner cases.
- A.1 ATTENTION PATTERNS: Attention patterns in LLaDA-8B-Base exhibit distinct head-specific structure and strong consistency across denoising steps.Figures 7(a–c) show head-specific patterns, while Figure 7(d) shows cross-step consistency.
- A.1 ATTENTION PATTERNS: Rare corner cases reduce cross-step similarity, but later steps remain strongly similar in Figure 7 and similarity exceeds 60% in Figure 8.These cases occur in a minority of attention heads.
A.2 EXPERIMENTAL DETAILS
The evaluation covers language understanding, reasoning, code generation, and long-context benchmarks, using multiple DLMs and sparse or cache-based comparison methods.
- A.2 EXPERIMENTAL DETAILS: Experiments evaluate general language understanding, mathematical reasoning, code generation, and long-context performance across MMLU, GSM8K, HumanEval, and RULER.RULER provides the long-context evaluation component.
- A.2 EXPERIMENTAL DETAILS: RULER uses 13 subtasks with subtask-specific T and l parameters, detailed in Table 3.The settings vary to enable efficient evaluation across different subtasks.
- A.2 EXPERIMENTAL DETAILS: SparseD is compared with slide-window attention, StreamingLLM, dKV-Cache, and Fast-dLLM.FlexAttention accelerates the slide-window and StreamingLLM methods.
- A.2 EXPERIMENTAL DETAILS: The evaluation specifies model- and method-dependent configurations, including cache refresh intervals, thresholds, block sizes, and batch settings.Fast-dLLM uses threshold 0.9, with block size 8 for MMLU and 32 for other datasets.
A.3 EVALUATION DETAILS ON RULER DATASET
This section provides detailed accuracy results for each RULER dataset subtask, complementing the aggregate results in Tables 1 and 2.
- A.3 EVALUATION DETAILS ON RULER DATASET: Tables 4 and 5 provide detailed RULER subtask accuracy results corresponding to the aggregate results in Tables 1 and 2.Different evaluation configurations are used for the RULER subtasks.
A.4 LIMITATIONS
The paper’s main limitation is its algorithmic focus, leaving system-level optimization and integration with cache-based methods for future work.
- A.4 LIMITATIONS: SparseD’s algorithmic focus leaves system-level optimization of sparse-pattern computation and head-specific sparse attention as open directions.The limitation concerns both computing sparse patterns and accelerating their head-specific execution.
- A.4 LIMITATIONS: Combining lossless sparse attention with cache-based methods remains an unresolved implementation challenge.Future work must retain the advantages of both approaches simultaneously.