Source-linked AI summary
SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference
Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, Jianfei Chen
TL;DR
Long sequences make attention costly, and existing sparse methods do not reliably generalize across diverse tasks while balancing accuracy and efficiency. SpargeAttn uses two online filtering stages with quantization to accelerate attention, achieving broad speedups without sacrificing end-to-end metrics.
Problem
Existing sparse attention methods have limited universality because task-specific patterns are difficult to generalize, while accurate prediction and low filtering overhead are difficult to achieve simultaneously.
Method
SpargeAttn combines selective token compression for universal sparse-mask prediction, an online softmax-aware filter, and 8-bit quantization.
Results
2.5x to 5x faster than existing dense and sparse attention models, SpargeAttn retains end-to-end model performance across language, image, and video generation tasks.
Takeaways & Limitations
SpargeAttn provides training-free sparse attention that can be adopted across language, image, and video generation models without compromising accuracy.
Takeaways & Limitations
Existing sparse attention methods remain limited in universality, with task-specific patterns that are difficult to generalize across tasks.
Abstract
from arXiv · showhide
An efficient attention implementation is essential for large models due to its quadratic time complexity. Fortunately, attention commonly exhibits sparsity, i.e., many values in the attention map are near zero, allowing for the omission of corresponding computations. Many studies have utilized the sparse pattern to accelerate attention. However, most existing works focus on optimizing attention within specific models by exploiting certain sparse patterns of the attention map. A universal sparse attention that guarantees both the speedup and end-to-end performance of diverse models remains elusive. In this paper, we propose SpargeAttn, a universal sparse and quantized attention for any model. Our method uses a two-stage online filter: in the first stage, we rapidly and accurately predict the attention map, enabling the skip of some matrix multiplications in attention. In the second stage, we design an online softmax-aware filter that incurs no extra overhead and further skips some matrix multiplications. Experiments show that our method significantly accelerates diverse models, including language, image, and video generation, without sacrificing end-to-end metrics. The code is available at https://github.com/thu-ml/SpargeAttn.
1. Introduction
Longer sequences make attention a significant part of inference latency, while sparse attention seeks to skip computations associated with near-zero attention values. SpargeAttn targets universal, training-free acceleration across tasks while retaining end-to-end metrics.
- 45K-128K sequence lengths in video generation and language models make attention a significant portion of inference latency.
- Attention sparsity arises because softmax often produces many values approaching zero, enabling sparse methods to omit corresponding computations.
- Existing sparse attention methods have limited universality because task-specific patterns vary significantly across language, image, and video generation tasks.
- Accurate and efficient sparse attention must predict sparse regions precisely while keeping prediction overhead minimal.
- SpargeAttn combines universal sparse-mask prediction, sparse online softmax filtering, and 8-bit quantization for training-free attention acceleration.
- 2.5x to 5x faster than existing dense and sparse attention models, SpargeAttn robustly retains end-to-end model performance across generative tasks.
2. Related Work
Related sparse-attention methods differ in how they construct sparse masks, while other attention accelerations optimize kernels, quantize, distribute workloads, or use linear-time designs. The cited approaches vary in task requirements, retraining needs, and compatibility.
- Pattern-required methods use fixed attention patterns such as sliding windows or attention sinks to construct sparse masks.
- Dynamic sparse attention computes sparse masks on-the-fly from the inputs.
- DitFastAttn is restricted to simple diffusion transformers and is incompatible with language models and MMDiT models such as Flux.
- Training-based methods modify attention computation and require retraining the entire model, making them more expensive than training-free methods.
- Kernel optimization, quantization, workload distribution, and linear-time attention are presented as approaches orthogonal to SpargeAttn.
3. SpargeAttn
SpargeAttn implements sparse FlashAttention with a two-stage online filter: compressed attention prediction first identifies blocks to skip, and softmax-aware filtering skips negligible updates during execution.
- Sparse FlashAttention: SpargeAttn tiles attention as FlashAttention and uses binary block masks to skip filtered QiK⊤j and ePijVj computations.The masks select which query-key blocks are computed during sparse FlashAttention.
- Selective Token Compression: Neighboring query and key tokens often have high similarity, enabling similar blocks to be compressed into representative tokens for sparse prediction.This pattern-free strategy avoids relying on fixed attention patterns across models.
- Selective Token Compression: Fix blocks are always computed, while selective blocks compute QiK⊤j and ePijVj only when compressed attention assigns them sufficiently high cumulative score.The method retains non-self-similar blocks because a single representative token may not capture their critical information.
- Selective Token Compression: The first-stage predictor averages self-similar blocks, computes compressed QK⊤, applies softmax, and selects positions whose cumulative attention reaches τ.Non-self-similar query or key blocks are treated as fix blocks rather than being omitted.
- Sparse Warp Online Softmax: The second-stage online softmax filter skips ePijVj when rowmax(Sij) is sufficiently below the running maximum, making the update negligible.The condition is evaluated across GPU-warp partitions using a threshold λ.
- Quantized Implementation: SpargeAttn integrates sparse computation with SageAttention quantization and uses CUDA implementation plus kernel fusion to reduce prediction overhead.The sparse and quantized operations are described as orthogonal and can be combined directly.
4. Experiment
Experiments evaluate SpargeAttn across language, image, and video generation models using end-to-end quality, speed, sparsity, overhead, and ablation analyses. Across these settings, SpargeAttn retains performance while improving attention and end-to-end efficiency.
- Setup: Experiments cover Llama3.1, CogvideoX, Mochi, Open-Sora-Plan, Flux, and Stable-Diffusion3.5 across text, image, and video generation.Evaluations include zero-shot language tasks, open-sora prompt sets, and end-to-end generation metrics.
- Quality and efficiency evaluation: SpargeAttn incurs almost no end-to-end metric loss across models and surpasses sparse baselines at various sparsity levels.Visible comparisons on CogvideoX, Flux, Stable-Diffusion3.5, Mochi, and Open-Sora-Plan show no performance loss and outperform baselines.
- Quality and efficiency evaluation: SpargeAttn achieves faster attention than Full-Attention and sparse baselines across varying sparsity levels.Kernel comparisons evaluate sequence length 22K and head dimension 128, including deployment with FlashAttention2.
- Ablation Study and key Insights: Prediction overhead is minimal compared with attention execution latency.Table 3 compares dynamic sparse block prediction overhead with attention latency.
- Quality and efficiency evaluation: 1.83x end-to-end speedup is achieved on Mochi with SpargeAttn.End-to-end latency is also evaluated on CogvideoX and Llama3.1.
- Ablation Study and key Insights: Hilbert Curve permutation improves block self-similarity and sparsity with only a marginal accuracy difference on Mochi.The self-similarity judge is reported to guarantee end-to-end accuracy in an ablation.
- Ablation Study and key Insights: Sparsity increases with sequence length on Llama3.1 and with sample timesteps in diffusion models.Diffusion-model sparsity also varies across layers and heads, motivating different hyperparameters for each.
5. Conclusion
SpargeAttn is presented as a universal sparse and quantized attention method that improves attention efficiency and accuracy across inputs. It uses two online filtering stages to skip matrix multiplications while preserving end-to-end metrics across language, image, and video generation models.
- SpargeAttn is a universal sparse and quantized attention method designed to execute attention efficiently and accurately for any input.
- Its first online filter rapidly predicts the attention map to skip some attention matrix multiplications.
- Its second online softmax-aware filter adds no overhead while skipping further matrix multiplications.
- Experiments report acceleration across language, image, and video generation models without sacrificing end-to-end metrics.
Impact Statement
The paper states that it aims to advance machine learning. It identifies potential societal consequences but does not highlight any specific ones.
- The work aims to advance the field of Machine Learning.
- The authors acknowledge many potential societal consequences of the work.
- The authors do not identify any societal consequence as requiring specific emphasis.
A.1. Detailed Explain and results of permutation ablation
The permutation ablation evaluates methods for rearranging visual tokens before attention. HilbertCurve consistently improves block self-similarity and sparsity with only marginal precision loss, while random permutation preserves precision but sacrifices sparsity.
- Permutation setup: Visual-token permutations transform tensors with shape T × H × W × d into flattened tensors with shape L × d, where L = T × H × W.The compared permutation methods are described in Table 8.
- Permutation results: HilbertCurve consistently achieves superior block self-similarity and sparsity on CogvideoX and Mochi, with only marginal precision loss.Table 9 reports Sim-q as query block self-similarity and Sim-k as key block self-similarity.
- Permutation results: Random permutation retains precision metrics but sacrifices sparsity, indicating dynamic adjustment to complex token sequences.
A.2. Ablation Study of Self-Similarity Judge
The self-similarity judge is examined through a targeted ablation that retains approximately 2% of tensor cases with substantial precision-error differences. It mitigates extreme precision loss with only a marginal sparsity reduction, especially for random permutations in CogvideoX.
- Ablation setup: The ablation retains approximately 2% of tensor cases where precision-error differences exceed 0.05 between judged and unjudged settings.Precision error and sparsity are used to evaluate the judge's influence.
- Findings: The self-similarity judge mitigates extreme precision loss while introducing only a marginal reduction in sparsity.
- Findings: Many cases with notable differences originate from the Random permutation category in CogvideoX, highlighting the judge's role in robustness to complex token sequences.
A.3. Additional Experiments
Additional experiments evaluate SpargeAttn on long-context Llama3.1, Open-Sora-Plan, and Mochi. The results include needle-in-a-haystack testing and visible generation comparisons with reported sparsity levels for several methods.
- Llama3.1 is evaluated on the Needle-in-a-Haystack task with sequence lengths of 16–28K.
- Visible examples are presented for Open-Sora-Plan.
- Mochi comparisons report sparsity of 0.47 for SpargeAttn, 0.3 for MInference, and 0.4 for FlexPrefill.
A.4. Sparsity analysis over diffusion model
The paper analyzes CogvideoX sparsity across layers, denoising timesteps, input samples, and attention heads. These dimensions exhibit diverse attention behavior and may inform diffusion-algorithm design.
- CogvideoX sparsity is analyzed across model layers, denoising timesteps, input samples, and attention heads.
- Layer-wise sparsity: Figure 14 presents layer-wise sparsity in CogvideoX.
- Timestep-wise sparsity: Figure 15 presents timestep-wise sparsity in CogvideoX.
- Sample-wise sparsity: Figure 16 presents sample-wise sparsity in CogvideoX.
- Head-wise sparsity: Figure 17 presents head-wise sparsity, illustrating diverse attention behavior across heads.