SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration

Authors: Jintao Zhang, Jia wei, Pengle Zhang, Jun Zhu, Jianfei Chen
Primary area: foundation or frontier models, including LLMs
Keywords: Attention, Quantization, quantized attention, inference acceleration

Abstract:
The transformer architecture predominates across various models. As the heart of the transformer, attention has a computational complexity of $O(N^2)$, compared to $O(N)$ for linear transformations. When handling large sequence lengths, attention becomes the primary time-consuming component. Although quantization has proven to be an effective method for accelerating model inference, existing quantization methods primarily focus on optimizing the linear layer.
In response, we first analyze the feasibility of quantization in attention detailedly. Following that, we propose SageAttention, a highly efficient and accurate quantization method for attention. The OPS (operations per second) of our approach outperforms FlashAttention2 and xformers by about 2.1x and 2.7x, respectively. SageAttention also achieves superior accuracy performance over FlashAttention3. Comprehensive experiments confirm that our approach incurs almost no end-to-end metrics loss across diverse models—including those for large language processing, image generation, and video generation. The code is available at https://github.com/thu-ml/SageAttention.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730592774231

thread_type:
official_review

<thread_object>
creator: reviewer_xy6z

<content>
summary:
This work explores post-training quantization of the attention mechanism. It claims that it's not well explored in the literature and the existing works focus either on quantized training or post-training quantization of non-attention layers (of course, among other attention acceleration methods, like linear attention, algorithmic accelerations, etc.). The work observed that the Keys matrix has similar values per-token, and proposed to smooth it by applying a per-token bias. It explores optimal quantization strategies and combines it with flash attention algorithmic techniques to obtain the best performance. SageAttention achieves very high reconstruction quality and does not seem to lose much performance on downstream metrics.

strengths:
- The performance boosts seem to be impressive compared to 
- The work carries a valuable observation that the keys matrix has huge correlation per-token. It can be valuable even outside of 
- I find the evaluation section to be very thorough, including many domains (text2image, text2video, LLMs). I especially appreciate text2video (where the sequences are the longest) and performance exploration on second-tier GPUs (e.g. 3090), which are more suitable for production use cases.
- I am not an expert in quantization, but I checked some prior works and didn't find similar approaches. This makes me consider the paper as novel.
- The paper provides the source code: I have not tried running it but it is valuable.

weaknesses:
- I think that writing should be improved. A table I would like to see the most is the end-to-end model speed improvement and metrics degradation after integrating the proposed attention mechanism. First, for many models accelerating the attention op itself does not necessarily lead to much improved overall speed when the sequence length is short and it's MLPs who carry the main burden. Second, it's not entirely clear when looking at various tables which variant is being used at the end (and why there are so many of them throughout the paper instead of being just in some restricted ablation section). Finally, there are quite many typos and grammatical errors, e.g.:
    - L084: "matrice" => "matrices"
    - L138 (and elsewhere): "IO" => "I/O"
    - L144: "it compute" => "it computes"
    - L142=144: missing line or equation
    - L151: "The σ()" => "The" should be ommited or it should be "The function σ()".
    - L184: "First, We" => "First, we"
    - L264-L266: ~3 typos in 3 lines
    - etc.
- There are no qualitatives provided for image/video generation.

questions:
- Is there a final attention variant among the proposed ones which performs the best uniformly across all tasks?
- L154: "The expression diag(l_i^j)^{-1}" would produce inf-s since l_i^j is initialized as 0. Or how is O_i^j initialized?
- L151: What is "online softmax" (i dont see it to be defined explicitly)? How does it differ from the standard softmax?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730608052953

thread_type:
official_review

<thread_object>
creator: reviewer_wwcu

<content>
summary:
This paper presents SageAttention, an innovative quantization method designed to optimize the computational efficiency of attention mechanisms in transformer models by employing an 8-bit integer (INT8) quantization. It employs the FlashAttention-wise quantization and matrix smoothing with FP16 for Matmul. With all these proposed components,  the method demonstrates improved computational performance and maintains accuracy compared to existing solutions like FlashAttention2 and xformers.

strengths:
+ The paper is well-written, with a logical structure and organization that facilitates understanding.

+ SageAttention shows competitive performance, outperforming FlashAttention2 and xformers by approximately 2.1x and 2.7x, respectively.

+ The method exhibits almost no end-to-end metrics loss across a variety of models, including large language models (LLMs), text-to-image (T2I), and text-to-video (T2V).

+ The discovery of channel-wise consistency, as illustrated in Figure 4, is particularly noteworthy and adds depth to the research.

