SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Authors: Jintao Zhang
Haofeng Huang
Pengle Zhang
Jia wei
Jun Zhu
Jianfei Chen
Venue: ICML 2025 poster (2025)
Keywords: Efficient attention
quantization
inference acceleration
Dit
video generation
AI Infrastructure
LLM Infrastructure
Primary area: deep_learning->attention_mechanisms
TL;DR: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization.

Abstract: Although quantization for linear layers has been widely used, its application to accelerate the attention process remains limited. To further enhance the efficiency of attention computation compared to SageAttention while maintaining precision, we propose SageAttention2, which utilizes significantly faster 4-bit matrix multiplication (Matmul) alongside additional precision-enhancing techniques. First, we propose to quantize matrixes $(Q, K)$ to INT4 in a hardware-friendly thread-level granularity and quantize matrixes $(\widetilde P, V)$ to FP8. Second, we propose a method to smooth $Q$, enhancing the accuracy of INT4 $QK^\top$. Third, we propose a two-level accumulation strategy for $\widetilde PV$ to enhance the accuracy of FP8 $\widetilde PV$. The operations per second (OPS) of SageAttention2 surpass FlashAttention2 and xformers by about **3x** and **4.5x**, respectively. Moreover, SageAttention2 matches the speed of FlashAttention3(fp8) on the Hopper GPUs, while delivering much higher accuracy. Comprehensive experiments confirm that our approach incurs negligible end-to-end metrics loss across diverse models, including those for language, image, and video generation.

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Jintao Zhang" *1
Haofeng Huang 12
Pengle Zhang'
Jia Wei
Jun Zhu
Jianfei Chen
Abstract
increase in real-world applications (Jiang et al. 2024 Zhang
et al. 20256). Several strategies have been developed to
Although quantization for linear layers has been widely used, its application to accelerate the atten- tion process remains limited. To further enhance the efficiency of attention computation compared to SageAttention while maintaining precision, we propose SageAttent ion2, which utilizes sig- nificantly faster 4-bit matrix multiplication (Mat- mul) alongside additional precision-enhancingg techniques. First, we propose to quantize matri- ces (Q, K) to INT4 in a hardware-friendly thread- level granularity and quantize matrices (P, V) to FP8. Second, we propose a method to smooth Q. enhancing the accuracy of INT4 QK Third, we propose a two-level accumulation strategy for PV to enhance the accuracy of FP8 PV. The oper- ations per second (OPS) SageAttent ion2 surpass FlashAttention2 and xformers by about 3x and 4.5x. Moreover, SageAitent ion2 matches the speed of FlashAAteentiion3((p8)) on the Hopper GPUs, while delivering much higher ac- curacy. Comprehensive experiments confirm that our approach incurs negligible end-to-end met- rics loss across diverse models, including those for language, image, and video generation. The code is available at https: ptittp::/githutt github..hhug.eo.g com/ thu m tention
mitigate the computational demands of attention such as
(1) Linear attention methods (Wang et al. 2020 Choro
manski al. 2021] Yu et al. 2022 Katharopoulos et al.
2020) that reduce complexity to ((N) and (2) Sparse ar-
lention methods (Liiu et al. 202 Chu et al. 2021 Li et al.
2022
Xiao et al., 2024b Chen et al. 2024; Jiang et al.
2024
Venkataramanan al.
2024
Gao et al. 2024] Fu
ct
2025 Yang et al. Zhang al. 2025t h that selectively process parts of the context these methods are only suitable for a lim- ited range of models and tasks. The widely used attention methods optimize attention by exploiting hardware capaci- ties, such as FlashAttention V1, V2, V3 (Dao et al. 2022 Dao 2024 Shah et al., 2024), xformers (Lefaudeux et al. 2022). ) and SageAttention (Zhang et al. 2025c g gl) These works do not omit computations across the entire sequence and achieve impressive speed and accuracy performance across various tasks.
2024| Zhang et al. 2025e Xi et al.
2025a
Motivation. For the two matrix multiplication (Matmul) operations in attention: QK T and PV, SageAttention ac- celerates them by quantizing the QKT to INT8 and uses FP16 Matmul with FP16 accumulators for PV. More- over, to maintain attention accuracy, SageAttention pro- poses smoothing K by eliminating its channel-wise out- liers. SageAttention achieves 2 speedup compared with FlashAttention2 and is the first quantized attention that in- curs a negligible end-to-end metrics loss across language, image, and video generation models. However, SageAt- tention has two weaknesses. (W1) INT8 Matmul achieves only half the speed of INT4. (W2) FP16 Matmul with FP16 accumulators provides speedup only on RTX 4090 and RTX 3090 GPUs. To leverage the faster INT4 tensor cores for QKT T and using a method that can accelerate PV on a broader range of GPUs, we propose to quantize Q, K to INT4 and P, to FP8.
1. Introduction
Due to the quadratic time complexity of attention, its effi- cient implementation becomes crucial as sequence lengths
'Equal contribution 'Dept. of Comp. Sci. and Tech., Insti- tute for AI, BNRist Center, THBI Lab, tsinghua-Bosch Joint ML Center, Tsinghua University "Institute for Interdisciplinary Infor- mation Sciences, Tsinghua University. Correspondence to: Jianfei Chen
Challenges. Quantizing Q, K to INT4 and P, V to FP8 presents significant challenges. For example, when only per- tensor quantizing Q, K to INT4, the text-to-video model CogvideoX will generate a completely blurry video, and
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
RTX4090
L20
Overal Score: 0.834
Overal Score: 0.854
11.3
11.3
FlashAttention2
F
22.0
22.0
33.0
33.0
SageAttention2
SageAttention2
44.0
44.0
3x
481
39s
1.7x
67.0
56.0
56.0
67.9
4.4
163
23s
78.0
FlashAttention2
SageAttention2
100
100
Kernel Speed (TOPS)
End-to-end Speedup
FlashAttention2:
End-to-End Latency:
1040s on RTX4090
SageAttention2:
End-to-End Latency:
577s RTX4090
1.8x Speedup
Figure 1. The upper left figure shows the kernel speedup on RTX4090 GPU. The upper right figure shows the end-to-end inference speedup of generating the first token and performance metrics for the Needle-in-a- Beeyle-in-a-aaystack task (Kamradi 2023) with sequence length of 100K on L1ama3. 1 on GPU. The figure below shows two videos generated by CogvideoX (1.S-5B) using FlashAttentio22 SageAttent on2 on RTX4090, showing that SageAttention2 accelerates generation by 1.8x with vidco quality loss.
Llama3 only achieves random--uessing-llved accuracy of 25% on MMLU. After an in-depth investigation, we identified three primary challenges: (C1) The numerical range of INT4 is quite limited compared FP16 INT8, leading significant quantization errors when Q and K have some abnormal values. (C2) We discover that the FP32 accumulator designed for FP8 matrix multiplication in the tensor core (mma. f32. f8.. f8. f32) is actually FP22, specifically with sign bit, 8 exponent bits, and 13 mantissa bits. This will lead to an accuracy loss of PV.
sion of Q. referred to as Q m. Then, we add Q Q after the QK Matmul to ensure the correctness of the attention computation. To address (C2), the accuracy loss caused by the 22-bit accumulator for FP8 Matmul of PV, we propose two-level accumulation strategy that uses an FP32 buffer to accumulate the values from the 22-bit accumulator after each block Matmul of PV. This method confines the errors to the block range. Additionally, we design an optional technique to enhance the accuracy of the 22-bit accumulator. Specifically, we could smooth V by subtracting the average value of its channel dimension and adding the subtracted item to the attention output to maintain the correctness.
Our approach. To address (C1), we first discover that the per-block quantization of Q, K in SageAttention does not achieve sufficient accuracy for INT4 quantization. Simul- taneously, to avoid the extra latency caused by per-token dequantization, where each GPU thread corresponds to mul- tiple quantization scales, we propose an adequately precise quantization method that incurs no additional dequantization overhead. Specifically, we introduce a per-thread quanti- zation approach based on the mapping between the GPU threads and the memory layout of matrices dictated by the PTX mma instructions. This method groups tokens cor- responding to the same thread for quantization and dequan- tization, ensuring that each thread is associated with a single quantization scale. This approach achieves much better ac- curacy performance than per-block quantization with no additional dequantization overhead. Second, for the signifi- cant channel-wise outliers in matrices Q and K, we adopt smoothing K in SageAttention and further propose an effec- tive method to remove these outliers in Q. Specifically, we propose subtracting the average value of the channel dimen-
Performance. We offer a high-performance implementation of SageAttentions on RTX4090 and L20 GPUs. This implementation achieves a peak performance of 481 TOPS on the RTX4090, outperforming FlashAttention2 and xform- ers by approximately 3x and 4.5x, respectively. To support NVIDIA Hopper GPUs, which lack native INT4 tensor core support, we additionally provide EageAttentions variant that quantizes Q, K INT8. FlashAttention3,, in contrast, is tailored to and only compatible with the Hopper architecture. Moreover, SageAttent ion2 -bb matches the speed of FlashAttentiin on Hopper GPUs, while delivering much better accuracy. For example, on popu- lar video generation models, our method does not compro- mise end-to-end accuracy, whereas (fassatteentioo33(pp3)) brings noticeable degradation, as visualized in Fig. 7 and We extensively evaluate the end-to-end metrics of state-of-the-art text, image, and video generation models. SageAttent ion2 can accelerate models in plug-and-
2

---

SagcAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
play way with negligible loss in end-to-end metries.
À = [A/6A]] 84 max( A The dequantization pro- is an element-wise scaling. For example, for per-tensor dequantization, &5,62 (AB) AB x ÓĄ0B.
2. Preliminary
2.1. FlashAttention
Table 1. Specdup compared to matrix multiplication in FP16 with an FP32 accumulator.
The attention computation can be formulated as: S =
QKT /vd. P = (S)), O = PV, where (S)), = exp(S,g)/Ex exp(Su)). The matrices Q, K, and V each have dimensionality N d, and S, P are N x N. d is typically small, e.g., 64 or 128, and N can be thousands or millions. The time complexity of attention is O(N2), pri- marily due to two matrix multiplications (QKT and PV). both with complexities of dd. FlashAttention (Dao 2024) is a GrU- attention implementation, which tiles Q. K, and from the token dimension into blocks =2 {K. )2g) [Vi]2= with block sizes of by by. by tokens, respectively, where ng, fgg, ny are the number of tiles, and b2- by- FlashAttention computes the output ma- trix O in parallel in tiles. Each streaming multiprocessor (SM) computes block O; (corresponds to Q:) by itera- tively loads Kj, Vj for each j. and update the output with online softmax MMilakov& & Gimelshein] 2018):
GPU
MM Input
MM Accumulator
Speedup
FP16
FP16
2x
RTX4090
FP8
FP32
2x
140,, L20
FPI6
FP16
1x
H100
FP8
FP32
2x
2.3. SageAttention
Based on the block tiling in FlashAttention (Dao et al. 2022). SageAttention (Zhang et al., (025c) quantizes Q, K to INT8 in a per-block granularity, that is, each Qi, K; has a separate scalar scale: &Q, = max(|Q4))/ (227))27,.. ōk; ( In this way, the product Sty in Eq. can be approximated as Siy x Q.KJ x (6,,5K//o /vd.. To maintain accuracy, SageAttention proposes a preprocess- ing technique by subtracting the token- token-wise mean from K. Additionally, SageAttention keeps P.y and V3 in FP16, but utilizes an FP16 accumulator (rather than FP32) for com- puting the product Piş V,. Reducing accumulator precision can accelerate matrix multiplication (MM) on the RTX4090 GPU. However, other GPUs, such as L.20, L40,, or H100, do not exhibit this behavior, as shown in Table I
(1) S. -KKK /Vd, (ma, Pvg) = ā (m Siy), 1,j exp(ex: (x. m.y Mi.j- 1 + rowsum(P.))
Oi- diag (exp(mi.j-1 m 0 +
where may and lyy are by dimennaiional vectors, initialized with -oo and 0 respectively. a(-) an online softmax operator: myj = max rS:)) Psj = exp(S.j - m.j). Finally, the output is computed as O, =
3. SageAttention2
In this section, we introduce SugeAttention2. an efficient and accurate quantized attention. The workflow of SageAt- tention2 is shown in Fig. 3 We quantize Q, K to INT4 and P, V to FP8 to maximize the efficiency and propose several techniques, including -moostoo,,,,,, phr-thread quantization, and two-level accumulation to preserve the accuracy, which we shall discuss in subsequent subsections.
diag (aag( -10mm
2.2. Quantization
A matrix multiplication C' = AB can be accelerated with quantization as:
(6A, A) (AA, (6g, = B C b (AB)
3.1. Smooth Q
db is quantizer which converts high-precision (e.g.. FP32 or FP16) matrix A to low-precision format A (e.g., INT4 or FP8) with a seale 5A, and 2b-1 is a de- quantizer to convert back to high-precision. We should have (A) x A. The actual matrix multiplication ÁB is performed with low precision. In modern GPUs, low- precision matrix multiplication is usually multiple times faster than higher-preeision ones. Quantizers differ in nu- merical format and granularity, e.g., how many elements ("quantization group"") share a common scale factor. For example, an INT4, per-tensor quantizer first computes the scale as the maximum absolute value of the entire tensor, scales the elements to the maximum representable range of INT4 [-7, +7], and then casts to INT4 with rounding:
First, we discuss how accurately compute QKT with INT4. The numerical range of INT4 is notably restrictive. This affects quantization due to the presence of outliers (Lin et al. 2025). Given the INT4 range [-7, +7], any element will be quantized to zero if it is more than 14 times (0.5 vs 7) smaller than the largest element in the group. Since outliers are much larger than other elements, it is likely that many non-outlier elements are quantized to zero, leading to substantial accuracy degradation. Therefore, to keep the quantization accurate, we need to keep the largest element small, making the magnitude of elements as uniform as possible. Such technique is called smoothing.
Here, we propose smoothing technique inspired by
3

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
LLama3.1
Token
Tole
Cogvideox
Taben
Token
Tokon
Channel
Channel
Channel
Channel
Channel
Channel
Figure 2. Typical examples of tensors' data distrihution in attention.
Qreis
KF P16
VPP16
quantization, and GEMV (general matrix vector multiplica- tion) for computing AS can be fused into single kernel, which reads the off-chip Q and K only once.
Optional
K moon
Vmoan
II
(2) attention: execute the low-p owwprecisionn GEMM, dequan- tize, and add back the vector AS: Sij L6a 6x, (Q.K))+
QFP16
KpPI6
VFP16
Mat Mui
ASi These operations are all done on chip, and the dequan- tization and vector addition only add a marginal overhead compared to the expensive mma operation for MM.
Per-thread Quantization
Per-channel Quantization
9
VEPS
Importantly, (QQ), ((;)) are quantized rather than Q₁ K,, Since the smoothed matrices are much smaller in magnitude and contain fewer outliers, the quantization accuracy be significantly improved. A theoretical analysis of the benefit of smoothing is included in Appendix A.5
1NTTA
KINTA
SageAttention2 Keme!
Optional
Remark, Classical techniques to improve the activation- weight MM, such as per-channel quantization, or SmoothQuant (Xiao et al., 2023) are not applicable here for the query-key MM in attention. Per-channel quantiza- tion cannot be applied to Q, K because the quantization must be conducted along the outer axis (token dimension) of QKQ On the other hand, both Q and K have significant outliers, so trading the quantization accuracy between them with SmoothQuant cannot work effectively, as shown in Sec. 4 Here, we utilize the unique token similarity pattern in attention to derive dedicated quantization method for Q and K. The previous work SageAttention only smooths K, it is less accurate than our method.
Figure 3. Workflow of tontont Smooth Q, K, V. GEMV to obtain AS. Per-thread quantize Q, K and per- channel quantize V. Perform the SageAttention2 kernel. Correct the output.
SageAttention (Zhang et al. 2025c). SageAttention ob- served that Q, K for all tokens are actually highly similar, with only small variations between different tokens (Fig. 2 shows the heatmap distributions of the Q, K, and V ran- domly sampled from I. ama3.. 1 (Dubey et al.) 2024) and Cogvi deox Yang et al. (025b)). We propose to smooth K as SageAttention does and further smooth Q by subtract- ing a common mean of each block:
A.9 shows an ex- ample from Cogvi deo of the distribution of Q with and without smoothing Q. We can find that with smoothing Q, the range of INT4 is utilized more uniformly and fully. Ta- ble presents end-to-end metries for different quantization methods with and without smoothing Q+K on L,l ama and CogvideoX (2b). The results demonstrate that smooth- ing Q+K offers significant accuracy benefits. Also, Table and 17 show that the order of effectiveness is smoothing Q+K smoothing Q > smoothing K > other baselines.
Empirical results. Fig. 20 in Appendix
)(2)) = QQ Ti (K₃) = K5- -
(2)
where & = mean (Q:), mean(K) are x D vectors, the mean is conducted along the token axis, and qi, & are broadcasted to tokens in a block and a tensor for subtraction.
With the decomposion, we have Siz- = Q.KT (qs (T + (Kx)) = q.ET + [T(K)) + *(Q)) ((Q))((K)) = (Q)))(K,)) (T)))++Ss,, +AS +b. Here, AS.y- = (K(K)) an 1 N vector, and b q + (Q,)) is an N x vector. We do not need compute b since adding common bias to an entire row of S does not affect the result after softmax. Therefore, we can accelerate KJ with INT4 by the following two stages:
3.2. INT4 Per-thread Quantization
Orthogonal to smoothing, we can mitigate the problem of outliers by refining the quantization granularity so that the number of affected elements by outliers becomes smaller. Although per-token quantization offers de- tailed level of granularity, it results in significant overhead
(1) preprocessing: smooth Q, K according to Eq. (2) ap-
ply quantization (6g., Q.) = v)((((.)), (6x, K2) = dh (x)))) and compute AS,, = (K,) Smoothing,
4

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Quantization each Warp
MMA.m16n8k64
C AB
B
Lokena
Toker
Tokena
Tokens
Tokena
Tokens
Token6
Token?
Dequantization
0
Tokeno
K
(K
Token1
80l [0]
Token1
Token3
Token2
k [0]
TokenO
Token1
T4
T4
T6
T8
0
TO
To
Token2
Token8
Token3
3 K X
Token4
Token8 Token9
Token9
Token5
Token10
Token6
6Q
Token11
A
Token7
T28
T29
129
T20
130
T80
C
Tokenä
[1
0
k [
(
Token9
(QRT)
p 3
Token10
Token56
X
Token11
Token24
Token57
Token12
0
TEL TTL
Token25
Token58
Token13
STA1 T81
8d[7]
Token59
6x [3]
Token14
Kj
Token15
T28
128
129
120
130
130
751
1 16 (QRT)
w
Each Thread Corresponds Scale
Figure 4. An example of per-thread quantization. The left figure shows the correspondence between the quantization scales and the tokens each GPU warp. The right figure shows the correspondence hetween quantization tokens and GPU threads a MMA. m 6m 8 instruction, showing that each GPU thread only corresponds to one quantization scale in li and óg in dequantization.
during dequantization. Specifically, each GPU thread in per-token quantization must handle multiple quantization scales, leading to high latency of the dot product of the quantization scale vectors & and og- SageAttenntinn uses per-block quantization, where each block Qi (b, tokens) and K; (b, tokens) have a single quantization scale. Such quantization strategy could achieve an accuracy perfor- mance close to per-token quantization and avoid the high dequantization overhead. However, quantizing Q and K to INT4 demands finer quantization granularity. To address this, we propose per-thread quantization, more precise and granular approach than the per-block quantizer, also without the additional overhead of the vector dot product between &e and 6K-
Empirical results. As shown in Table 6 and Table 15 we compare the average and worst accuracy of INT4 quantiza- tion at per-token, per-thread, per-block, and per-tensor gran- ularity using real Q, K, V across all layers of Cogv iieex. Results indicate that the accuracy of per-thread quantization is very close to per-token and significantly outperforms other granularities. Moreover, Table 19 shows that per-thread quantization incurs almost no speed degradation, while per- token quantization introduces noticeable overhead due to the reduced hardware efficiency.
3.3. FP8 quantization for PV
We now to the MM PV, where P:j exp(Sij mij) is the unnormalized quantity according Eq. ① The dis- tribution of P is unique and differs from other activations. First, we note that Sij maj S 0, so Pij E [0, 1] (S and E apply element-wise). We find that P often consists of many small elements, but their sum is non-negligible (e.g., 5000 elements around 10-4), In this case, we must represent small elements accurately. INT quantization is unsuitable for this setting, since it distributes the quantization points evenly within the numerical range. SageAttention (Zhang et al. 2025c) choose to retain P and V in FP16, and ac- celerate the MM by decreasing the accumulator precision. However, this strategy is only effective on very few GPUs.
Specifically, each block of Q, i.e., Qú. in SageAttention will be split into Cur segments and processed by eu GPU warps in GPU streaming processor (SM). We call each segment of Q: as Q.., and k. K₃ since K, is shared among warps. Then, each warp containing 32 threads uses the mma m1 64 PTX instruction (NVIDIA) for the Kn According to the layout requirement of this instruc- tion, we find that Q [8k + i] could share one guantization scale, and K8[8k= 24] together with K, [8k 2i + 1] could share quantization scale. Such quantization method is more fine- -reiigd with additional overhead. This is because it assigns different GPU threads to distinct quanti- zation groups based on the MMA instruction layout, with each thread performing dequantization only using a single quantization scale value. We show an example of per-thread quantization in Fig. The detailed formulation is shown in Equation g and Fig. [8(p (please see Appendix f. for more detail).
We propose to quantize P. V to FP8 with 4 exponent bits and 3 mantissa bits (E4M3). The numerical range of E4M3 is [-448, +448]. We quantize P with a static scale: 8p 448 since the original P elements are already in [0, 1]. address the channel-wise outliers shown in Fig. Empirical results in Table and Table To the average and worst accuracy of different
We quantize V per-channel
5

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Algorithm 1 Implementation of SageActention2.
Input: Matrices ((FP1)), (FP16), ((FP6)) E R block size by: bEv, warp count Preprocessing: K = K - mean(K), ). (ōv, V) = vV(V). /pperchhel
Divide Q to T.m. = N/bb blocks [Q:]; divide K, and V to T. = N/bb. blocks (K:). (Vi): for io 1 to Ti, do
Q: =een(Q:); (fq: Q.) vQ( q.) //pertthread;
for j in [1, [,] do
(ōx, Kÿ) (K;) /ppr-thread, w range(cw), st = w Our
Sylst at + Cwl
Sa (Matmul (Qi at + rwl, KJ + GEMV KJ) ) / Paralleled by warps. The Vod0x is
illustrated in Fig.
maj = max (m. (owmax((Ss)) P.j- = exp(Si: - mss)) lij- = e 1 mi rowsum ; (FP22)) = Matmul((P, P., * 448) t0FF8..c43)) Vg);
(FP32)) = diag(e 1-:-----m- -))) 1 + 01,
end for Load By into SM; O diag(h,T,)) Os., (FP32) /418 dv Write O:: end for return 0 = {0.}
data types used for P, V across all layers of Cogvi i deox. The accumulator is always 32-bit. We can see that the ac- curacy of E4M3 very close to that of FP16 and superior to E5M2 and INT8. Most modern GPUs have tensor cores that support FP8 Matmul operations, which are twice fast as those using FP16.
accumulated to Oig in the high FP32 precision.
Remark. The two-level accumulation strategy also implemented in CUTLASS (NVIDIA) 2023) and Deep- Gemm (DeepSeek-Al et al., 2024) for computing weight- activation products in linear layers. To the best of our knowl- edge, we are the first to discover and investigate the effect of the FP22 accumulator and implement the two-level accu- mulation for attention.
3.4. FP32 MMA Buffer for FP22 Accumulator
While FP8 quantization for PV above is theoretically accu- rate in simulation, we observe that the actual CUDA imple- mentation suffers consistent accuracy degradation. After narrowing down the problem, we find that the accumulator for the mma (f32fBf8f3) instruction on the Ada and Hopper architecture is actually FP22, specifically with sign bit, exponent bits, and 13 mantissa bits. Specifically, for mma (f32f8f8f32) instruction C AB + D, where ini-
Optional smooth V technique. We also figure out an- other way to mitigate the accuracy loss due to the FP22 accumulator when V possesses channel-wise biases: Vm mcan axis 0), V = the correctness of the attention computation, it is only neces- sary to add Vm to the final calculation of O: O O + Vm This is because the sum of each row of the P matrix equals 1, so PV = =Vn.
V----VV V-V-- Furthermore, maintain
A, B are FP8 matrices and C, D are FP32 matrices, tialize the A, B to 7ero and vary D) to test the data type of the accumulator. When D is initialized with 1 sign bit, 8 exponent bits, and 13 mantissa bits, the value of C' exactly matches D. However, when D is initialized with more than 13 mantissa bits, the value of C' is equal to D with its least significant 10 mantissa bits zeroed out (i.e., truncated). Con- sequently, matrix multiplication of PV, quantized to FP8, incurs a certain degree of accuracy loss compared to using an FP32 accumulator.
Remark. For details on smoothing V, see Appendix A.3 This technique is optional and not employed in our main experiments, as it provides significant benefits only when V exhibits channel-wise bias, which are absent in some models, such as I.lama.. 1 (see Fig. 2
4. Experiment
Main result. SageAttention2 is faster than FlashAtten- tion2 and xformers by about 3x and 4.5x. Moreover, SageAttention! matches the speed of FlashAttn3(fp8) on the Hopper GPUs and is much more accurate than FlashAtt3(((88), SageAttentio22 maintains end-to-end met- ries across language, image, and video generation models.
To mitigate this accuracy loss, we adopt a two-level accumu- lation strategy, which uses an FP32 buffer to accumulate the values of V3 in FP22. Specifically, we rewrite Eq. 11 as Ra,- = V, Ozj = diag (exp(mi,,-1 m.:)) O 1 + Rij. Here, two sets of accumulators Ray and Oij are maintained in the register. Raj computed with the mma (F32f8f8:32) instruction, providing 22 effective bits, which is sufficient since we only accumulate over small number of by tokens (e.g., by 64). Then, R.y is
4.1. Setup
We validate the effectiveness of SageAttention2 across a diverse set of repre-
Models.
6

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
RTX4090, (Head dim 128, causal =
False)
RTX4090, (Head dim 128, causal =
-Tuee)
Speed (TOPS)
800
xformers
Torch
SageArtn1
FlashAltn2
SageAttn2-4b
SageAtin2-8o
800
Torch
xformers
x>
FlashAltn2 SageAttn11
SageAttn2-BB
SageAttn2-4b
600
1424
177
1430
1981
600
4733
3ns
522
323
322
400
108
B22
200
DO
DOM
105
DOM
1016
200
G0
CON
LOS
105
0
Sequence Length
Sequence Length
Figure 5. Speed comparison
between SageAttent inniteentiionn on2 and haselines (RTX4090, headdim= 28).
Table 2. End-to-end metrics across text, image, and video generation models. X indicates inability to generate results for evaluation.
Model
Attention
WikiText (Ppl.) J.
Lambda (Acc.)
MMLU (Acc.) t
Longbench
Full-Precision
6.013
0.815
0.635
49.40
HadmdAttn
7.872
0.762
0,500
44.07
SmoothAttn
7.180
0.783
0,541
44.69
Llama3.
SageAttention
6.017
0,812
0.634
49.55
SageAttn2-4b
6.256
0,798
0.607
48.79
SageAttn2
6.019
0,811
0.634
49.59
Full-Precision
7.241
0,432
0.743
49.78
HadmdAttn
7.989
0,435
0.669
45.97
SmoothAttn
8.943
0,449
0.592
42.20
GLM4
SageAttention
7.243
0.433
0.744
49.79
SageALtn2-4b
7.352
0.433
0.725
49.23
SageAttn2-8b
7.242
0.432
0.745
49.60
Model
Attention
CLIPSIM
CLIP-T
VQA-a f
VQA-t 1
FScore 1
Full-Precision
0.1778
0.9979
70.231
70.928
2.507
HadmdAttn
0.1576
0.9933
8.990
2.299
CogvideoX (1.5-BB)
SmoothAttn
0.1559
0.9950
8.812
2.277
SageAttention
X
FlashAttn3-Ep6
0.1562
0.9918
6.531
2.181
X
Sage
0.1721
0.9978
57.729
52.989
2.884
SageAttn2-B5
0.1775
0,9980
69.492
74.415
2.487
Full-Precision
0.1783
0.9995
82.516
75.934
0.604
HadmdAttn
0,1727
0.9989
7.514
0.762
0.175
Hunyuan Video
SmoothAttn
0.1739
0.9988
6.987
0.609
0.148
SageAttention
0.1786
0.9995
82.496
79.843
0.597
FlashAttn3-Cp8
0.1742
0.9941
4,433
1,460
X
SageAttn2-4b
0.1751
0.9995
81.478
65.371
0.610
SageALtn2-8b
0.1782
0.9996
81.786
75.354
0,586
Full-Precision
0.1798
0.9986
45.549
65.4116
1.266
HadmdAttn
0.1733
0.9980
9.053
25.133
0.704
SmoothAttn
0.1687
0.9978
3.383
3,480
0.241
Mochi
SageAttention
0.1800
0.9987
48.707
63.763
1.269
FlashAttn3-fp8
0.1762
0.9982
14.964
13.711
0).457
SageAttn2-4b
0.1783
0.9986
35.955
43.735
1.137
SageAttn2-88
0.1797
0.9986
46.760
64.901
1.255
Model
Attention
FID 1
sFID 1
CLIP 1
Full-Precision
10,960
16.648
26.180
1.009
EadndAttn
11.353
18.495
26.123
0.965
SmoothAttn
11.149
19.017
26.109
0.959
Flux
SageAttention
10.944
16.641
26.171
1.008
SageAttn2-4b
10.577
17.497
26.141
0,998
SageAttn2-8b
10.927
16,723
26.175
1.009
Full-Precision
14.105
15.646
25.505
0.902
EadndAttn
14,259
15,909
25.513
0.886
Stable-Dif fusion3.5
SmoothAttn
14.161
15.649
25.510
0.887
SageAttention
14.140
15.678
25.503
0.902
SageAttn2-4b
14.097
15.397
25.487
0.895
SageAttn2-8b
14.106
15.647
25.499
0.901

---

SagcAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Full precision Attention
SageAttention2-80
SageAttention244b
SmoothAttn
HadmdAttn
Figure 6. Comparison examples from Hur.yanVideo, prompts sampled from open sora prompt sets.
Full- Precisioo
SageAttn2-88
FlashAttn3((p8)
Figure comparison example using SageAttn2-8b and FlashAttention3 CogvideoX-I.5
sentative models from language, image, and video generation. Specifically, we conduct experiments on ten models: Llama2 (7B) (Touvron et al.] [023),
for NVIDIA Hopper GPUs, which lack native INT4 tensor
core support, and incorporates all techniques described in Sec. exxcept for smoothing Q.
Llama3.1 (8B) (ubeyy et al.) (024). and GLM4
Baselines. (1) SmoothAttn. Following Qserve (Lin
(9B) (GLM et al. 2024) for text2text, Cogvideox
et al. 2025), we apply smooth quant for Q, K with smoothing factor a 0.5. (2) HadmdAttn. Follow- ing Quarot (Ashkboos et al. 2024), we apply random Hadamard transformation for Q, K before INT4 quanti- zation. (3) SageAttention (Zhang et al. 2025c), which uses smoothing K, INT8 per-block quantization for Q, K, and FP16 for P, V. (4) FlashAttn's (fp8), the FP8 sion of FlashAttention33 which only
(2B), CogvideoX (1.5-BB) (Yang et al. 2025b).
HunyuanVideo (Kong et al. 2024). and Mochi 1 (Team 2024) for text2video, Flux (schnell) (Black Forest Labs
2023) and Stable-Diffu fftsiifoffooo.... (turbo) (Stability AI 2023) text2image, and TIMM (Wightman, 2019) for image classification.
Datasets and metrics. For Details about the datasets and metrics we used, please refer to Appendix. A.7
Hopper GPUs.
Table Two kernel implementations of SageAttention2.
4.2. Speed and Accuracy of Kernels
Kernel (Q(()), vx(K)
vp((),, (V)
Kernel Speed. We compare the speed of SageAttention2 against baselines using head- dim=64 and headdim=128 both with and without Causal
SageAttn2-4b INT4 per-thread FP8 per-block and per-channel SageAttn2 Sb INTS per-thread FP8 per-block and per-channel
Implemetation. We implement two attention kernels as shown Table using CUDA, The 8-bit variant adapted
Mask (Vaswani 2017). Detailed setup can be found
8

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Accuracy. Table 3 and 17 show the average accuracy of different methods with INT4 4 Q, K and FP8 P. V across all layers of CogvideoX. The results indicate the accuracy of SageAttn2-4b is superior to other baselines.
Table Average accuracy across all layers of Cogvideox using different smoothing methods.
Method
CosSim
Relative L1 |
RMSE
None
80.04%
0.3906
0.2223
MadmdAttn
79.77%
0.3782
0.2180
SmoothAttn
90.21%
0,3383
0.1952
4.3. End-to-end Performance
Smooth K
98.07%
0,1493
0.0743
Smooth h 98.30%
0,1250
0.0712
Metrics loss. We assessed the end-to-end metrics of various models using EageAttents ion2 compared to baselines. Detailed evaluation results are presented in Table 2 The results indicate that bageAttn2
Smooth Q+K
99.46%
0.0648
0.0334
outperforms all base- lines and maintains most of the end-to-end accuracy across all models. Additionally, SageAt
Table 5. End-to-end metries comparison, where Q, K are quan- lized into INT4, while P, stay in full precision.
incurs almost metrics loss across various models. More experiment results other models are shown in Appendix A.9
Smooth (Q+K)
Llama3.1 (Lambda)
Llama3.1 (WikiTeet))
CogVidenX (vqa-t)
Q.
Full-Precision
81.5%
6.013
75.360
INT4
x
72.6%
11.698
24.670
Visible image and video examples. Fig.
Quantization
/
80.8%
6.219
75.147
and 9 show some visible comparison examples from HunyuanVideo, Mochi and Cogvideox. We can ob- serve that SageAttn2 does not introduce any visible differences compared full-precision attention, whereas SageAttn2-42 has minor differences but is much better than the baselines.
Table Average accuracy across all layers of Cogvideed using different quantization granularities.
Method
Cos Sim
Relative
RMSE
Per-token
99.45%
0.0649
0.0335
Per-thread
99.45%
0.0622
0.0313
End-to-end speedup. We compared the original genera- tion latency and the latency using Lageettentio for models with long sequence lengths Table observing sig- nificant speedup effects. For instance, SageAttentt on2 achieved 1.8x speedup in CogvideoX (1.5-5B) without any metrics loss (SageAt 8b)). SageAt fur- ther accelerated these models but with little metrics loss.
Per-block
98.03%
0.1492
0.0744
Per-tensor 97.15%
0.1 1..8
0.0865
Table Average accuracy using different data types of (P, V)
across all layers of CogvideoX, where (Q, K") are smoothed.
Q. K
P.V V
Cos Sim t
Relative L1 +
RMSE |
INTS
77.05%
0.5618
0.5044
E5M2
99.20%
0.0905
0.0903
4.4. Ablation Study
INT4
E4M3
99.44%
0.0683
0.0347
Overhead of techniques we proposed. As shown in Ta-
FP16
99.45%
0.0649
0,0335
ble 18 the overhead on kernel speed of per-thread quantiza- tion, smoothing Q, and two-level accumulation are 0.35%, 3.7%, and 0% compared to the attention kemnl.
Table End-to-end generation latency using SageAttent
(The latency of L L.ama3.1 is the time to first token generation using different sequence lengths).
Benefit of smoothing V. The experiment showing the bene- fit of smoothing is shown in Appendix. A.4
Saggettn
Model
GPU
Original
2-8b
2-4b
5. Conclusion
Cogvideos (2B)
RTX4090
86
54 s
52
Cogví deox (1.5-5B)
RTX4090
1040 s
577
555
We introduce SageAttention? an efficient and accu- rate quantized attention. First, we propose to quantize matrices (Q. K") in a thread-level granularity and (P, V) to FP8. Second, we propose a method to smooth Q, en- hancing the accuracy of QKT Third, we propose two- level accumulation strategy enhance the accuracy of FP8 PV. tag ion2 is faster than FlashAtten- tion2 and xformers by approximately 3x and 4.5x, respec- tively. Moreover, SageAt tentiitentions matches the speed of ( on the Hopper GPUs, but offers significantly higher accuracy. Extensive experiments con- firm that our approach maintains end-to-end metrics across language, image, and video generation models.
HunyuanVii deo
1.20
2221 s
1486
1435
Moch1
1.20
2336 s
1316
1190 s
T.1 T.1ama3.1 1 (48K token)
RTX4090
9.2 s
5.7
5.6
Llama3. 1. (100K token)
L20
39.9
25.4
23.2
in Appendix A.8 Specifically, Fig. shows the speed across varying sequence lengths on RTX4090, indicat- ing that SageAttn2-4b and SageAttn2-8hb are approximately 3x and 2.7x faster than FlashAttention2, and about 4.5x and 4x faster than xformers, respectively. Fig.
15 and i6 in Appendix A.2 show more speed results on RTX 4090, 1.20, H20, H100 GPUs.

---

SagcAttention2:
Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Acknowledgment
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Re, C. Flashat- tention: Fast and memory-efficiet exact attention with IO-awareness. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Pro- cessing Systems, 2022.
This work was supported by the NSFC Projects (Nos. 92270001, 62376131). J.Z is also supported by the XPlorer Prize.
Impact Statement
DeepSeck-Al, Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., Dai, D., Guo, D., Yang, D., Chen, D.. Ji, D., Li, E., Lin, F., Dai, F., Luo, F., Hao, G., Chen, G., Li, G., Zhang, H., Bao, H., Xu, H., Wang. H., Zhang, H., Ding, H., Xin, H., Gao, H., Li, H., Qu, H., Cai, J. L., Liang, J., Guo, J., Ni, J., Li, J., Wang. J., Chen, J., Chen, J., Yuan, J., Qiu, J., Li, J., Song, J., Dong, K. Hu, K., Gao, K., Guan, K., Huang, K., Yu, K., Wang, L.,, Zhang, L.., Xu, L., Xia, L.,, Zhao, L., Wang. L., Zhang, L., Li, M., Wang. M.. Zhang, M., Zhang, M., Tang. M., Li, M., Tian, N., Huang, P., Wang, P., Zhang. P.. Wang. Q., Zhu, Q.. Chen, Q., Du, Q., Chen, R. J., Jin, R. I.., Ge, R., Zhang. R., Pan, R., Wang, R., Xu, R., Zhang, R., Chen, R., Li. S. S., Lu, S., Zhou, S., Chen, S., Wu, S., Ye, S., Ye, S., Ma, S., Wang. S., Zhou, S., Yu, S., Zhou, S., Pan, S., Wang, T., Yun, T., Pei, T., Sun, T., Xiao, W. L., Zeng. W., Zhao, W., An, W., Liu, W., Liang, W., Gao, W., Yu, W. Zhang, W., Li, X. Q. Jin, X., Wang, X., Bi, X., Liu, X., Wang. X., Shen, X., Chen, X., Zhang, X., Chen, X., Nie, X., Sun, X., Wang, X., Cheng, X., Liu,
This paper presents work that aims to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
References
Ashkboos, S., Mohtashami, A., Croci, M. L., Li, B., Cameron, P., Jaggi, M., Alistarh, D., Hoefler, T., and Hensman, J. Quarot: Outlier-free 4-bit inference in ro- tated LLMs. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., and Li, J. LongBench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 2024.
I: Long Papers), pp. 3119-3137,
X., Xie, X., Liu, X., Yu, X., Song, X., Shan, X., Zhou, X., Yang, X., Li, X., Su, X., Lin, X., Li, Y. K., Wang,
Black Forest Labs. Flux. black-forest-lappe flux 2023.
com/
Y. Q., Wei, Y. X.. Zhu, Y. X., Zhang, Y., Xu, Y., Xu, Y., Huang, Y., Li, Y., Zhao, Y., Sun, Y., Li, Y., Wang, Y. Yu, Y., Zheng, Y., Zhang, Y., Shi, Y., Xiong, Y., He, Y., Tang, Y., Piao, Y., Wang, Y., Tan, Y., Ma, Y., Liu, Y., Guo, Y., Wu, Y., Ou, Y., Zhu, Y., Wang, Y., Gong, Y., Zou, Y. He. Y., Zha, Y., Xiong, Y., Ma, Y., Yan, Y., Luo, Y., You, Y., Liu, Y., Zhou, Y., Wu, 7. F., Ren, 7. 7., Ren, 7., Sha, Z., Fu, Z., Xu, Z., Huang, Z., Zhang, Z., Xie, Z., Zhang, Z., Hao, Z., Gou, Z,, Ma, Z., Yan, Z., Shao, Z., Xu, Z., Wu, Z. Zhang, Z., Li, Z., Gu, Z., Zhu, Z., Liu, Z., Li, Z., Xie,
Chen, Y. Qian, S., Tang, H., Lai, X., Liu, Z., Han, S., and Jia, J. Longlora: Efficient fine-tuning of long-context large language models. In The International Conference on Learning Representations, 2024.
Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q., Mo- hiuddin, A., Kaiser, L.. Belanger. D. B., Colwell, L. J., and Weller, A. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
Z., Song, Z., Gao, Z., and Pan, Z. Deepseek-vä v3 technical report. arXiv preprint arXiv:2412.19453 2024.
Deng, J., Dong, W., Socher, R., Li, J..JJ,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEFE conference on computer vision and pattern recognition, pp. 248 2255, leee, 2009.
Chu, X., Tian, Z., Wang, Y., Zhang, B., Ren, H., Wei, X., Xia, H., and Shen, C. Twins: Revisiting the design of spatial attention in vision transformers. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Ad- vances in Neural Information Processing Systems, 2021.
Dettmers, T., Pagnoni, A., Holtzman, A., and Zettlemoyer, L. QLORA: Efficient finetuning of quantized LLMs. In Thirty-sevenin Conference on Neural Information Pro- cessing Systems, 2023.
Chu, Y., Xu, J., Yang, Q., Wei, H., Wei, X., Guo, Z Leng, Y., Lv, Y., He, J., Lin, J., Zhou, C., and Zhou, J. Qwen2- audio technical report, 2024.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXi::2407.2178. 2024.
Dao, T. Flashattention-2:
22: Faster attention with better paral- lelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024.
10

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Fu, T., Huang, H., Ning, X., Zhang, G., Chen, B., Wu, T., Wang, H., Huang, Z.,, Li, S., Yan, S., Dai, G., Yang. H., and Wang, Y. Moa: Mixture of sparse attention for automatic large language model compression. arXiv preprint ar arXiv:2406. 14909, 2024.
Jiang, H., LI, Y., Zhang, C., Wu, Q., Luo, X., Ahn, S., Han, Z., Abdi, A. H., Li, D., Lin, C.-Y., Yang, Y., and Qiu, L. MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Gao, Y. Zeng, Z., Du D., Cao, S., So, H. K.-H., Cao, T., Yang. F., and Yang. M. Seerattention: Learning intrinsic sparse attention in your llms. arXiv preprint arXiv:2410.13276 2024.
Kamradt, G. Limtest needle in a haystack pressure testing 1lms. [attps://github. com/ gkamradt/
LLMTest NeedleInAlaystaack 2023.
GLM, T., Zeng. A., Xu, B., Wang, B., Zhang, C., Yin, D.,
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference mu- chine learning, pp. 5156-5165, PMLR, 2020.
Rojas, D., Feng, G., Zhao, H., Lai, H., Yu, H., Wang, H.,
Sun, J., Zhang. J., Cheng. J., Gui, J., Tang, J., Zhang, J., Li, J. Zhao, L., Wu, L.. Zhong, L., Liu, M., Huang, M., Zhang. P.. Zheng. Q., Lu, R., Duan, S., Zhang, S., Cao, S., Yang. S., Tam, W. L.., Zhao, W., Liu, X., Xia, X., Zhang, X., Gu, X., Lv, X., Liu, X., Liu, X., Yang, X., Song, X., Zhang, X., An, Y., Xu, Y., Niu, Y., Yang, Y., Li, Y., Bai, Y., Dong. Y., Qi, Z., Wang. Z., Yang, Z., Du, Z., Hou, 7., and Wang, 7. Chatglm: A family of large language models from glm-130b to glm-4 all tools. arXiv preprint arXiv:2406. 12793, 2024.
Kong, W. Tian, Q., Zhang, Z., Min, R., Dai, Z., Zhou, J., Xiong. J., Li, X., Wu, B., Zhang, J., Wu, K., Lin, Q. Wang, A., Wang, A., Li, C., Huang, D., Yang, F., Tan, H., Wang, H., Song. J., Bai, J., Wu, J., Xue, J., Wang. J., Yuan, J., Wang, K., Liu, M., Li, P., Li, S., Wang, W., Yu, W., Deng, X., Li, Y., Long, Y., Chen, Y., Cui, Y., Peng, Y., Yu, Z., He, Z., Xu, Z., Zhou, Z., Xu, Z., Tao, Y., Lu, Q., Liu, S., Zhou, D., Wang, H., Yang, Y., Wang, D., Liu, Y., Jiang, J., and Zhong, C. Hunyuanvideo: A systematic framework for large video generative models. arXiv preprint 2:Xiv:2412.03600 2024.
Hendrycks, D., Basart, S., Mu, N., Kadavath, S., Wang, F., Dorundo, E., Desai, R., Zhu, T., Parajuli, S., Guo, M., et al. The many faces of robustness: A critical analysis of out-of-distribution generalization. In Proceedings of the IEEE/CVF international conference computer vision, pp. 8840-834444 2021a.
Lefaudeux, B., Massa, F., Liskovich, D., Xiong, W., Caggiano, V., Naren, S., Xu, M., Hu, J., Tintore, M., Zhang. S., Labatut, P., Haziza, D., Wehrstedt, L., Reizen- stein, J., and Sizov, G. xformers: A modular and hack- able transformer modelling library, hhttps://githhi com/ htte//accebooreeeaccc//////
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021b.
2022
Li, D., Kamko, A., Akhgari, E., Sabet, A., Xu, L., and Doshi, S. Playground v2.5: Three insights towards enhanc- ing aesthetic quality in text-to-image generation. arXiv preprint arXiv: 17245, 2024.
Hessel, J., Holtzman, A., Forbes, M., Le Bras, R., and Choi, Y. Clipscore: A reference-free evaluation metric for image captioning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7514-7528, 2021.
Li, K., Wang, Y., Peng, G., Song, G., Liu, Li, and Qiao, Y. Uniformer: Unified transformer for efficient spatial-temporal representation learning. In International Conference on Learning Representations, 2022.
Heusel, M., Ramsauer. H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 2017.
Lin, Y., Tang, H., Yang. S., Zhang, Z., Xiao, G., Gan, C., and Han, S. @Serve:w4a8Ke V4 quantization and system co- design for efficient LLM serving. In Eighth Conference on Machine Learning and Systems, 2025.
Hu, Y., Huang, W., Liang, Z., Chen, C., Zhang, J., Zhu, J., and Chen, J. Identifying sensitive weights via post- quantization integral. arXiv preprint arXiv:2503.01901, 2025.
Liu, Y., Cun, X., Liu, X., Wang, X., Zhang, Y., Chen, H., Liu, Y., Zeng, T., Chan, R., and Shan, Y. Evalcrafter: Benchmarking and evaluating large video generation models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22139 22149, 2024.
Jelinek, F., Mercer, R. L., Bahl, L. R., and Baker, J. K. Per- plexity a measure of the difficulty of speech recognition tasks. The Journal of the Acoustical Society of America, (((11):633SS63 1977.

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Liu, Z.,, Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Venkataramanan, S., Ghodrati, A., Asano, Y. M., Porikli, F, and Habibian, A. Skip-attention: Improving vision transformers by paying less attention. In The Twelfth International Conference on Learning Representations, 2024.
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. In International Conference on Learning Representations, 2022.
Milakov, M. and Gimelshein, N. Online normalizer cal- culation for softmax. arXiv preprint 2018.
Wang. H., Ge. S., Lipton, Z., and Xing, E. P. Learning robust global representations by penalizing local predic- tive power. Advances in Neural Information Processing Systems, 32, 2019.
NVIDIA. 8.5.
Parallel thread execution isa version
https::/docs.avice www.ii.ccccccuuaaaaaaaaaaaa cuda/ parall -tw.eeh-eeecuttiooooo
Wang. S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 04760, 2020.
NVIDIA. CUTLASS: CUDA Templates for Linear Algebra Subroutines and Solvers. GitHub repository, 2023. URL
Wang, Y., Chen, Z., Chen, X., Zhu, J., and Chen, J. Frame- bridge: Improving image-to-video generation with bridge models, arXiv preprint arXiv:2410.15377,, 2024.
/etps://githule www/WWIII////uatet cuttlasss
Panayotov, V., Chen, G., Povey, D.. and Khudanpur, S. Librispeech: An asr corpus based on public domain au- dio books. In 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. doi: 10.1099/CAASP22115577735
Wightman, R. Pytorch models. http:t://sts::::/t/tttt hoooodhimmmmmodee..ttps/s////1/ com rwi foom/wwighene wtwp:/a//pp...ooooieeeeeeeeeeeee 2019.
5206-5210, 2015.
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, N.-Q., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The lambada dataset: Word prediction requiring broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Compu- tational Linguistics (Volume I: Long Papers), pp. 1525- 1534, 2016.
Wu, H.,, Zhang, E., Liao, L.., Chen, C., Hou, J., Wang, A., Sun, W., Yan, Q., and Lin, W. Exploring video quality assessment user generated contents from aesthetic and technical perspectives. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 20144 20154, 2023.
Xi, H., Yang. S., Zhao. Y., Xu, Li, Lin, Cai, H., Zhang, J., Li, D., et al. Sparse videogen: Accelerating video diffusion transformers with spatial-temporal spar- sity. In International Conference on Machine Learning (ICML), 2025.
Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Shah, J.. Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate attention with asynchrony and low low-precision. The eeihth- Annual Conference on Neural Information Processing Systems, 2024.
Xiao, C., Zhang, P., Han, X., Xiao, G., Lin, Y., Zhang, Z., Liu, Z., and Sun, M. Infllm: Training-fre ffree long-cone long-context extrapolation for lims with an efficient context memory. First Workshop on ICMI. 2024, 2024a.
Foundation Models@
Stability AL 3.5.
Introducing stable diffusion https://stabiiity.at news/
Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38099, PMLR, 2023.
trtroducing-stable--de
2023.
Team, G.
Mochi 1.
com/
genmoa 1 //wodlollss.sl 2024.
Touvron, H., Martin, I.., Stone, K., Albert, P., Almahairi, A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288. 2023.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ei- ficient streaming language models with attention sinks. In The Tweifth International Conference on Learning Representations, 2024b.
12

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Xu, J., Liu, X., Wu, Y., Tong, Y., Li, Q., Ding, M., Tang, J., and Dong, Y. Imagereward: Learning and evaluat- ing human preferences for text-to-image generation. In seventh Conference on Neural Information Pro- cessing Systems, 2023.
Zhang, P., Chen, Y., Su, R., Ding, H., Stoica, I., Liu, Z., and Zhang, H. Fast video generation with sliding tile attention. arXiv preprint
Zhang, P., Huang, H., Chen, Y., Lin, W., Liu, Z., Sto- ica, I., Xing, E. P., and Zhang, H. Faster video dif- fusion with trainable sparse attention. arXiv preprint arXiv:2505. 13389, 2025i.
Yang, S., Xi, H., Zhao, Y., Li, M., Zhang, J., Cai, H., Lin, Y., Li, X., Xu, C., Peng, K., et al. Sparse videogen2: Accel- erate video generation with sparse attention via semantic- aware permutation. arXiv preprint arXiv:2505. 18875, 2025a.
Zhang, P., Wei, J., Zhang, J., Zhu, J., and Chen, J. Accurate int8 training through dynamic block-level fallback. arXiv preprint arXiv:2503.0050 050000000 2025j.
Yang. Z.. Teng. J., Zheng, W., Ding, M., Huang, S., Xu, J., Yang, Y., Hong, W., Zhang, X., Feng, G., et al. Cogvideox: Text-to-video diffusion models with an pert transformer. In The Thirleenth International Confer- ence on Learning Representations, 2025b.
Zhang, X., Chen, Y., Hu, S., Xu, Z., Chen, J., Hao, M., Han, X., Thai, Z.,, Wang, S., Liu, Z.,, and Sun, M. ooBench: Extending long context evaluation beyond 100K tokens. In Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Pro- ceedings of the 62nd Annual Meeting of the Association for Computational Linguisties (Volume I: Long Pupers), pp. 15262-15277, 2024.
Yu, W., Luo, M., Zhou, P., Si, C., Zhou, Y., Wang, X., Feng, J., and Yan, S. Metaformer is actually what you need for vision. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10819- 10829, 2022.
Zhao, M., Zhu, H., Xiang, C., Zheng, K., Li, C., and Zhu, J. Identifying and solving conditional image leak- age in image-to-videe video diffusion model. arXiv preprint arXiv:2406. 1500255,..1735 2024.
Zhang. J., Huang. H., Zhang, P., Wei, J., Zhu, J., and Chen, Sageattention2: Efficient attention with smoothing q and per-thread quantization. 2025a.
Zhao, M., He, G., Chen, Y., Zhu, H., Li, C., and Zhu, J. Riflex: A free lunch for length extrapolation in video diffusion transformers. arXiv preprint orXiv:2502.1584, 2025a.
Zhang, J., Li, G., and Su, J. Sage: framework of precise retrieval for rag. In 2025 IEEE 41st International Confer- ence Data Engineering (ICDE), pp. 1388-1401. IEEE Computer Society, 2025b.
Zhao, T., Fang, T., Huang, H., Liu, E., Wan, R., Soedar- madji, W., Li, S., Lin, Z., Dai, G., Yan, S., Yang, H., et al. Vidit-q: Efficient and accurate quantization of dif- fusion transformers for image and video generation. In International Conference on Learning Representations, 2025b.
Zhang. J., Wei, J., Zhang, P., Chen, J., and Zhu, J. Sageatten- tion: Accurate 8-bit attention for plug-and-pla and-play inference acceleration. In The International Conference on Learn- ing Representations, 2025c.
Zhang. J., Wei, J., Zhang, P., Xu, X., Huang, H., Wang, H., Jiang, K., Zhu, J., and Chen, Sageattention3: Microscal- ing fp4 attention for inference and an exploration of 8-bit training. arXiv preprini ar arXiv:2505. 1594, 2025d.
Zheng, K., Lu, C., Chen, J., and Zhu, J. Dpm-solver-v3: Im- proved diffusion ode solver with empirical model statis- tics. Advances in Neural Information Processing Systems, 36:55502-55542, 2023.
Zhang. J., Xiang, C., Huang, H., Wei, J., Xi, H., Zhu, J., and
Zheng, K., Chen, Y., Mao, H., Liu, M.-Y., Zhu, J., and Zhang, Q. Masked diffusion models are secretly time- agnostic masked models and exploit inaccurate categori- cal sampling. arXiv preprint arXiv:2409.02908, 2024a.
Chen, J. Spargeattn: Accurate sparse attention accelerat- ing any model inference. In International Conference Machine Learning (ICML), 2025e.
Zhang, J. Xiang, C., Huang, H., Wei, J., Xi, H., Zhu, J., and Chen, J. Spargeattn: Training free sparse attention accelerating any model inference. 2025f.
Zheng, K., He, G., Chen, J., Bao, F,, and Zhu, J. Diffusson bridge implicit models. arXiv preprint 15X. 2024b.
Zhang, J., Xu, X., Wei, J., Huang, H., Zhang, P., Chendong, X., Zhu, J., and Chen, J. Sageattention2++: A more efficient implementation of sagcattention2. arXiv preprint arXiv:2505.21 136, 2025g.
Zheng, K., Chen, Y., Chen, H., He, G., Liu, M.-Y., Zhu, J., and Zhang, Q. Direct discriminative optimization: Your likelihood- based visual generative model secretly gan discriminator. arXiv preprint arXiv:2503.0 103, 2025.
13

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Zheng, Z., Peng, X., Yang, T., Shen, C., Li, S., Liu, H., Zhou, Y., Li, T., and You, Y. Open-sora: Democratiz- ing efficient video production for all. arXiv preprint arXiv:2412.20404, 2024c.
14

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
A. Appendix
A.I. Visible Comparison Exmaples
Full precision Attention
SageAttention2-81
SageAttention2-43
SmoothAttn
HadmdAttn
Figure 8. Comparison examples from CogvideoX (2B), prompts sampled from open prompt sets.
A.2. Additional Kernel Speed Comparison
Fig. 10
13
4
15 and 16 compare the speed of SageAttent 2io2 against baselines using headdim=64 and
128,, both with and without Causal Mask (Vaswani 2017), on RTX4090, L20, H100, and H20 GPUs.
Table summarizes the performance gain of different attention methods against baselines on various modern GPUs.
Table 9. Speedup of different attention methods on various GPUs.
Method
3090
4090
A100
L40
L20
H100
H20
FlashAttention2
1.00
1.00
1.00
1.00
1.00
1.00
1.00
Flastinten
X
X
X
X
X
1.37
1.57
FlashAttention3 (fp8)
X
x
x
2.63
3.06
Sage
1.97
1.96
1.37
1.45
1.24
1.53
1.52
SageAttention2
2.93
2.60
2.46
2.61
3.12
Table 10. An accuracy example real tensors of Cogvideox model with without smoothing V.
Smooth V
Cos Sim
Relative L1
RMSE
X
98.25%
0.1980 0.0406
0.2387
99.75% %
0.0773
A.3. Smoothing
As shown in Fig. 17 this strategy could enhance the accuracy of FP22 for values in PV for the following reasons: Each row P spans value range from to 1, and each column V in some models consistently features channel-wise biases that are exclusively positive or negative, for instance, ranging between and in CogvideoX. Consequently, the values of PV could be quite large. However, the floating-point number representation range uniform-it is denser near zero. Therefore, by subtracting the mean V. along the channel dimension from V, the values of PV will be closer to higher representational precision (see Fig. for a visual demonstration). Additionally, to maintain the correctness of the attention computation, it is only necessary to add V to the final calculation of O: O = O V This is because the sum of each of the P matrix equals 1, so PV," Vn In other words, this method decomposes V into two parts: Vm and V. For V. it centers the values of each column around zero, which leads to the dot product result between row from the quantized P matrix and a column from the quantized matrix being closer to zero. This makes the representation of FP22 more accurate. Meanwhile, Vm is retained in FP16 and is added to O at the end without causing loss of precision for the Vi part.
zero, resulting
15

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Mochi
Full-Precision
SageAttn2-8b
FlashAttention3 (fp8)
HunyuanVideo
Full-Precision
SageAttn2-8b
FlashAttention3
(fp8)
Figure 9. comparison example using SageAttn2-8h FlashAttention3 on Mochi and HunyuanVidec.
A.4. Experiment of Smoothing V
Table 10
shows the attention accuracy on real tensors sampled from CogvideoX with and without smoothing V. It demonstrates that smoothing V could improve the accuracy of EageAttention2 when quantizing Q. K to INT4 and
P.V to FP8. We find that smoothing generally effective for diffusion models (Zheng et al. 2023: 2024b a 2025
Zhao
al.
2024
2025a Wang al. 2024).
A.5. Theoretical Analysis of Smoothing
In this section, we analyze the benefit of smoothing from a theoretical perspective. Let X Nxd be activation tokens of dimension d. Following (Dettmers al. 2023), we suppose that an activation token follows an Gaussian distribution
16

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
RTX4090, (Head dim 64, causal = False)
RTX4090, (Head dim 64, causal True)
800
Torch
FlashAttn2
SageAttn2-80
800
Torch
FlashAltn2
SageAttn2-80
Speed (TOPS)
xformers
SageAttn1
SageAttn2-4b
xformers
SageAttn1
SageAttn2-48
600
400
223
E334
33699
212
326
449
320
1303
1445
400
600
368
289
350
406
3115
577
338
11435
125
122
GOM
DOX
GOM
77
2K
8K
16K
32K
2K
8K
16K
32K
Sequence Length
Sequence Length
Figure 10. Speed comparison between SageAttent onn and baselines (RTX4090, headdim=64).
L20, (Head dim 64, causal
False)
L20, (Head dim 64, causal True)
Torch
FlashAttn2
SageAttn2-ee 8b
Torch
FlashAltn2
SageAltn2 8b
Speed (TOPS)
400
xformers
SageAttn1
SageAttn2-44 4b
400
xformers
SageAltn1
SageAttn2 4b
30
230
248
1202
259
2205
162
2105
266
235
99
745
202
2044
254
200
140
142
200
30
69
DOM
DDW
DDM
DDM
DDM
00
1K
2K
4K
8K
16K
32K
1K
2K
4K
8K
16K
32K
Sequence Length
Sequence Length
Figure 11. Speed comparison between SageAttention2 and baselines (L20, headdim=64).
L20, (Head dim
False)
L20, (Head dim 128,
Torch
FlashAttn2
SageAttn2-8b
FlashAttn2 SageAttn1
SageAttn2-89
Speed (TOPS)
xformers
SageAttn1
SageAttn2-4b
xformers
SageAttn2-4b
249
262
269
1273
238
264
214
273
201
207
188
200
133
200
165
1.08
1226
1122
GOO
LOOM
DOM
.OM
2K
8K
32K
2K
8K
16K
32K
Sequence Length
Sequence Length
Figure 12. Speed comparison between SageAttention2 and baselines (L20, hcaddim=128),
H100, Head Dim 64, causal = False
H100, Head Dim = 64, causal = True
Xformers
SageAttention FlashAttn3-11p) fp16
Fpps
Xformers FlashAttn2
SageAttentioni
RlshAttna-fa f-f
Speed (TOPS)
1000
FlashAttn2
SageAttn2-89 2-bb
1000
FlashAttn3-fp11
SageAttn2-8b
750
596
750
588
274
482
500
383
500
2007
800
1268
1339
250
250
1K
2K
4K
8K
16K
32K
1K
2K
4K
8K
16K
32K
Sequence Length
Sequence Length
Figure 13. Speed comparison between SageAttentions and baselines (H100, headdim=64).
,,,). Mt, Σ²), where μ (F1, P2,...,.. Pa) and E is diagonal matrix with Σ² diag(o?, oz, .ff.....)).. Further, suppose that different token X; is i.i.d. sampled from the same distribution.
Suppose the absolute maximum value in quantization group (Hu et al. 2025 Zhang et al. 2025] is M, and the bit width 6, then there are 2 quantization levels. Under the round-to cond-to-nearess strategy. the expected quantization error is 2 2 which is proportional to the maximum absolute value in the quantization group. So a smaller absolute maximum value leads to smaller quantization error.
2M
17

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
H100, Head Dim 128, causal
False
H100, Head Dim 128, causal True
1500
Xformers
SageAttention1
3lashAttn3-fp8
1500
Xformers FlashAltn2
SageAttentionl
FlashArta--pa fp8
Speed (TOPS)
FlashAttn2
FlashAttn3-1pp fp16
SageAttn2-88
FlashAltn3-fp1e
SageAttn2-Bb
1000
89
900
1000
77
329
500
166
L8L
500
1180
1K
2K
4K
8K
16K
32K
1K
2K
4K
8K
16K
32K
Sequence Length
Sequence Length
Figure 14. Speed comparison between SageAttention2 and baselines (H100, headdim=128).
H20, Head Dim 64, causal False
H20, Head Dim 64, causal True
Xformers
SageAttention
FlashAttn3-fp8
Xformers
SageAttention1
FlashAttn3-fp8
Speed (TOPS)
400
FlashArtn2
FlashAttn3-f116
SageAttn2-8b
400
FlashArtn2
FlashAttn3-fp16
SageAttn2-80
133
271
222
216
250
260
200
200
10:00
2K
8K
16K
32K
1K
2K
8K
16K
32K
Sequence Length
Sequence Length
Figure 15. Speed comparison between SageAttenticn2 and baselines (H20, headdim=64).
H20, Head Dim 128, causal
False
H20, Head Dim
causal
Xformers
SageAttention1
FlashAtti3--f88
Xformers FlashAttn2
SageAttention1
FlashAttn3-fp8
Speed (TOPS)
400
FlashAttn2
FlashAttn3-fpl)
SageAttn2-80
400
FlashAttn3-fr
SageAttn2-88
266
260
264
200
200
60
62
1K
2K
4K
8K
16K
32K
1K
2K
4K
8K
16K
32K
Sequence Length
Sequence Length
Figure 16. Speed comparison between SageAttention2 and baselines (H20, headdim=128).
Error of FP22 comapred FP32
FP22 Error for PV
Heatmap of column of
FP22 Error
value PV without smoothing
value PV with smoothing v
Heatmap of a row of
0.8
Heatmap of a column V mean(V)
13.2410
95065.6208
0
20000
40000
60000
80000
0.4 -0.2
0.0 0.2
0.4
Value
Figure 17. An example of dot product precison of 1 and column of V presented by FP22 data type.
After smoothing, we have:
N
Yú = Xij
1
Σ
X
(3)
N
km
18

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
and Yy also follows a Gaussian distribution. The mean and variance of Yop can be calculated as follows:
1 13
E B
E[Xai]
Pj =0
(4)
N
N
Var[Yiy] = Varl
N
Xi5l
+
Σ
N
Varl
Xayl
(N 112
N2 of + (N - N2 0²
(N
N of
1)
(5)
N
N
km 1,pgh
So Yyy has smaller mean and variance compared to Xty. Following the property of Gaussian distribution, we have:
P([Yy] < >
p( c). Ve
(6)
So the distribution of Xı is more concentrated toward 0 after smoothing. Then we know that
P(absmax(Y)) <) II
P(Y <d>>
II
P([X;) <) P(absmax(X)) <
j=1
j=1
So this makes the distribution of absolute max value in token more concentrated towards 0, leading to smaller quantization error.
A.6. Per-Thread Quantization Formulation
R|C
T1: (c0,
TB: (c0,
1
T4: (co, c1)
TS: (c0, c1)
Te: (c0, c1)
T7: (c0, ct)
2
7
T28: (c0, c1)
T29: [c0, c1)
T30: (c0, c1)
T31: [1)
[c2, c5)
c5)
T4: (c2, c3)
T5: [c2, cB)
T7: (3)
128: (x2, c3)
129:
130: (c2, c3)
T31: (2,
Figure Memory layout of INT4/INT8 a warp. C and D is of shape 16x8, Fach thread only holds 4 out of the 128 elements.
for accumulator matrix C and D in D A ammmmnn 32 threads (TO T31)
To further clarify the per-thread quantization, we first introduce the INT4 MMA instruction of Tensor Core, and then give the formulation of per- wer-thread quantization.
Tensor cores, first introduced in NVIDIA's Volta architecture, are specialized units designed for efficient matrix-multiply- and-aceumulate (MMA) operations. Their usage significantly enhances computational efficiency and performance in Al and performnrfe computing (HPC) workloads. Tensor cores compute small tiles of MMA operations, specifically D A B C' warp (32 contiguous threads) basis. Each thread in the warp holds a fragment input matrices and will get fragment of output matrix computation result. The INT4 mma m1 64 tensor
operation computes the product of 16 x 64 INT4 matrix A and 64 x 8 INT4 matrix B, both stored in registers. It accumulates the result into a 16 x INT32 matrix C, also stored in registers, and returns the final product matrix D, which has the same shape (16 x 8), data type (INT32), and storage location (registers). Each thread holds only 32 of the input and output data. Fig. 18 extracted from the PTX document (NVIDIA) shows the memory layout of matrix and D among 32 threads in a warp. Each thread only holds of the 128 result elements.
19

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
i.g = [(n * 8 Cw /w/b,)]
(m)) b
qi [6gg] = ={8 x (n%8) + 1(n *
).
n NO
bu cw
max(1 [qa:(āsolll))
8q [og1
[q: [i8g]]
Q19.(16u)
= [26g]
dq [d6y]
i8k = (n * 4/ 4/bx]]
(8)
kn (is]
= 18 x (n%8) + [n/bx] * 3u
{8 x (n%8) + + [n/bal by3. e N]
max(1 K [kn [isx]]])
5k [1ax]]
7
K R [
R =
5x [ox)
By ensuring results held by each thread share common dequantization scale (belong to the same quantization group), we can avoid the overhead associated with per-token quantization. Leveraging this observation, we design per-thread quantization as formulated in Eq. 8 where Car is the count of GPU Warps, b, and by are the block size of Q, K, and n is the token index of Q, K. For typical block size of b, 128, by = 64 and warp number cw
(as used in FlashAttention2)), 0, 1, ... .7 can be made into one quantization group and key tokens j, 1 + + 8 + j, + j, ... 56 + j. 57 + j (j = 0, 2, 3) can be made into one quantization group, as visualized in Fig. This design aligns with the memory layout of output matrix D of tensor core shown in Fig. 18] ensuring that each thread only needs one Q scale and one K scale for dequantization.
each warp processes a tile of 32 query tokens and 64 key tokens. Query tokens 8 + i, 16 + i, 24 i
As result, this approach creates 32 quantization groups for Q (8 for each of the 4 warps) and quantization groups for K in 128x64 block, providing 32x and 4x finer granularity compared to per-block quantization for query tokens and key tokens, respectively. Table Table 15 show the accuracy gains by using per-thread quantization. Per-thread quantization achieves accuracy that closely matches per-token quantization, without introducing any kernel speed degradation (see
Table 18 and 19
A.7. Datasets and Metrics in Experiments
Datasets. Text-to-text models are evaluated on four zero-shot tasks: WikiText (Merity et al. 2022) to assess the model's prediction confidence, LAMBADA (Paperno et al. 2016) evaluate contextual understanding. MMLU (Hendrycks et al., 2021b) for measuring knowledge across various subjects, and Longbench (Bai et al. 2024) for comprehensive assessment of long context understanding capabilities. Text-to-video models are evaluated using the open-sora (Zheng et al. (024c) prompt sets. Text-to-image models are assessed on MJHQ-30K (Li et al. 2024). TIMM is evaluated on on three image datasets: ImageNet (Deng al. 2009), ImageNet-Sketch (Sketch) (Wang al. 2019). and ImageNet-Rendition (ImageNet- (Hendrycks et al. 202 aa).
End-to-end metrics. For text-to-text models, we use perplexity (ppl.) (Jelinek et al. 1977) for WikiText, Accuracy (Acc.)
for LAMBADA and MMLU, and Longbench score (Bai et al.) 2024). For text-to-video models, following Zhao et al.
(2025b), we evaluate the quality of generated videos on five metrics: CLIPSIM and CLIP-Temp (CLIP-T) (Liu et al. 2024)
to measure the text-video alignment; (VQA-a) and (VQA-t) to assess the video aesthetic and technical quality, respectively; and Flow-score (FScore) for temporal consistency (Wu et al. 2023). For text-to-image models, generated images are compared with the images in MJHQ-30K dataset in three aspects: FID (Heusel et al.) (017)) and sFID (Salimans et al. 2016) for fidelity evaluation, Clipscore (CLIP) (Hessel et al. 2021) for text-image alignment, and ImageReward (IR) (Xu et al.
2023 for human preference. For TTMM, we use classification accuracy.
Accuracy metrics. We use three metrics to assess the accuracy of quantized attention output O' compared to atten- tion output in full-precision O: First, we flatten O' and O into vectors in the shape of 1 x n. Then, Cosine similar- ity: CosSim 0000/VO 2F OZ, Relative LI distance: L1 = 1 - O'l/ [O0, Root mean square error:
20

---

SageAttention2:
Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
RMSE ((/m)
A.8. Kernel Benchmark Setup
We benchmark kernel speed with a batch size and 32 attention heads across variety of sequence lengths. Benchmarks conducted using head dimensions of 64 and 128, both with and without Causal Mask (Vaswani 2017 To generate input tensors for benchmarking, For floating-poin data types, inputs drawn from Gaussian distribution with mean 0 and standard deviation 1, while for integer data types, inputs are uniformly sampled within the representation range: 128, 127] for INT8 and [-8, 7] for INT4.
follow standard practices adopted in prior works such as FlashAttention (Dao et al.) 2022]
Table 11. End to-end metrics on Llama2 (7B).
Model
Attention
WikiTeet (Ppl.)
Lambda (Acc.)
MMLU (Acc.) 1
Full-Precision
5.823
0.886
0,439
HadmdAttn
6.771
0.860
0.360
SmoothAttn
6.717
0.867
0.392
Llama2
SageAttention
5,824
0.887
0.439
SageAttn2-4b
5,912
0.881
0.428
SageAttn2 8b
5.828
0.886
0.438
Table 12. End-to-end metrics on Cogvideox (2B).
Model
Attention
CLIPSIM t
CLIP-T
VQA-a
VQA-t f
FScore
Full-Precision
0.1836
0.9975
77.605
75.360
3.006
HadmdAttn
0.1742
0.9877
29.780
23.985
0.499
CogvideoX (2B)
SmoothAttn
0.1741
0.9870
41.703
47.043
0.624
SageAttention
0.1833
0.9976
76.997
71.360
2.988
SageAttn2-4b
0.1821
0.9973
77.368
74.906
2.603
SageAttn2 8b
0.1829
0.9977
76.532
74.281
2.941
Table 13, End-to-end metrics an image classification model.
Model
Attention
ImageNet (Acc.) t
Sketch (Acc.) 1
ImageNet-r (Acc.) 1
Full-Precision
84.79%
45.32%
59.55%
HadmdAttn
84.50%
44.89%
58.80%
SmoothAttn
84.40%
44.68%
58.73%
TIMM
SageAttention
84.74%
45.38 %
59.95% %
SageAttn2-4b
86.67 %
45.24%
59.29 %
SageAttn2-81
84.79%
45.39
59.57%
Table 14. Comparison with ((ast))ttemtioo3((((()))) on Llama-3-2640 (8B) on InfiniBench
Zhang et al. 2024)
(H100 GPU).
Attention
Eng.Sur Sum
Eng.QA
Eng.MC
Code.Debug
Math.Find
Retr.PassKey
Retr.Num
Retr.KV
Avg.
Full-Precision
18.03
12.5
64.19
24.37
18.29
100.0
100.0
7.0
43.05
F @1ashAttn3-ip
19.03
11.73
55.90
22.59
22.57
100.0
100.0
0.4
41.53
SageAttention2
18.17
12.46
64.19
25.63
17.43
100.0
100.0
6.6
43.06
Table 15, Worst accuracy across all layers of CogvideoX using difffernt quantization granularities.
Method
Cos Sim f Relative L1 +
RMSE $
Per-token
96.76%
0.1916
0.0775
Per-thread
96.72%
0.1932
0,0776
Per-block
90.68%
0.3615
0.1490
Per-tensor
85.85%
0.4687
0.2261
21

---

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
NIAH Llama-3-8--262k
NIAH Lama-3-88-2622 w FlashAttn3-1p8,
NIAH Llama-3-88-262k w/ SageAttn2
1.0
1.0
0.0
11.0
11.0
11.0
0.8
0.8
0.8
22.0
22.0
22.0
(%)
33.0
0.6
33.0
0.6
0.6
Depth
56.0
56.0
56.0
67.0
67.0
67.0
78.0
78.0
78.0
02
0.2
0.2
89.0
6f
30x
70k
102K
134K
1666
1986
2306
262k
5ox
70k
107K
134k
108k
199k
2306
262k
y0k
100k
1346
166k
198k
2306
2622
0.0
Token Limit
Taken Limit
Token Limit
(a) Full Precision
(b)
(c) SageAttention2
Pigure 19, Needle In Haystack results on L1 amama-- 3-262k (8B).
x10
Quantized value without smoothing Q
Quantized value with smoothing Q
3.0
2.5
Frequency
0.0
5
Figure An example quantized value distribution before after smoothing Q.
Table 16. Worst accuracy using different data types of (P, V) across all layers of a Cogvideox model, where (Q, K) are smoothed.
Q.K
V.V
Cos Sim
Relative L1 |
RMSE
INT8
19.52%
0.9579
1.4483
E5M2
94.94%
0.2327
0.2361
INT4
E4M3
96.70%
0.1956
0.0779
FP16
96.76%
0.1916
0.0775
Table 17. Worst accuracy all layers of CogvidecX using different smooth methods.
Method
CosSim f
Relative L1 $
RMSE +
None
4.83%
0.9979
0.7784
HadmdAttn
4.85%
0.9978
0.7785
SmoothAttn
64.49%
0,9262
0.7204
Smooth 90.86%
0.3565
0,1464
Smooth 93.10%
0.2989
0.2195
SageAttn2-4b
96.71%
0.1956
0.0779
A.9. Additional Experiments and Analysis
Additional Results. Table II
12 and
13
show results of SageAttention2 and other baselines on Llama2 (7B),
CogvideoX (2B) and TIMM.
Results of Suner-Long Context. We further conduct experiments on super-long context using Llama 1amma-3-666 -3-262k (8B
22

---

SageAttention2:
Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
Table 18. Overhead of per-thread quantization, smoothing and two-level accumulation techniques measured on L20 GPU.
Table 19. Comparison different quantization granularities mea-
sured on L20 GPU, with QK in INT4 and PV in FP8.
Method
TOPS
Granularity
TOPS
Attention (INT4 + FP8)
284
Per-tensor
286
Per-thread quantization
283
Per-block
284
+ Two-level accumulation
283
Per-thread
283
+ Smoothing Q
273
Per-token
268
InfiniBench (Zhang al. 2024) and Needle-i--a-Haystack (NIAH) (Kamradi 2023), with sequence lengths reaching up to 262k tokens an H100 GPU. Since Hopper GPUs lack native INT4 tensor support, SageAttent ion2-8b for this evaluation. We compare it against FlashAttenniion3((ps)) ensuring both methods operate under the same bit width. Results are shown in Table 14 and Fig T9 SageAttent on maintains model performance even under super-long context, while (lashAttentionn3(p5)) suffers from end-to-end accuracy degradation.
Results of Audio Tasks, We evaluate Owen2-Audic (7b) (Chu et al. 2024). a speech-to-text model, on the ASR task using the Librispeech (Panayotov et al. 2015) test split and measured performance with the WER metric (Word Error Rate). As shown in Table 20 SageAttent ion2 consistently outperforms the baselines, highlighting its effectiveness in muditerelated models and benchmarks.
Table 20. End-to-end metrics on Qwen2-Audi (7B).
Model
Attention
Test-Clean +
Test-Dev +
Full-Precision
1.74
4.01
HadmdAttn
1.77
4.05
SmoothAttn
1.76
4.01
Qwen2-Audh
SageAttention
1.74
4.02
SageAttn2-4D
1.73
3.99
SageAttn2-8b
1.72
4.03
23


Reviews:
[{"review_id": "YFZ25S182o", "reviewer": "Reviewer_1zQu", "summary": "SageAttention2 introduces a new way to enhance the accuracy and efficiency of attention through a 3 pronged process: firstly, it introduces an INT4 matrix-multiplication technique for query-key and FP8 matmul technique for attention weight and values; secondly, it proposes a smoothing technique for queries to reduce loss of accuracy due to outliers in INT4 quantization; thirdly, it investigates the reason for loss in accuracy due to FP8 matmul for attention weights and values, and proposes a way to overcome the loss.\n\nThe main algorithmic contributions are as follows:\n1. Providing an efficient INT4 quantization scheme that adapts to thread-level granularity to move from multiple quantization scales in a GPU Warp to single quantization scale per warp to reduce dequantization overhead. This technique takes adavantage of the Nvidia's PTX Tensor Core warp-level `mma` instruction layout to partition Q and K thread blocks into quantization groups such that each quantization group shares a single quantization scaling factor to perform dequantization.\n\n2. To reduce the accuracy overhead from loss of precision due to INT4 range, the paper adapts a technique to smooth Q similar to a previous technique of smoothing K from SageAttention paper. This technique subtracts the mean along the token dimension to utilize the INT4 range `[-7, 7]` more uniformly. This reduces the effect of outliers by making them smaller in magnitude, which in turn allows for a significant preservation of accuracy.\n\n3. For `P*V`, the paper choses FP8 quantization over INT or FP16 due to better representation of P values numerically and prevalence of FP32 accumulator in wide variety of GPUs respectively. The authors observe that the loss in accuracy due E4M3 FP8 quantization of PV (P per-block, V per-channel quantization) is due to the nature of FP22 accumulator in the CUDA implementation, and they propose a 2 stage FP32 accumulation buffer during the computation of the online-softmax PV matmul to counteract the loss in accuracy due to the CUDA implementation using FP22 accumulator.\n\n4. Lastly the authors propose the same smoothing technique for V by subtracting mean along the token dimension to counteract the accuracy loss due to FP22 accumulation, but its benefit is only found to be observed when V possesses channel-wise bias, so this smoothing is kept as optional.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 1zQu,  \nThank you for your valuable suggestions and questions. Below, we address each point raised.\n\n---\n>### Comment1\n\n**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.\n\n|Attention|Cossim ↑|Relative L1 ↓|\n|-|-|-|\n| SageAttention|0.99995|0.00733|\n| Sage2-8b|0.99982|0.01573|\n| Sage2-4b|0.99460|0.06480|\n| Sage2-4b (without smooth Q)|0.9498|0.27305|\n\n---\n>### Comment2\n\n**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:\n\n|Attention|test-clean ↓|test-other ↓|\n|-|-|-|\n|Full-Precision|1.74|4.01|\n|HadamardAttn|1.77|4.05|\n|SmoothAttn|1.76|4.01|\n|SageAttention|1.74|4.02|\n|Sage2-4b|1.73|3.99|\n|Sage2-8b|1.72|4.03|\n\nWe can see that SageAttention2 consistently outperforms the baselines, highlighting its effectiveness in audio-related models and benchmarks.\n\n---\n>### Comment3\n\n**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:  \n1) The attention speedup.  \n2) The proportion of total latency attributed to attention.\n\nFor example:  \n- If a model spends 50s on attention and 50s on other operations, a **2×** attention speedup reduces latency by **25s** (total: 75s).  \n- A **2.5×** attention speedup yields a **30s** reduction (total: 70s).\n\nMoreover, 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.\n\n---\n>### Comment4\n\n**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:\n\n|Model|Cossim ↑|Relative L1 ↓|\n|-|-|-|\n|Sage2-8b|0.9997|0.02133|\n|Sage2-8b (without two-level accumulation)|0.9939|0.17843|\n\n---\n>### Question1\n\n**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.\n\n---\n>### Question2\n\n**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:\n\n1. **Speed Impact**:\n   - It requires additional CUDA Core operations for dequantization and has the same Tensor Core count as INT8. Therefore it won't bring speed gains.\n\n2. **Accuracy Impact**:\n   - INT4 quantization introduces greater accuracy loss than INT8.\n\n---\n>### Question3\n\n**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.\n\nA 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.\n\n---\n[1] AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration  \n[2] SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models  \n[3] OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models\n\n---\n---\nIf you feel your concerns have been resolved, we would greatly appreciate it if you consider raising the score.", "claims_and_evidence": "Most of the claims made in the paper over accuracy and performance are backed by sufficient evidence. \n\nObservations:\n1. Smoothing of Q+K seems to be an effective technique in preservation of accuracy with INT4/8 matmul of QK.\n2. The granular per-thread quantization seems to be on-par with per-token quantization strategy for INT4 QK matmul in terms of accuracy, though enough evidence hasn't been presented on the said overhead of per-token dequantization. Theoretically, per-thread quantization shouldn't have any performance overhead, which is what has been presented in the table 18, and evidence of it being able to handle outliers in preserving accuracy has been tabulated in Table 4.\n3. The 2-level FP32 buffer used for FP22 accumulation leads to circumvent the accuracy loss due to mma(f32f8f8f32) CUDA implementation. Though this claim hasn't been backed by sufficient evidence in terms of accuracy over multiple tasks. Table 7 and 16 fail to mention if this 2 level was employed in FP8 matmul of PV.", "methods_and_evaluation_criteria": "The paper covers many modalities across multiple tasks involving language, image, and video generation, along with image classification involving an extensive suit of models that have attention at their core for computation.\n1. The SageAttention2 is benchmarked against multiple different renowned efficient attention algorithms, in text2text, text2video, and text2image and compared in terms of the metrics that are majorly indicative of accuracy in literature.\n2. Each concept in SageAttention2 is tested separately for average accuracy and worst accuracy, while keeping the other parameters of the attention frozen. These include but not limited to:\n```\n    a. Smoothing of Q and K, separately and together and comparing against other methods of preserving INT4 accuracy, like SmoothAttn, Hadamard transformation etc.\n    b. Differing the INT4 quantization granularities among per-tokens, per-thread, per-block, and per-tensor.\n    c. Differing the PV matmul calculation with FP8 E4M3, FP8 E5M3, FP16, and INT8 precisions while keeping the KQ matmul  INT4 methodology the same.\n```\n3. For performance, varying the different techniques of employing per-thread granularity in quantization of QK, employing the 2-level accumulation strategy for PV matmul, and smoothing of Q. Each of these techniques is ablated against the baseline of INT4 QK (K smoothed), PV FP8 matmul (with default FP32 accumulation.)", "theoretical_claims": "I checked for correctness of the following equations and their assumptions:\n1. Q.K matmul, adjusting for smoothed Q and smoothed K. All conclusions are correct based on the formulations.\n2. Online softmax equation adjusting for 2-level accumulation strategy with FP8 matmul\n3. Optional smoothing of V, and adjusting for `V_mean` addition to the output.\n4. Indices calculation for Q and K, along with the scaling factor calculations for the same, with correct assumptions for per-thread grouping of Q and K tokens.\n5. The algorithm 1 in Page 4 seems to reflect the text astutely.", "experimental_designs_or_analyses": "The paper goes over multiple experiments with regards to kernel speed comparing with xformers, SageAttention, FlashAttention2, and FlashAttention3 (for Nvidia Hopper architecture). The analysis for the experiments are found to be accurate in text with regards to TOPS. There is some room for improvement over the TOPS claim, as the end-to-end generation latency across LLMs show that switching to SageAttention2-4b over 8b might not be the best use of hardware (Table 8), given that 8b suffers almost no loss in accuracy over the full-precision attention, and has quite similar end-to-end generation latency as compared to 4b.\n\nThe paper also mentions the metric loss of SageAttention2-4b compared to 8b, which is apparent once the tabular and visual results are inspected upon.", "relation_to_broader_scientific_literature": "More and more papers are targeting the GPU hardware by leveraging the instruction set, inherent knowledge of how GPU executes code blocks. We are seeing more of warp-level algorithms, algorithms designed to reduce any unnecessary communications overhead. Algorithm co-design with innate hardware knowledge are leading to better algorithms, that are turning out to be optimal.\n\nThis paper stands together with lot of papers that are structuring up the methodology of hardware-software co-design, which is a must for lot of real-world applications.", "essential_references_not_discussed": "Not related to understanding the topics in this paper directly, but would be nice to get a short reference and methodology discussed in FlashAttention2. FlashAttention2 also implements methodology to eliminate per-warp communication which is similar to the computation done in SageAttention2 per-thread INT4 quantization with single scaling factor to eliminate dequantization overhead.", "other_strengths_and_weaknesses": "This paper provides a bright way towards deploying low-latency transformer models. With great INT4/8 and FP8 attention mechanism, lot of real world real-time deployments of LLM could be possible as long as hardware supports such tensor cores. The SageAttention2 INT4 Q*K matmul algorithm seems highly intuitive, and given the per-thread granularity, it seems to be an optimal way to design quantization groups, which takes into account both performance and accuracy.\n\nThe paper restricts itself to only one technique to deal with quantization outliers, namely, smoothing. It would be great to see a comprehensive analysis of more methods that don't introduce any overhead, but that could be better suited to handle outliers for a reduced quantization error due to INT4 precision. The fact that INT8 version of SageAttention-2 matches or is on-par with the full-precision attention in accuracy, seems to indicate that there is room for improvement in the QK INT4 matmul front.", "other_comments_or_suggestions": "1. Would like to see more comparisons between SageAttention2-4b and SageAttention-8b and SageAttention, the visual results on HunyuanVideo and CogvideoX seems to indicate the weakness of INT4 range. Maybe related to loss of information to INT4 range since the smoothed Q distribution seems to indicate a normal curve, so there might be some unrecoverable accuracy due to INT4 range.\n\n2. Could also include few benchmarks against speech related tasks, like audio generation, TTS and/or speech-to-text. \n\n3. The end-to-end generation latency of SageAttention2-4b seems to be on par with SageAttention2-8b for 4090 and L20 GPUs. Wheras the TOPS tell a different story. When it comes to performance, I suggest relying on end-to-end latency in milliseconds instead of TOPS. Since it aligns more with real-world use cases.\n\n4. More detailed ablation studies needed for the 2-level FP32 buffer accumulation in regards to accuracy, like FP8 E4M3 2-level accumulation vs baseline FP8 E4M3.", "questions_for_authors": "1. Is the reason for not including the SageAttention2-4b in H100 and H20 benchmarks due to lack of INT4 tensor cores in Hopper GPUs used?\n2. If the above is true, would it be possible to pack 2 INT4 values in INT8 for the above 2 gpu architectures?\n3. Would learning for or calibrating for these per-thread quantization scaling factors lead to reduce the gap between SageAttention-8b and SageAttention-4b instead of calculating the scaled-max-of-absolute values as scaling factors?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Dxa6wtGXbW", "reviewer": "Reviewer_KKjK", "summary": "This paper proposed several improvements on SageAttention to make it comparable with FlashAttention 3 in terms of speed but better in accuracy. The enhancements mainly focused on enabling lower precision compared to previous SageAttention, i.e. move from INT8 to INT4 for Q*K^T and from FP16 to FP8 for P*V. To address the outlier problem in Q*K^T, SageAttention2 extended the smoothing technique used in SageAttention from K-only to both Q and K and showed great improvement in accuracy. In order to further improve accuracy and reduce the overhead from INT4 dequantization, author proposed an interesting per-thread quantization scheme which will enforce each thread to read and use only one set of Q K scale. This method shows comparable accuracy with per-token, i.e. better than per-block, and costs almost no overhead. The challenge for using FP8 in P*V calculation is caused by Hopper TensorCore design, which uses FP22 accumulator instead of FP32. To alleviate the impact of FP22, author employed the two-level accumulation technique and optionally applied smoothing on V. Resulting SageAttention2 shows very promising results on language, image, and video models.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer KKjK,  \nThank you for your valuable suggestions and questions. Below, we address each point raised.\n\n---\n>**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.\n\n**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.\n\n---\n>**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.\n\n**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. \n\nWe measure the TOPS of per-token INT4 quantization and the result is as follows, with the settings aligned with Table 18:\n\n|Quantization|TOPS|\n|-|-|\n|per-block (line 1094 in our paper)|284|\n|per-thread (line 1095 in our paper)|283|\n|per-token|268|\n\nPer-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.\n\n---\n>**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.\n\n**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.\n\n---\n>**C1.** typo on Line 115/116 right column\n\n**Reply**: Thank you for pointing out the typo and we will revise our paper.\n\n---\nIf you feel your concerns have been resolved, we would greatly appreciate it if you consider raising the score.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "NA, no theoretical claims.", "experimental_designs_or_analyses": "Yes, no issues.", "relation_to_broader_scientific_literature": "It's an improvement on SageAttention and competitive alternative to FlashAttention2/3.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strength\n1. describe the key concepts clearly with good supporting data.\n2. tested on a wide range of models.\n3. demonstrated speed-up and accuracy improvement.\n\nWeakness\noverall, a nice, solid work. Here are a few minor suggestions:\n1. 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. \n\n2. 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.       \n\n3. In Appendix Table 9, 14, and Fig 18, it wasn't clear which version of SageAttention2 (4b or 8b) were used in the experiments.", "other_comments_or_suggestions": "type on Line 115/116 right column, \"... preprocessing technique **to by** subtracting the token-wise mean...\"", "questions_for_authors": "Please see Weakness above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "30PbaHzCJr", "reviewer": "Reviewer_t1YD", "summary": "This paper makes the attention computation more efficient. It uses INT4 quantization of Q and K, instead of INT8 quantization. To enhance the accuracy of INT4, this paper proposes a outlier smoothing strategy, which is well-motivated. The overall design and implementation take hardware characteristics into account. The accuracy and efficiency of proposed method are evaluated by various tasks and settings: on RTX4090, the proposed method is much faster than FlashAttention2; on Hopper GPUs, the proposed method matches the speed of FlashAttention3 while delivers much higher accuracy.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer t1YD,  \nThank you for your valuable suggestions and questions. Below, we address each point raised.\n\n---\n>**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.\n\n**Reply**: Thank you for the valuable suggestion. We analyze the quantization error as follows:\n\n**Proof**:\nLet $X \\in \\mathbb{R}^{N \\times d}$ be $N$ activation tokens with dimension $d$.\n\nFollowing [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)$. \n\nFurther, we suppose that different token $X_i$ is i.i.d. sampled from the same distribution.\n\nSuppose 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.\n\nAfter smoothing, we have:\n$$\nY_{ij} = X_{ij} -  \\frac{1}{N} \\sum_{k=1}^N X_{kj}\n$$\n$Y_{ij}$ also follows a Gaussian distribution. The mean and variance of $Y_{ij}$ can be calculated as follows:\n\n$$\n\\mathbb{E}[Y_{ij}] = \\mathbb{E}[X_{ij}] - \\frac{1}{N} \\sum_{k=1}^N \\mathbb{E}[X_{kj}] = \\mu_j - \\frac{1}{N} \\sum_{k=1}^N \\mu_j = 0\n$$\n\n$$\n\\mathrm{Var}[Y_{ij}] = \\mathrm{Var}[\\frac{N-1}{N} X_{ij}] + \\sum_{k=1, k\\neq i}^N \\mathrm{Var}[\\frac{1}{N} X_{kj}] \n= \\frac{(N-1)^2}{N^2} \\sigma_j^2 + (N-1) \\frac{1}{N^2} \\sigma_j^2 = \\frac{(N-1)}{N} \\sigma_j^2\n$$\n\nSo $Y_{ij}$ have smaller mean and variance compared to $X_{ij}$. Following the property of Gaussian distribution, we have:\n$$\nP(|Y_{ij}| < \\epsilon) > P(|X_{ij}| > \\epsilon),\\ \\forall \\epsilon > 0\n$$\n\nSo the distribution of $X_{ij}$ is more concentrated towards 0 after smoothing. Then we know that\n\n$$\nP(\\mathrm{absmax}(Y_i) < \\epsilon) = \\prod_{j=1}^d P(|Y_{ij}| < \\epsilon) > \\prod_{j=1}^d P(|X_{ij}| > \\epsilon) = P(\\mathrm{absmax}(X_i) < \\epsilon)\n$$\n\nSo this makes the distribution of absolute max value in a token more concentrated towards 0, leading to smaller quantization error.\n\n---\n>**Comment1.** The experimental settings of Figure 5 are not clear. What is number of heads? Which dataset is used?\n\n**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.\n\n---\n>**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?\n\n**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.\n\n[1] QLoRA: Efficient Finetuning of Quantized LLMs\n\n---\nIf you feel your concerns have been resolved, we would greatly appreciate it if you consider raising the score.", "claims_and_evidence": "All claims are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense.", "theoretical_claims": "There are no theoretical claims in this paper.\n\nThis paper does some mathematical derivations about quantization in section 3.1 and I have checked its correctness. I also have checked the implementation of SageAttention2 in algorithm 1.", "experimental_designs_or_analyses": "I have checked all experiments represented in main text. I think the experimental designs are sound and sufficient. \n\n1. It compares with FlashAttention and recent proposed INT4 quantizations under various tasks and settings.\n2. The ablation study is sufficient to demonstrate the effectiveness of its key designs, including Q/K smoothing strategies, Q/K quantization granularities, and numerical precision choices of P/V.", "relation_to_broader_scientific_literature": "algorithm and hardware co-design", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths: The proposed takes hardware characteristics into account. The proposed INT4 Per-thread Quantization is novel and it is a great example of algorithm/hardware co-design.\n\nWeaknesses: 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.", "other_comments_or_suggestions": "The experimental settings of Figure 5 are not clear. What is number of heads? Which dataset is used?", "questions_for_authors": "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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rMGHuH6SaI", "reviewer": "Reviewer_r8J3", "summary": "The authors propose SageAttention2, where they manage to quantize key and query matrices in the attention computation to INT4 while the softmax outputs and value matrices are quantized to FP8. They show that the quality degradation is managable in this configuration. Meanwhile, if the key and query matrices are quantized to INT8 instead, then there is almost no measurable quality degradation. The key is to perform asymmetric quantization of the key and query matrices.\n\n## update after rebuttal\nMy score is unchanged after the discussion period. Keeping my score to champion for the authors' work.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer r8J3,  \nThank you for your valuable question. Below, we address the question raised.\n\n---\n>**Question1.** Do the authors think that it is possible for PV computation to be quantized to sub-8-precision as well?\n\n**Reply**: \nThank 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.\n\nSpecifically:\n- We keep Q and K in INT8 (since Blackwell GPUs don't support INT4)\n- We quantize P and V to NVFP4\n- Using our smooth Q and smooth K techniques\n\n|Attention|Cossim ↑|Relative L1 ↓|\n|-|-|-|\n|Sage2-8b|0.99982|0.01573|\n|Sage2-4b|0.99460|0.06480|\n|FP4 PV Attention|0.99674|0.03250|\n\nBased on these preliminary results, we believe FP4 quantization for P and V is possible and practical.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "There aren't any theoretical claims except for simple mathematical derivations. They seem correct.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "Yes. Attention is a widely used operation. This work will be impactful.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "* The paper is well written.\n* The experimental design is comprehensive.\n* This work will be the new state-of-the-art attention kernel.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "* Do the authors think that it is possible for PV computation to be quantized to sub-8-precision as well?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FtOc5SUS9R", "reviewer": "Reviewer_aDo7", "summary": "In Transformer-based models, SageAttention accelerates self-attention through quantization, but its use of INT8 for queries and keys is slower than INT4. Moreover, its acceleration is limited to specific Nvidia architectures due to FP16 computations. To address this issue, this paper proposes a thread-level granularity quantization method that considers GPU architecture and applies outlier smoothing to queries. Furthermore, the study identifies architectural issues in FP16 matrix multiplication accumulation operations and introduces two-level accumulation strategies to mitigate these problems. As a result, SageAttention2 achieves superior performance and acceleration compared to various existing attention mechanisms, including FlashAttention, across diverse models and tasks.\n\n##  update after rebuttal\nThank you for the authors' response. After careful review, I believe the experiments and responses are sufficient. Therefore, I have revised the score for this paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer aDo7,  \nThank you for your valuable suggestions and questions.\n\n---\n>### Weakness1\n\n**Reply**: We appreciate the valuable question. We argue that:\n\n1. We first discover the critical role of accumulator precision in PV matrix multiplication for Attention; this is an essential insight for the attention operator.\n\n2. 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.\n\n3. Beyond addressing hardware limitations, The technique can enhance accuracy for efficiency techniques using reduced accumulator precision in matrix multiplication.\n\n---\n>### Weakness2\n\n**Reply**: Thank you for your question. We argue that:\n\n1.  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.  \n\n2. 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.\n\n3. Per-thread quantization is a highly innovative approach, and is non-trivial, requiring sophisticated algorithm-hardware co-design.\n\n---\n>### Weakness3\n\n**Reply**: \nThank 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.\n\n|Method|TOPS|Cossim↑|Relative L1↓|Ppl.↓|\n|-|-|-|-|-|\n|Attention (INT4 + FP8)|284|0.8004|0.3906|7.963|\n|+ Per-thread quantization|283|0.9249|0.3127|7.459|\n|+ Two-level accumulation|283|0.9498|0.2731|7.345|\n|+ Smooth Q|273|0.9946|0.06480|6.256|\n\n---\n>### Weakness4\n\n**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.\n\nFurthermore, 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:\n\n|Quantization|TOPS|\n|-|-|\n|per-tensor|286|\n|per-block (line 1094)|284|\n|per-thread (line 1095)|283|\n|per-token|268|\n\n---\n>### Comment1\n\n**Reply**: We appreciate the reviewer’s suggestion. We evaluate the end-to-end inference latency of Mochi on L20 below:\n\n|Attention|Latency (s)|\n|-|-|\n|HadamardAttn|1198|\n|SmoothAttn|1208|\n|SageAttentoin2|1190|\n\nSageAttention2 achieves a similar inference speed to HadamardAttn and SmoothAttn while delivering significantly better accuracy performance.\n\n---\n>### Comment2\n\n**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.\n\n---\n>### Question1\n\n**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:\n\n|Attention|TOPS|\n|-|-|\n|Torch|10.9|\n|xformers|94.1|\n|FlashAttn2|142.5|\n|SageAttn1|255.3|\n|SageAttn2-8b|329.9|\n|SageAttn2-4b|352.6|\n\nThe 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.\n\n---\n>### Question2\n\n**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:\n\n|Quantization|Cossim↑|Relative L1↓|\n|-|-|-|\n|Per-Channel|0.9946|0.0648|\n|Per-Token|✗|✗|\n|Per-Block|0.9930|0.0651|\n|Per-Tensor|0.9922|0.06777|\n\n---\nIf you feel your concerns have been resolved, we would greatly appreciate it if you consider raising the score.", "claims_and_evidence": "The proposed method effectively leverages low-bit operations for acceleration, demonstrating its suitability for efficient computation. Additionally, it visualizes and compares the data distribution of queries, keys, and values, providing a solid rationale for determining the granularity of each quantization. Through these analyses, this paper further demonstrates that the proposed model achieves high acceleration while maintaining reasonable performance, as evidenced by comparative examples of actual image generation.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are appropriate for the problem at hand. Experiments were conducted across multiple benchmarks, including text, video, and various other tasks.", "theoretical_claims": "The paper does not make significant theoretical claims that require proof verification. Instead, it focuses on empirical analysis and performance evaluation rather than theoretical derivations.", "experimental_designs_or_analyses": "The experiments are valid, as SageAttention2 has been compared across various models and tasks against multiple existing attention mechanisms.", "relation_to_broader_scientific_literature": "This paper builds on prior research in self-attention acceleration, particularly SageAttention, which introduced quantization for improved efficiency. However, its reliance on INT8 for queries and keys and FP16 for attention maps limited acceleration benefits to specific Nvidia architectures.\n\nBy introducing thread-level granularity quantization and extending outlier smoothing to queries, this study refines existing quantization techniques, aligning with prior work on mixed-precision inference and INT4 quantization. Additionally, the identification of FP16 matrix multiplication accumulation issues and the proposed two-level accumulation strategy contribute to ongoing research on numerical precision challenges in GPU-based deep learning.\n\nThe study’s experimental validation across various models, including LLMs, reinforces the effectiveness of quantization-aware strategies in optimizing Transformer inference across different GPU architectures.", "essential_references_not_discussed": "All essential references have been included.", "other_strengths_and_weaknesses": "1. Strenths\n\n• This paper identifies and analyzes the discrepancy between the theoretical accuracy of the proposed quantization method and its actual performance during GPU inference. By addressing this issue, the study demonstrates effective model acceleration while maintaining the intended accuracy. \n\n• This study analyzes the limitations of existing research and proposes improvements, conducting experiments on various models, including LLMs. The results demonstrate effective acceleration of model inference across a broad range of GPU environments.\n\n2. Weaknesses\n\n• This paper identifies errors arising during the accumulation process in matrix multiplication between the attention map and V. To mitigate this issue, it introduces a two-level accumulation technique. However, as acknowledged by the authors, this approach is not entirely novel. Additionally, the root cause of the problem stems from a design flaw in Nvidia GPU architecture. Since the proposed method primarily addresses a hardware-related limitation rather than introducing a fundamentally new algorithmic contribution, its impact on advancing the field is limited.\n\n• The proposed method enhances performance compared to SageAttention by reducing the quantization bit and mitigating the resulting performance degradation. This is achieved by extending outlier smoothing—previously applied only to the key—to the query and modifying the quantization granularity to a thread-wise level. However, the first approach is a straightforward extension of an existing technique rather than a novel innovation. Moreover, the change in granularity does not lead to a substantial improvement over previous methods, making the overall contribution appear insufficient.\n\n• A proper ablation study is needed. The ablation study provided in the appendix only compares TOPS, which focuses solely on computational speed. This is insufficient to fully explain the impact on actual model performance and makes it difficult to evaluate overall trade-offs.\n\n• While the paper explains why per-thread quantization was chosen, it lacks comparisons with other quantization methods.", "other_comments_or_suggestions": "• This paper compares model performance with existing techniques, including Smooth Attention and Hadamard Attention. However, it does not evaluate their inference speed, which would be a valuable addition.\n\n• The figures lack readability and explanation. For example, the authors mention converting P and V to FP8, but Figure 2 and other visual materials do not clearly illustrate this process.", "questions_for_authors": "• SageAttention2 performs well in computations with long sequences. However, can it still provide meaningful speed advantages in cases with short sequences or small batch sizes?\n\n• In the paper, value uses channel-wise granularity, unlike query and key. This appears to be due to value being computed with FP8 precision. Are there any results comparing the effects of using different quantization granularities specifically for value?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Employing low precision MMA on new hardware for key operations in neural networks can improve model inference performance. Previously, SageAttention accelerates the attention operation by quantizing and using INT8 MMA for QK MatMul, and performing FP16 MMA (FP16 accumulation) for the PV MatMul, which is quite suitable for the platforms that supports INT8 and FP16 MMAs, such as NVIDIA Ampere GPUs. Later, new platforms that supports more advanced low precision MMAs, such as INT4 and FP8, have been released. For instance, NVIDIA Ada GPUs (SM89) support INT4, INT8, and FP8 MMAs, and NVIDIA Hopper GPUs (SM90) support INT8 and FP8 MMAs. Alternative attention acceleration techniques, such as FlashAttention, have been extended to support these new low precision MMAs and have achieved better performance. Therefore, it is worthwhile to explore the possibility of using even lower precision MMAs, including INT4 and FP8, for SageAttention to further improve the performance of the attention operation. This paper proposed SageAttention2, which is a continuation of the previous SageAttention work for new hardware platforms.

In this paper, the attention QKV data distribution has been studied to explore the way of performing quantization that balances the performance and accuracy. Building on top of the K quantization smoothing in SageAttention, tha paper found it is vital to do quantization smoothing for Q as well for SageAttention2 that performs compute in low precision.

Due to the QK MatMul for SageAttention2 is computed using INT4 MMA, the per-block quantization used for the INT8 QK MatMul in SageAttention would result in a large quantization error, whereas the much finer per-token quantization would result in lower performance due to each thread in a warp has to handle multiple scale factors in the implementation. The paper analyzed PTX MMA instruction layout requirements and found that by using the mma.m16n8k64 warp instruction and grouping the tokens from Q and K in a certain way, which is finer than per-block but coarser than per-token, each thread in a warp only has to handle one scale factor. This is equivalent as improving the QK quantized MatMul accuracy without sacrificing performance and is critical for the success of SageAttention2.

When it comes to the PV MatMul for SageAttention2, the paper found that using FP8 MMA results in similar accuracy as FP16 MMA, while the performance is significantly improved. An additional key finding from the paper is that the paper found the FP8 MMA(f32f8f8f32) instruction on the Ada and Hopper architecture is actually FP22, which can cause accuracy decay for PV compute. To mitigate this issue, the paper implemented a two-stage accumulation strategy, that is using FP22 (the apparent data type is still FP32 though) for the atom FP8 MMA accumulation (it has to be like this because of the MMA instruction), and use FP32 for the final accumulation in a separate buffer using accumulation instructions that have no such FP22 issue.

SageAttention2, developed specifically for the Ada architecture, surpass FlashAttention2 and xformers, originally developed for the Ampere architecture, by about 3x and 4.5x measured on RTX4090. On the Hopper architecture, despite the fact that INT4 MMA is not supported, SageAttention2 was adapted to use INT8 MMA for QK MatMul. The performance of SageAttention2 on the Hopper architecture is on par with FlashAttention3, which is originally developed for Hopper architecture, while delivering much higher accuracy.

Most of the questions from the reviewers are related to or originated from the following two aspects:

1. The motivation of some design choices, such as data types, quantization granularity, and quantization smoothing necessity, etc., which requires experimental result support in the presence and absence of those choices.
2. The lack of introduction of supported precisions on different hardware platforms. Many reviewers were not aware that the Hopper architecture, even though newer than Ada, does not support INT4 MMA, and they had been asking why some implementations such as SageAttention2-8b is necessary and why some experiments were missing for SageAttention2-4b on Hopper architecture, etc. The authors should clarify this in the introduction and related work sections.

The authors of the paper have adequately addressed the comments and concerns from the reviewers and will add the missing experimental results and clarifications in the final version of the paper as requested.

Overall, this is a solid paper that presents the improvement of a previous accelerated computing algorithm for a new hardware platform. It will have a somewhat large impact to the deployment of Transformer models and it should be accepted by ICML 2025.
Decision: Accept (poster)