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

Q3: L151: What is "online softmax" (I dont see it to be defined explicitly)? How does it differ from the standard softmax?

Reply: We explain online softmax, and its use in FlashAttention as follows.

Online softmax, introduced by Maxim Milakov and Natalia Gimelshein from NVIDIA, is a kernel fusion method designed to accelerate the softmax computation on actual hardware. For a vector $a = \\{a_i\\},{i=1..n}$, the traditional softmax first computes $a_{max} = \max(a)$ and then $a' = a - a_{max}$ (to maintain numerical stability), followed by computing $\text{exp}(a') / \sum_{i=1}^n \exp(a'_i)$, which requires multiple CUDA kernel launches.

Online softmax fuses the computation of $a' = a - a_{max}$ and $\sum_{i=1}^n \text{exp}(a'_i)$ into a single CUDA kernel:

The calculation process initializes with $m_0 = -\infty$ and $d_0 = 0$, and iteratively updates with $m_i = \max(m_{i-1}, a_i)$ and $d_i = d_{i-1} \times \exp(m_{i-1} - m_i) + \exp(a_i - m_i)$. In the end, $\text{max}(a) = m_n$ and $\sum_{i=1}^n \exp(a'_i) = d_n$, and the result can be computed by $\text{exp}(a-m_n)/d_n$. This method considers each newly read $a_i$​​, updating the maximum and rescaling the partial sum based on the updated maximum to ensure equivalent computation.

Flash Attention integrates online softmax by combining it with $QK^\top$ and $PV$ multiplications, enabling a single kernel to perform the entire attention computation. Likewise, the $S_i^j$ in L151 corresponds to the $a_i$ in the about explanation. The rows of $S_i^j$ represent multiple vectors requiring softmax computation, which, in flash attention, corresponds to the number of tokens in a tile of $Q$. The columns of $S_i^j$ indicate the number of elements in a continuous subsequence of vector $a$ processed in parallel, corresponding to the number of tokens in a tile of $K$.

Dear Reviewer xY6z,

Thank you for recognizing the potential and effectiveness of our work and providing constructive comments. Below, we address each point raised.

W1-(1): 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.

Reply: Thank you for your suggestion and insightful conclusion. Table8 in our paper already displays the detailed degradation of end-to-end metrics after integrating SageAttention. We add a Table to illustrate the end-to-end speedup of various models as follows.

We highlight two examples in the following Table that demonstrate the end-to-end latency improvements achieved by SageAttention when the sequence length exceeds 10,000.

W1-(2): 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).

Reply: We apologize for any confusion caused. In all tables presented in our paper, 'SageAttention' denotes the combination of SageAttn-B and SageAttn-vB, as described in Section 4.5. Notably, SageAttn-B is already twice as fast as FlashAttention and delivers the highest accuracy. We have modified the clarification in Section 4.5. For practical applications, we recommend SageAttn-B due to its superior accuracy and speed.

W1-(3): Finally, there are quite many typos and grammatical errors.

Reply: Thank you for thoroughly reviewing our work and pointing out the typos. We have addressed and corrected the typos you identified, along with others we discovered. All corrections have been highlighted in purple.

W2: There are no qualitatives provided for image/video generation.

Reply: Thank you for your suggestion. We add some image and video comparison examples in the Appendix in our paper. To better show the video, we also add these examples at https://anonymous.4open.science/r/image_video_examples-3E44/README.md.

Q1: Is there a final attention variant among the proposed ones which performs the best uniformly across all tasks?

Reply: Yes, SageAttn-B consistently performs the best across all tasks. Although SageAttention combines both SageAttn-B and SageAttn-vB, it is important to note that SageAttn-vB is only 4% faster than SageAttn-B, yet it exhibits inferior accuracy. Therefore, SageAttn-B remains the preferred choice due to its superior performance.

Q2: 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?

Reply: Actually, the $diag(l_i^j)^{-1}$ in L154 refers to the state after completing all iterations from Lines L144-149, i.e., when $j = b_{kv}$. Initially, only $diag(l_i^0)=0$, while all subsequent $diag(l_i^j)\neq 0$. We also clarify this point by revising this part in our paper.

Dear Reviewer xY6z,

