SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization

Dear Reviewer aDo7,
Thank you for your valuable suggestions and questions.

Weakness1

Reply: We appreciate the valuable question. We argue that:

1. We first discover the critical role of accumulator precision in PV matrix multiplication for Attention; this is an essential insight for the attention operator.

2. The technique and the methods mentioned in the paper were developed concurrently (within a two-month interval). Additionally, we were the first to discover and report this hardware issue, and we can provide evidence after the rebuttal.

3. Beyond addressing hardware limitations, The technique can enhance accuracy for efficiency techniques using reduced accumulator precision in matrix multiplication.

Weakness2

Reply: Thank you for your question. We argue that:

1. Smooth Q and Smooth K are distinct. Specifically, smooth Q is a per-block level smoothing and introduces an additional matrix-vector multiplication to compute ΔS added in the GPU kernel. Moreover, despite its simplicity, this approach demonstrates exceptional improvement. We argue that contribution should not be dismissed based on methodological simplicity.

2. The improvement from per-thread quantization over SageAttention’s per-block quantization is significant, not marginal. For example, Table 6 clearly shows superior accuracy with per-thread quantization. Also, using per-block quantization for Sage2-4b will get completely blurred videos of Cogvideo.

3. Per-thread quantization is a highly innovative approach, and is non-trivial, requiring sophisticated algorithm-hardware co-design.

Weakness3

Reply: Thank you for your valuable suggestion. We add the accuracy ablation results and end-to-end perplexity of Llama-3.1 alongside the TOPS and will revise Table 18.

Weakness4

Reply: We appreciate your suggestion about comparing other quantization methods. We want to highlight that our paper already includes comparisons with per-tensor, per-block, and per-token quantization methods in Table 6 and 15, with discussions in lines 250–256. The results demonstrate that per-thread quantization achieves accuracy comparable to per-token quantization and outperforms per-tensor and per-block quantization.

Furthermore, we provide a speed comparison of different quantization methods, showing that per-thread quantization introduces almost no computational overhead compared to per-tensor and per-block methods, while per-token quantization is notably slower:

Comment1

Reply: We appreciate the reviewer’s suggestion. We evaluate the end-to-end inference latency of Mochi on L20 below:

SageAttention2 achieves a similar inference speed to HadamardAttn and SmoothAttn while delivering significantly better accuracy performance.

Comment2

Reply: Thank you for your suggestion. First, V's FP8 conversion is already shown in Figure 2. Second, the P matrix only exists internally within the GPU kernel and is difficult to visually represent in the overview figure. Therefore, we mainly describe its FP8 conversion process in detail in Algorithm 1 and Section 3.3. We will revise Figure 2 in our paper to explicitly indicate that P is converted to FP8.

Question1

Reply: We appreciate the reviewer’s question. Our speed evaluation already covers a range of sequence lengths from 1,024 to 32,768, demonstrating speedup across all sequences (e.g., 1,024). To further address your concern, we provide additional results for batch size = 1 and sequence length = 1024 on RTX4090, as shown below:

The results show that SageAttention2 consistently delivers higher throughput with short sequences and small batch sizes, maintaining similar speedup as observed in long-sequence and large-batch scenarios.

Question2

Reply: Thank you for your valuable suggestion. First, per-token quantization can not be applied to V because the quantization must be conducted along the outer axis of $PV$. We compare the accuracy of per-channel, per-block, and per-tensor quantization methods, showing that per-channel quantization achieves the best accuracy:

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

Dear Reviewer r8J3,
Thank you for your valuable question. Below, we address the question raised.

Question1. Do the authors think that it is possible for PV computation to be quantized to sub-8-precision as well?

Reply: Thank you for the insightful suggestion. The answer is yes - we have tried quantizing P and V to FP4 using Blackwell GPUs' micro-scaling quantization and obtained preliminary accuracy results across CogVideo layers.

Specifically:

• We keep Q and K in INT8 (since Blackwell GPUs don't support INT4)
• We quantize P and V to NVFP4
• Using our smooth Q and smooth K techniques

Based on these preliminary results, we believe FP4 quantization for P and V is possible and practical.

Dear Reviewer t1YD,
Thank you for your valuable suggestions and questions. Below, we address each point raised.

Weakness1. The effectiveness of proposed Q/K smoothing is evaluated empirically. The evaluation is sufficient but it might be better to analysis the theoretical benefits of Q/K smoothing. For example, the authors can analysis and compare the quantization error bounds with/without smoothing for specific input distributions. This provides the paper with stronger theoretical support.

Reply: Thank you for the valuable suggestion. We analyze the quantization error as follows:

Proof: Let $X \in \mathbb{R}^{N \times d}$ be $N$ activation tokens with dimension $d$.

Following [1], we suppose that an activation token follows an Gaussian distribution $\mathcal{N}(\boldsymbol{\mu}, \Sigma^2)$, where $\boldsymbol{\mu} = (\mu_1, \mu_2, \ldots, \mu_d)$ and $\Sigma^2$ is a diagonal matrix with $\Sigma^2 = \mathrm{diag}(\sigma_1^2, \sigma_2^2, \ldots, \sigma_d^2)$.

Further, we suppose that different token $X_i$ is i.i.d. sampled from the same distribution.

Suppose the absolute maximum value in a quantization group is $M$, and the bit width is $b$, then there are $2^b$ quantization levels. Under round to nearest strategy, the expected quantization error is $\frac{1}{2} \cdot \frac{2M}{2^b}$ which is proportional to the maximum absolute value in the quantization group. So smaller absolute maximum value leads to smaller quantization error.

After smoothing, we have:

$Y_{ij}$ also follows a Gaussian distribution. The mean and variance of $Y_{ij}$ can be calculated as follows:

So $Y_{ij}$ have smaller mean and variance compared to $X_{ij}$. Following the property of Gaussian distribution, we have:

So the distribution of $X_{ij}$ is more concentrated towards 0 after smoothing. Then we know that

So this makes the distribution of absolute max value in a token more concentrated towards 0, leading to smaller quantization error.

Comment1. The experimental settings of Figure 5 are not clear. What is number of heads? Which dataset is used?

Reply: We apologize for the lack of clarity in the experimental settings of Figure 5. We used a head size of 32 and a batch size of 4. Since we are benchmarking kernel speed, we follow standard practice as FlashAttention1/2/3, i.e., using Gaussian input (with mean 0, standard variance 1) for floating-point inputs. For integer inputs, we use uniform random values within the representation range: [-128, 127] for INT8 and [-8, 7] for INT4. We will clarify this in the final manuscript.

Question1. The authors mentioned that the accumulator for the mma(f32f8f8f32) instruction is actually FP22. I am not familiar with mma instruction. Is the precision mismatch some kind of software "bug"? Or perhaps there are deeper design reasons behind it?

Reply: Thank you for your question. This is a hardware-level issue, not a software bug. The exact reason for this precision mismatch is not publicly documented, but we hypothesize that it is due to chip area constraints. Given the limited space, the designers may have had to make trade-offs, and reducing the accumulator precision for the FP8 tensor appears to have been one such compromise.

[1] QLoRA: Efficient Finetuning of Quantized LLMs

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

Dear Reviewer KKjK,
Thank you for your valuable suggestions and questions. Below, we address each point raised.

W1. Author highlighted several times about the comparison with FlashAttention2/3 throughout the paper, however, FlashAttn were missing in many cases in the main experimental results, for example, only the middle part of Table 2 (for video models) shows FlashAttn3-fp8, but the other two categories only have SageAttn family. Since it was mentioned in Abstract, Introduction, and Fig. 1, readers would be expecting a quantitative benchmark with both FlashAttn2 and 3. Especially when SageAttn2 provides two modes, i.e. one is faster (INT4) while the other is more accurate (INT8), a dedicated table for a fair comparison to FlashAttn family would be beneficial.

Reply: We apologize for any confusion regarding the comparison with FlashAttention2/3. The "full precision" row in our tables represents FlashAttention2/3’s unquantized attention performance, making Table 2 a direct and fair comparison between SageAttn2 and FlashAttention2/3. Note that since FlashAttention3 can only run on Hopper GPUs (H100/H20), so the speed comparisons with it are naturally limited to these GPUs. We will revise the manuscript to explicitly clarify this points.

W2. In Section 3.2, author stated that per-token method would result in "significant overhead" but didn't specify/quantify how serious this problem is. On the other hand, in Table 18, it only shows per-thread method added no overhead. Maybe author can elaborate a bit in Sec 3.2 and give readers a better idea how much improvement were made by this technique.

Reply: Thank you for your valuable feedback. We appreciate your point about the lack of clarification on the overhead introduced by the per-token method in Section 3.2.

We measure the TOPS of per-token INT4 quantization and the result is as follows, with the settings aligned with Table 18:

Per-thread introduces about 0.4% overhead, while per-token quantization results in approximately 6% overhead, which is 15 times the overhead of per-thread quantization.

W3. In Appendix Table 9, 14, and Fig 18, it wasn't clear which version of SageAttention2 (4b or 8b) were used in the experiments.

Reply: We apologize for the ambiguity. On Hopper GPUs (H100/H20), INT4 tensor core is not available. So in Figure 9, the columns of H100 and H20 report the result of SageAttn2-8b, while all other entries of SageAttn2 use the 4b version. For Table 14 and Fig. 18, the experiment was conducted on H100 (as stated in line 1041), so we use SageAttn2-8b and compared it with FlashAttention3-FP8, which has the same bit width, for a fair evaluation. We will clarify this in the final manuscript.

C1. typo on Line 115/116 right column

Reply: Thank you for pointing out the typo and we will revise our paper.

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

Dear Reviewer 1zQu,
Thank you for your valuable suggestions and questions. Below, we address each point raised.

Comment1

Reply: Thank you for your valuable suggestion. We compared the accuracy of sageattention, sage2-8b, sage2-4b, and sage2-4b without smooth Q across CogVideo layers (see table below). The limited INT4 representation range causes precision loss, but this is unrelated to Smooth Q. Before smoothing, the distribution was less uniform. While smoothing makes it closer to normal (though not perfectly uniform), it significantly reduces quantization error.

Comment2

Reply: Thank you for suggesting benchmarks on speech-related tasks. In response, we evaluated Qwen-2-Audio 7B, a speech-to-text model, on the ASR task using the Librispeech test split and measured its performance with the WER metric. The results are presented below:

We can see that SageAttention2 consistently outperforms the baselines, highlighting its effectiveness in audio-related models and benchmarks.

Comment3

Reply: We agree and recognize the significance of end-to-end latency for real-world performance evaluation, which is why we report it in our results. However, we emphasize that the end-to-end speedup depends entirely on:

1. The attention speedup.
2. The proportion of total latency attributed to attention.

For example:

• If a model spends 50s on attention and 50s on other operations, a 2× attention speedup reduces latency by 25s (total: 75s).
• A 2.5× attention speedup yields a 30s reduction (total: 70s).

Moreover, quantifying the attention kernel’s FLOPs is essential, as it directly measures computational efficiency and our method’s advancement. This approach aligns with the evaluation standards of FlashAttention 1/2/3.

Comment4

Reply: Thank you for your valuable suggestion. We compare the accuracy of Sage2-8b with and without 2-level FP32 buffer accumulation across all layers of CogvideoX:

Question1

Reply: Thank you for your question. Yes, the reason for not including SageAttention2-4B in H100 and H20 benchmarks is that Hopper GPUs do not have INT4 Tensor Core.

Question2

Reply: Thank you for your question. H20 and H100 are both Hopper architecture. Yes, it is technically feasible to pack INT4 into INT8, dequantize them to INT8 in the kernel, and use INT8 Tensor Core ops. However, this approach has significant drawbacks:

1. Speed Impact:

   • It requires additional CUDA Core operations for dequantization and has the same Tensor Core count as INT8. Therefore it won't bring speed gains.

2. Accuracy Impact:

   • INT4 quantization introduces greater accuracy loss than INT8.

Question3

Reply: Thank you for your valuable suggestion. We believe that learning the scaling factors would not help reduce the gap. Learned scaling factors are static and input-independent. However, the inputs exhibit significant fluctuations [1, 2], making static scaling suboptimal. For example, if all inputs are smaller than the predetermined scales, the representation range of INT4 is not fully utilized. Moreover, calibrating will compromise true plug-and-play compatibility.

A more promising approach may be learning clipping factors [3]. It scales the absolute maximum by a learned ratio, which helps suppress outliers while preserving the representation range.

[1] AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
[2] SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models
[3] OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models

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