weaknesses:
- The method relies heavily on FlashAttention, which may weaken its technical contribution and originality. What will the performance be if it does not employ the FlashAttention as the basis?

- The reported superiority over FlashAttention3 appears to be quite marginal, raising questions about the significance of the improvements.

- Another major weakness of this paper is that it does not compare SageAttention with other task-specific quantization methods, such as AWQ [1] for LLMs, Q-diffusion [2] for text-to-image, and ViDiT-Q [3] for text-to-video applications, which could provide a more comprehensive evaluation of its performance.

[1] AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

[2] Q-Diffusion: Quantizing Diffusion Models

[3] ViDiT-Q: Efficient and Accurate Quantization of Diffusion Transformers for Image and Video Generation

questions:
Could the authors elaborate on the "Llama" column of Table 1? Specifically, why does the number remain stable even with quantization? Understanding this aspect could provide valuable insights into the robustness of the proposed method.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730756229579

thread_type:
official_review

<thread_object>
creator: reviewer_tytt

<content>
summary:
This papers conducts in depth analysis of viability to quantize LLMs/Diffusion models into INT8 frameworks. It also proposes smoothing methods to alleviate the outlier pains in the QKV projection process, demonstrating viable tradeoffs. Comparisons to relevant work is strong.

strengths:
- Paper is well written.
- Experiments are thorough.
- Problem is challenging.

weaknesses:
- Full comparison to strong SOTA methods such as Flash attention 3, though slightly mentioned in the introduction and in Table 14, is not deeply explored. 
- Only targeted 4090/3000 series GPUs - it would be recommended to be tested on stronger GPUs at server level that is facing the strongest limitations.
- It would be great to test across VLMs too.

questions:
As above in weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730759657051

thread_type:
official_review

<thread_object>
creator: reviewer_utax

<content>
summary:
This paper introduces a new attention quantization method for speeding up transformer inference. To accelerate attention, the authors propose using INT8 quantization instead of FP8 for faster matrix multiplication on GPUs, along with a method for smoothing the K matrix to improve accuracy. Instead of quantizing the P and V matrices, they maintain them in FP16 and use a low-precision FP16 accumulator for faster multiplication without accuracy loss.  Finally, they offer different speed-accuracy trade-offs and a layer-wise selection method for optimal performance. Extensive experiments have been done on different transformer models for text, image and video generation tasks. The results show that the proposed SageAttention speeds up the FlashAttention2 and xformers by more than 2 times without losing any performance.

strengths:
1. The contributions of this work are well motivated. 

2. The proposed method seems to be novel although I am not an expert in this field. 

3. The experiments are quite extensive, covering two different GPUs (RTX4090 and RTX3090), representative models for language, image, and video generation, and a wide range of datasets. 

4. The results are quite impressive, showing more than two times speedup without performance degradation.

weaknesses:
1. Some design choices seem to be decided by the specific hardwares that are evaluated RTX4090 and 3090 (L271). Are those design choices also compatible with other GPUs like A100 and H100? 

2. Table 7 shows that different model/task has different speedup. How is the speedup related to the specific transformer architecture, model size, and complexity of the task?

questions:
I am not an expert in this field and do not have additional questions at this stage. I might have more questions at the discussion phase.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces SageAttention, a quantization technique for transformer attention, leveraging INT8 quantization and matrix smoothing to achieve substantial speedups on GPUs. Extensive evaluations show 2.1x–2.7x acceleration over FlashAttention2 and xformers without sacrificing accuracy across diverse models and tasks, including language, image, and video generation. SageAttention also outperforms FlashAttention3 in terms of accuracy for certain tasks and applicability on broader hardware.

Strengths:
* Performance Gains: SageAttention shows impressive speed-ups across different model architectures and GPUs, with consistent accuracy across a wide range of tasks.
* Broad Applicability: The method has been tested on a diverse set of models and GPUs, suggesting a robust and versatile approach.
* Thorough Evaluation: Comprehensive experiments cover both consumer and server-grade GPUs, providing insights into performance under varied conditions.

Weaknesses:
* Dependence on FlashAttention: While implementation builds on FlashAttention, reviewers accepted the justification of its orthogonality and its necessity for state-of-the-art performance.
* Comparison with Similar Methods: There's a noted lack of direct comparison with task-specific quantization methods like AWQ, Q-diffusion, and ViDiT-Q, though the authors have addressed this in rebuttals by adding comparisons.
* Presentation: Some reviewers pointed out areas in the manuscript where clarity and grammatical accuracy could be improved.

Overall, SageAttention is a well-executed work with broad applicability in accelerating transformer-based models. The thorough responses, additional experiments, and demonstrated utility across diverse tasks and hardware strengthen its case for acceptance.