Thank reviewer xY6z for the valuable comments and the strong acceptance of our work. We hope our reply has addressed the concerns and brought new insights to the reviewer.

In addition, please let us know if you have any other questions or concerns. Thank you very much!

I am thankful to the authors for such a detailed response. I have one more remaining question: why CogVideoX is referred as "Full Precision" in Table 8? As far as their README tells, fp16 was used for 2B and bf16 for 5B. Does CogVideoX use flash attention or vanilla torch attention?

Also, for the specified metrics in the first table of the above response, what metrics exactly are being used to determine metrics degradation (In Table 8, there are many metrics)? Is it average across all the metrics?

Q4: I have one more remaining question: why CogVideoX is referred as "Full Precision" in Table 8? As far as their README tells, fp16 was used for 2B and bf16 for 5B. Does CogVideoX use flash attention or vanilla torch attention?

Reply: Thank you so much for your question, and we apologize for any confusion. "Full Precision" in Table 8 refers to the original precision of the target model's attention before quantization. For the CogVideoX-2B model we used, this corresponds to FP16, as pointed out. We will modify Table 8 to specify this precision from "Full Precision" to "Origin (FP16)".

CogVideoX-2B originally uses FlashAttention, which is noted in Table 7. We will revise this Table to specify the precision further.

Q5: Also, for the specified metrics in the first table of the above response, what metrics exactly are being used to determine metrics degradation (In Table 8, there are many metrics)? Is it average across all the metrics?

Reply: Table 8 includes many metrics because a single end-to-end quality metric is often not sufficient for language, image, and video generation model evaluation. For example, text-to-image generation requires assessing both image quality and the alignment between text and image. Each metric reflects a potential metrics degradation, and these are commonly used measures in the field, briefly explained in Section 5.1. The "overall degradation" mentioned in the previous reply is indeed the average of the metrics degradation for each metric listed in Table 8.

Please let us know if you have any other questions or concerns. Thank you very much!

W3: 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.

Reply: Thank you for your advice. First, we argue that SageAttention is orthogonal to AWQ, Q-diffusion, and ViDiT-Q because those works are mainly used to quantize the linear layers. Second, AWQ is only used to compress the parameters of LLMs with no acceleration effect in computation. Q-diffusion has not reported its acceleration effect in their paper and provided codes with acceleration effect in its official repository. ViDiT-Q has not provided the codes with acceleration effect in its official repository. Nonetheless, we compare SageAttention with those works as follows.

We compare the perplexity of Llama2-7B on WikiText and the speedup in the prefilling stage as follows.

We compare SageAttention with Q-diffusion (W8A8) on Unidiffuser using the same experimental setting in our paper as follows.

We compare SageAttention with VIDIT-Q on CogvideoX using the same experimental setting in our paper as follows. Since the official repository does not provide acceleration code, we estimate a theoretical theoretical maximum: the Linear layer accounts for 24% of Cogvideo's latency, and W8A8 offers at most 4x speedup for Linear layer, resulting in a theoretical maximum end-to-end speedup of 100/ (100 - 24*3/4)% = 22%.

We emphasize that SageAttention is orthogonal to works that focus on quantizing Linear layers such as AWQ, Q-diffusion, and ViDiT-Q.

Q1: 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.

Reply: Thank you for your very insightful question. The reason is that the distribution of $Q, K$, and $V$ in the attention of Llama2-7B is relatively uniform. As a result, quantizing $Q, K$, and $V$ to INT8 or FP8 does not significantly impact the accuracy of attention. This insight inspires the idea that better control over outlier activations in models could lead to more precise quantization results. We add more data distribution figures for $Q, K$, and $V$ of Layer0, Layer10, and Layer20 in Llama at [https://anonymous.4open.science/r/sageattention_example-FB01/resource/rebuttal_llama_acc.png]. We also report their attention accuracy when quantizing $Q, K$ to INT8 or FP8 as follows.

Dear Reviewer wwcu,

Thank you for acknowledging our presentation, performance, and novelty. We address individual comments below:

W1: The method relies heavily on FlashAttention, which may weaken its technical contribution and originality.

Reply: We argue that the contributions of SageAttention are significant and distinct from FlashAttention for the following reasons:

1. The implementation of SageAttention is based on FlashAttention, but this is reasonable for the following reasons. Firstly, FlashAttention is a SOTA method designed to avoid reading/writing the Softmax matrix in attention. To conduct efficient standard Attention, using such a method is almost inevitable. Just like efficient software development relies on optimized algorithm libraries. Secondly, apart from the implementation, our method can also be applied based on other attention implementations, such as Torch. Lastly, some sparse attention works, such as [1] and [2], all based on FlashAttention, were published in NeurIPS 2024 and 2023.

   [1] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention, NeurIPS 2024.

   [2] Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention, NeurIPS 2023.

2. Our contribution is orthogonal with FlashAttention. We systematically explore the challenges of quantizing $Q, K, P, V$ on the speed and accuracy of attention across various models, and implement efficient quantized attention applicable across multiple tasks and models. Given that FlashAttention is currently the most widely used in the industry, we chose to implement our approach based on FlashAttention. Our main contributions include proposing novel quantization methods (such as smoothing K), choosing suitable numerical formats (INT8 for $QK$, and FP16 with FP16 accumulator for $PV$), and choosing suitable quantization granularity to preserve accuracy. On the implementation side, our implementation is compatible with FlashAttention and various GPUs, achieving a high acceleration ratio while ensuring accuracy. Some contributions are highly nontrivial and original.

W1-(2): What will the performance be if it does not employ the FlashAttention as the basis?

FlashAttention is the state-of-the-art and most commonly used standard attention; another commonly used attention is Torch attention. We report the implementation speeds based on Torch and FlashAttention as follows.

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

Reply: We argue that our superiority is not marginal for the following reasons:

1. As shown in Table1, SageAttn-B outperforms FlashAttention3 by 1.1x and 2.5x on end-to-end metrics on image and video generation tasks.
2. Although the improvement on Llama2 and TIMM is marginal, the reason is that the two tasks are easy to attention quantization.
3. The limitation of FlashAttention3 is obvious; that is, FlashAttention3 can only be used on and is optimized exclusively for expensive Hopper architecture GPUs, i.e., H100 and H800 GPUs. It cannot be used on commonly used GPUs such as the RTX 4090 and A100.
4. Figure3 only provides one visible comparison example with FlashAttention3. To supplement this, we provide more visible examples at [https://anonymous.4open.science/r/sageattention_example-FB01/README.md], showing that FlashAttention3 fails in image and video generation visibly.

Thank you so much for raising the score! We greatly appreciate your thoughtful suggestions, which significantly improve our work. We will revise the manuscript to include all the new experiments as recommended.

Thanks for the detailed response of authors. Most of my concerns are addressed and I would like to raise the score (I will give score 7 if there is such an option). I highly recommend the authors to revise the manuscript to include all the new experiments.

Dear Reviewer UTAx,

Thank you for recognizing the potential and effectiveness of our work and providing constructive comments. Below, we address each point raised.

W1: 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?

Reply: Yes. Our method is also effective on A100, A6000, A800, and H100 GPUs. We compare the TOPS of different attention methods on those GPUs in the following Tables.

A100:

A6000:

A800:

H100:

While the superiority of SageAttention on these GPUs is not as high as it is on consumer-level GPUs like RTX 4090 and RTX 3090, we argue that the performance degradation occurs as follows:

1. Tensor Core: On RTX 4090 and 3090, SageAttention utilizes the INT8 Tensor Core, whose throughput is 4x faster than the FP16 Tensor Core. However, on A100, A6000, A800, and H100, this speedup is reduced to 2x.
2. Accumulator: The speedup of the Tensor Core with FP16 accumulator is more marginal than on RTX 4090 and 3090.
3. Instance optimization: Some GPU architectures have unique speed optimization methods, which we need to take advantage of in future work. For example, the H100 allows for asynchronous computation between Tensor Cores and CUDA Cores.

W2: 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?

Reply: There is no specific relationship since the speedup refers solely to the acceleration of the attention part. This speedup is independent of the transformer architecture, model size, or task complexity, because it is only related to the attention operation, which is influenced by the original implementation, e.g., xformers or FlashAttention, and the shape of Q, K, and V. Different shapes have varying speeds primarily due to slight differences in hardware utilization caused by different sequence lengths.

End-to-end acceleration mainly depends on the proportion of time spent on attention. For example, attention accounts for 52% of the time in the CogVideo, and a 2x speedup in attention would result in an overall speedup of 1.35x (calculated as 100/(100-26)). If the sequence length is short, attention will occupy a smaller fraction of the time, limiting the potential for end-to-end acceleration.

A800:

H100:

While the superiority of SageAttention on these GPUs is not as high as it is on consumer-level GPUs like RTX 4090 and RTX 3090, we argue that the performance degradation occurs as follows:

1. Tensor Core: On RTX 4090 and 3090, SageAttention utilizes the INT8 Tensor Core, whose throughput is 4x faster than the FP16 Tensor Core. However, on A100, A6000, A800, and H100, this speedup is reduced to 2x.
2. Accumulator: The speedup of the Tensor Core with FP16 accumulator is more marginal than on RTX 4090 and 3090.
3. Instance optimization: Some GPU architectures have unique speed optimization methods, which we need to take advantage of in future work. For example, the H100 allows for asynchronous computation between Tensor Cores and CUDA Cores.

W3: It would be great to test across VLMs too.

Reply: Thank you for your suggestion. We add an experiment on a VLM (Llava1.6-Mistral-7B) as shown in the Table below, where the TextVQA, POPE, and VQAv2 are three commonly used datasets, and the metric is accuracy. We also add this experiment to our paper.

Dear Reviewer tyTT,

Thank you for recognizing the potential and effectiveness of our work and providing constructive comments. Below, we address each point raised.

W1: Full comparison to strong SOTA methods such as Flashattention3, though slightly mentioned in the introduction and in Table 14, is not deeply explored.

Reply: Firstly, from an accuracy perspective, SageAttention outperforms FlashAttention3, especially in video and image generation.

1. We have already compared the end-to-end metrics of SageAttention and FlashAttention3 in Table1 in our paper, where the 4-th row is SageAttn-T, and the 6th row is SageAttn-B, across language generation, image generation, video generation, and image classification models. The results show that SageAttention outperforms FlashAttention3.

2. Figure3 provides only one visible comparison example with FlashAttention3. To supplement this, we offer more visible examples at https://anonymous.4open.science/r/sageattention_example-FB01/README.md, showing that FlashAttention3 fails in image and video generation visibly.

3. We supplement the comparison of kernel accuracy for SageAttention versus FlashAttention3 when (Q, K, V) follows a standard normal distribution $\mathcal{N}(0, 1)$, with 1% replaced by values from a normal distribution $\mathcal{N}(1, 10)$. The results are in the Table below.

   Attention	Cos Sim$\uparrow$	Relative L1$\downarrow$	RMSE$\downarrow$
   FlashAttention3	98.93%	0.1393	0.0636
   SageAttention-T	99.92%	0.0386	0.0026
   SageAttention-B	99.89%	0.0467	0.0210

Secondly, in terms of speed and applicability, the main focus of our study is to accelerate attention on commonly used hardware, such as RTX4090, 3090, and A100. However, FlashAttention3 can only run on expensive Hopper architecture GPUs, i.e., H100 and H800 GPUs.

4. We compare the kernel speed of SageAttention and FlashAttention3 in the Table below.

Although FlashAttention3 is faster than SageAttention on H100, this is because its main contribution lies in utilizing the programming features of the Hopper architecture. In theory, our method can also be further optimized for the Hopper architecture.

W2: 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.

Reply: Yes. Our method is also effective on A100, A6000, A800, and H100 GPUs. We compare the TOPS of different attention methods on those GPUs in the following Tables.

A100:

A6000:

Dear Reviewer tyTT,

We want to thank reviewer tyTT for the valuable comments. If our reply has addressed the concerns and brought new insights to the reviewer, we will highly appreciate it if the reviewer considers raising the score.

In addition, please let us know if you have any other questions or concerns. Thank you very much!

Dear Reviewer tyTT:

We really appreciate your insightful and constructive comments, which helped us improve this paper. We have made our maximum effort to address all your concerns. As the discussion phase deadline is approaching, we would like to ask if our responses have addressed your concerns. We’d be happy to provide further clarification if there are any remaining issues.

If you feel that your concerns have been resolved, we would greatly appreciate it if you consider raising the score.

Best wishes!
The Authors