SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration
Authors: Jintao Zhang
Jia wei
Pengle Zhang
Jun Zhu
Jianfei Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Attention
Quantization
quantized attention
inference acceleration
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a quantization method for Attention that achieves speedups of 2.1x and 2.7x compared to FlashAttention2 and xformers, respectively, without lossing end-to-end metrics across various models.

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

Published as a conference paper at ICLR 2025
SAGEATTENTION: ACCURATE 8-BIT ATTENTION FOR PLUG-AND- PLAY INFERENCE ACCELERATION
Jintao Zhang, Jia Wei, Pengle Zhang, Jun Zhu, Jianfei Chen Dept. of Comp. Sci. & Tech., Institute for AI, BNRist Center, Bonghhhh-Bocchhh Bosch Joint MI. Center, THBI Lab, Tsinghua University (zhang-15 242 mails., jianfeice, deszje ts tsinghua. edu. cn
Tsinghua-
ABSTRACT
The transformer architecture predominates across various models. As the heart of the transformer, attention has computational complexity of O(N2), compared to O(N) for linear transformations. When handling large sequence lengths, at- tention becomes the primary time-consuming component. Although quantization has proven to be an effective method for accelerating model inference, existing quantization methods primarily focus on optimizing the linear layer. response, we first analyze the feasibility of quantization in attention detailedly. Follow- ing that, we propose SageAttent ion, a highly efficient and accurate quanti- zation method for attention. The OPS (operations per second) of our approach outperforms FlashAttention2 and xformers by about 2.1x and 2.7x. respectively. SageAt tentt on also achieves superior accuracy performance tention3. Comprehensive experiments confirm that our approach incurs almost no end-to-end metrics loss across diverse models--icllddne those for large lan- guage processing, image generation, and video generation. The code is available at https:///itthtt com/ www/tou-m/////eecce ttttttuummmmmmmm mmmammllmmteeeeeeeeemeeeeeeeeeee tention
FlashAt-
FlashAttention2
FlashAttention2
SageAttention
331
2x
164
Speed (TOPS)
SageAttention
FlashAttention2
SageAttention
90s
1.35x
End-to-end Speedup
Figure 1: An example of SageAttention on video generation (CogvideoX on RTX4090).
INTRODUCTION
Attention is the fundamental component of transformers (Vaswani 2017). and efficiently computing attention is crucial for transformer-based applications. Moreover, there is recent trend in process- ing longer sequences, which further strengthens the need for faster attention. In tasks like video generation (Yang et al. 2024) and language model prefilling (Dubey et al.) [2024), the sequence length can easily go up to 8K~ 128K. Due to quadratic complexity, the cost of attention domi- nates all other operations in such scenarios, illustrated in Figure
Quantization is an effective strategy for enhancing neural networks' computational and memory effi ciency by reducing the numerical precision. There abundant works accelerating training (un al. 2019 Xi al. 2024b, Peng et al. 2023) and inference (Jacob et al. 2018 Xiao et al. 2023a)
"Corresponding author.

---

Published as a conference paper at ICLR 2025
Latency of operators in transformer model 333 HeadDim=64)
800
Attention
FP16
Quantize K, to IMT8 directly
Latency (ms)
600
Linear
Others
400
Use
200
FlashAttemiond
1K
2K
4K
8K
16K
32K
Use
Sequence Length
SageAttention
Figure 2: Latency of attention.
Figure 3: A comparison example.
with low- low-precision numerical formats such as FP8, INT8, or INT4, However, existing works primar- ily focused quantizing the linear layer, where attention is left unaccelerated in high-precision, such as FP16. There is not yet a work that systematically investigates the quantization of atten- tion. Moreover, many quantization methods require extra training, and the cost can be prohibitive for large-scale models. While FlaenAttention3 (Shah et al.) 2024) was released recently and offers an FP8 version, it is tailored and can only be used with the Nvidia Hopper architecture. This exclusive optimization limits its broader applicability. Furthermore, our analysis demonstrates that directly implementing the FP8 version can lead to performance degradation, as detailed in Table
Quantizing attention is challenging. The computation of attention is more complex than that of linear operations. Attention includes a softmax operation and two matrix multiplication (Matmul) oper- ations: QK and PV. Direct 8-bit quantization and dequantization of the matrices (Q, K, P, V) in attention will result significantly degraded performance across various models. For example, the text-to-image model Unidiffuser (Bao et al. 2023) will generate completely blurry image with both INT8 and FlashAttention''ss s FP8 implementation (See Figure 3. and Llama2 only achieves random-guesse guessing--seed accuracy of 25.5% on the MMLU dataset with INT8 attention. After in- vestigating deeply, we identified two primary challenges: (C1) The matrix K exhibits a significant channel-wise outlier, leading substantial accuracy loss during quantization. (C2) Simply quantiz- ing (P, V) into INT8 does not consistently ensure the accuracy of PV across various scenarios.
quantization method to accelerate attention while preserving accuracy. SageAttent ion is easy-to-use. As a post-training quantization method, can be used in a alug-and-play manner in inference time by simply replacing the original high- precision implementation. We propose several techniques to achieve this goal. First, we opt to quantize the tensors attention to INT8 rather than FP8. This decision based on the fact that INT8 Matmul on some commonly used GPUs, e.g.. RTX4090 and 3090, are four times faster than in FP16 and two times faster than FP8. Moreover, INT8 quantization for matrices (Q, K) more precise than FP8 in attention (See Table To address (C1), we propose method to smooth the K matrix. This method significantly enhances accuracy with a negligible time overhead (<0.2%). To address (C2), an alternative quantizing (P, V) to 8-bit, we propose more accurate yet efficient method for the Matmul PV: we maintain (P, V) in FP16 and use a low- pwweiiiiiionn FP16 accumulator. This strategy doubles Matmul's speed without sacrificing any accuracy. Finally, we implement several versions of attention with different speed-accuracy tradeoffs and propose a method to select the fastest attention implementation for each layer while preserving accuracy.
In this paper, we propose SageAttention
We offer high-performance implementation of SageAttent in on RTX4090 and 3090 GPUs using Triton (Tillet al. 2019). Our implementation contains fused kernel combining ROPE with quantization and a fast self-attention kernel inspired by Alaetinntentton-sssyl tiling. The implemen- tation utilizes the fast INT8 mmma(u8.u8.s32) s32) and 1116-iit-FP11--cccmuuu instructions of Nvidia Tensor Core. Our kernel is about 2.1x and 2.7x faster than FlashAtten- tion2 and xformers, respectively. Notably, achieves 340 TOPS on RTX4090 headdim=64 and headdim=128, reaching 52% of the theoretical INT8 throughput. In contrast, the peak for the state- of-the-art FlashAttention2 is only 165 TOPS. Moreover, at headdim=64, our throughput on RTX 4090 is even close to the 490 TOPS throughput of Atteniitt, which is exclusive to the much more powerful and expensive Hopper GPUs We extensively evaluate the end-to-ent metrics of our approach
mmaff16,f16.f16)
oftte--f-thh--ee image/video generation, image classification, and language models. On all tasks, SageAtt ent ion can be directly adopted in a plug-and and and-play manner with negligible loss in model performance, while offering more than 2x speedup than Attention and xformers.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
We categorize efficient Attention works into three groups: (1) Sparse Attention, This strategy only selects parts of sequence from a given context for processing with standard Attention. Implemen- tations like Swin transformer (Liu et al. 2021), Twins (Chu et al. 2021), UniFormer (Li et al.), At- tentionsinks (Xiao et al. 2) InfLLM (Xiao et al., (2224, Long (Chen et al. 2023), Minfer- ence (Jian et al. 2024). ) and SkipAttention /Venkataramanan et al. 2023) show promise. However, that they only work in a few scenarios because omitted calculations are not always useless. (2) Linear Attention. Techniques that transform Attention computation to reduce time complexity, for example, Linformer (Wang et al. 2020)) Performer (Choromanski al.)
these methods' limitations
2020), MetaFormer (Yu et al. 2022), and LinearAttention (Katharopollos et al. 2020). which lower the time complexity of Attention from O(N2) into ((V). These methods excel in specific scenar- ios while standard Attention remains prevalent. (3) Kernel Optimization. Rather than simplifying calculations, these methods exploit hardware capacities to enhance speed. The xformers (Lefaudeux et al. (022) platform accelerates Attention with customizable blocks and dedicated CUDA kernels. FlashAttention (Dao et al. 2022) proposes tiling to reduce the memory reads/writes between GPU global memory and on-chip SRAM for significant speedups. FlashAttentioo22 (Dao 2023) refine the parallelism and warps partition of FlashAttention. Bikshandi & Shah (2023) further optimize FlashAttention2 by kernel fusion. FlashAttention3 (Shah al.) (222)) is proposed for Hopper archi- tecture. However, FlashAttention3 is exclusive to the Hopper GPU architecture, and the accuracy of its quantization version is significantly lower than our method (See Table T.. RingAttention (Liul eet a scales FlashAttention across multiple GPUs. I-bert (Kim et al.) (022)) quantizes all tensors a transformer block into INTS but is restricted RoBERTa. Our method falls under the third category, and is orthotopic with the first and second categories.
3 PRELIMINARY
Our method builds on FlashAttention-2 and adopts dynamic quantization. We will begin by review- ing FlashAttention-2, followed by a brief introduction to dynamic quantization techniques.
3.1 FLASHATTENTION
The computation of self-attention can be formulated as follows: S (QKT/VV P=o(S), O PV, where (5); = cxp(S.))/ Eg exp(Sut) is the softmax operation. The matrices Q. K, and V each have dimensions N d, while the matrices S, are N N. While is typically small, e.g., 64 or 128, N can be thousands if not millions. Therefore, the N x N matrices (S. P) are much larger than (Q, K, V), and a naive implementation suffers from the huge amount of global memory 1/O for (S, P) reads/writes. Attenttt (Dao (023) proposes to tile Q, K, and V from the token dimension into blocks {Q.}. ( [Vi with block sizes of by. Dkv, bkg. respectively. Then, to avoid the memory I/O for (S, P). it uses online softmax (Milakov & Gimelshein 2018 to progressively
compute each block of O, i.e., O as follows.
First, for each block of (K)), (V)), it computes the following equations iteratively:
sp Q-KK /d, (m,, P.) ã (m 1/ exp(mm_ 1 + rowsum (P.), Of = diag (exp(m? i- m²) 0f-1 + PVV, (2)
Where m² and & are by x 1 vectors, which are initialized to - and 0 respectively. ā() is an online (.roomxx())) softmax operator: m² max(m:-1 "rommax 8/)}. Pj = exp(S! m²)
Finally, after all iterations, i.e., bkc. the output O can computed by O = (
3.2 DYNAMIC QUANTIZATION
A matrix multiplication C = AB can be accelerated with quantization as: (5A. A) = (A1,, (6n, B) vB)) Ć = AB, C C= (c)
33)
3

---

Published as a conference paper at ICLR 2025
Here, 1 is quantizer which converts a high-precision (e.g., FP32) matrix A to low-precision format Å (e.g., INT8 FP8) with scale Đλ, and w is dequantizer to convert back to high- precision. We should have (A) A. The actual matrix multiplication AB is carried in low- precision. In modern GPUs, low-precision matrix multiplication is usually multiple times faster than higher-precision ones.
Many quantizers depend on the numerical format and granularity, e.g., how many elements share a common scale factor. For example, an INT8 per-tensor dynamic quantizer first computes the scale as the maximum absolute value of the entire tensor, scales the elements to the maximum repre- sentable range of INT8 [-127, +127], and then casts to INT8 with rounding: Á = [A/6A]] 1. = max (1 12212. Likewise, per-token quantizer assigns scale factor for each token of tensor: :] = : [i = max([AAi.:]]))
Also, per-channel quantizer assigns a scale factor for each channel of the tensor, i.e., along the channel dimension: Af:, i] [A[: 1,
Aa (]) 127. Based on the tiling approach of FlashAttention, we can apply per-block quan- tization correspondingly. per-block quantizer asigns scale factor for every b m n tokens:
Amm: n, [A[m n. n 0 max(| A m n, n2::ll//22 Dequantization simply involves
element-wist scaling: ,,, (A) = A.
4 SAGE ATTENTION
In this section, we propose Sageattention, a fast yet accurate method to accelerate attention computation with 8-bit quantization. Considering that most networks not natively trained with quantized attention, SageAttention is designed to be plug-and-pla
Unlike linear layers, which are easy to quantize, quantizing attention is more complicated. Extra treatment is required to ensure both good accuracy and fast speed. First, we will formulate quantized attention in Section 4.1 followed by introducing our approach.
4.1 FORMULATION
Based on the description of FlashAttention and dynamic quantization in Section 3.1 and 3.2 formulate the quantized attention as follows.
we
vo((//vd), (fk, R) = (K((K (6r, vrr (ōv, V) vVv (4) Quantization: (5Q, () =
Attention: S= Wb Vaqek (QRT), (m', ()m, O diag (exp(m' m)) Wapdy (PV) (5)
ΦK is a transformation to obtain quantized K, which we shall discuss in subsequent sections. For simplicity, omit all superscripts and subscripts, but the matrices used in attention still tiles, and the computation is still organized as FlashAttention described in Section 3.1 Compared to adds quantizers
the original full-precision version, shown in Eq. 4 5 SageAttent Q, K, P, V and dequantizers to the product to accelerate both Matmuls of QK" and PV. Online softmax is left in full-precise precision..
Talen
Talenn
Tokenn
Tolen
2006
Channel
Channel
Channel
Channel
Channel
Channel
14
12
-210
208 -151
347
(a) distribution Cogvideo
(b) distribution in Unidiffuser
Figure 4: Typical examples of data distribution of (Q. K. V).
4

---

Published as a conference paper at ICLR 2025
Table I: End-to-end metrics comparison of different quantization methods.
Quantization
Smoothing
Llama WikiText
Cog (Fscore)
Unidiffuser
UltraPixel
TIMM
(Q, K)
K
(FID)
(FID)
ImageNet
Full-Precision
5.823
3.768
163.33
179.78
84.79%
X
5.824
1.924
221.18
193.36
84.21%
Per-token
5.824
3.734
166.52
179.79
84.74%
X
5.825
2.014
229.08
195.67
84.18%
Per-hlock
5.824
3.718
166.93
179.98
84.76%
X
5.826
1.902
267.06
196.26
84.12%
Per-tensor
5.824
3.640
167.65
180.21
84.69%
FlashAttn3 (with quant)
5.850
3.394
394.13
383.61
84.70%
4.2
SMOOTH MATRIX K
Directly quantizing Q, K often results in large error. Particularly, quantizing Q, to INT8 yields completely blurry image/video in text-to-image/vider tasks. As shown in Figure 4.1 we visualize two typical groups of Q, K, V from a text-to-image model text-to-video model CogvideoX (Yang et al. (024). Notably, K exhibits distinct channel-wised outliers. However, per-channel quantization cannot be applied for K, because quantization can only be performed at the outer axis (token dim) of the Matmul QK Moreover, the previous smoothing technique proposed for linear layers (Xiao et al.) 2023a) cannot be applied since Q is also heavily affected by outliers. Fortunately, the channel outliers of K have pattern: Fach token's key is actually large bias shared by all tokens, plus small token-wise signal. Therefore, the outlier from large variation across tokens, but simply the large bias. Based this observation, propose to smooth the matrix K by transform which subtracts averaged K across all tokens:
(Bao et al., 2023) and
7(K K mean(K)
(6)
where mcan(K) = 1 EN K[t,;] is the average key, with a shape x d. Note that such N transformation does not change the attention score P. because for any query q, have a (q(K mean(K) (araan(K])) [) =qg
a(gh J). Finally, the transformation from full-precision K to quantized K can be written as Φx (K) 4/x o 7, where dhx is a quantizer. In other words, a full-precision K substracted with the mean, before eventually being quantized.
mean(K)
Table presents end-to-end metrics for different quantization methods with and without smoothing on various models. The results demonstrate that smoothing offers significant benefits aceu- racy. Moreover, the speed overhead of smoothing K for attention is less than 0.2 (See To)
Table 2: Average accuracy using different data types across all layers of real models.
Q.K
P. V.
Cos Sim
Relative
RMSE
Table 3: Worst accuracy using different data
B4M3
99.94%
0.0345
3.53e-3
types across all layers of real models.
INT8
E5M2
99.81%
0.0572
6.1
Q, P.
Cos Sim
Relative L1
RMSE
INT8
99.70%
0.1035
6.82e-3
E4M3
76.36%
0.5899
0.4311
B4M3
99.81%
0.0607
5.93c-3
E5M2
78.98%
0.4233
0.4371
E4M3
E5M2
99.68%
0.0769
7.72e-3
INT8
INT8
56.40%
0.7921
0.5405
INT8
99.58%
0.1199
8.31e-3
FP16
99.99%
0.0116
0.0091
E4M3
99.37%
0.1107
1.09c-2
E5M2
ESM2
99.22%
0.1213
1.20e-2
INT8
99.13%
0.1583
1.24e-2
4.3
QUANTIZATION FOR Q, V.
Quantization granularity for Q, K: vQ)(Q) and (k((K) be set with the granularity of per- token, block or per-tensor. This is because per-channel quantization not feasible, since the scale factors of the inner axis of QK? cannot be used to do dequantization (Xiao et al. 2023a
5

---

Published as a conference paper at ICLR 2025
Data type of Q, K: We choose INT8 for (Q( and dyx (K) for two reasons. First, Table 2 shows the average accuracy using different data types (INT8, E4M3, E5M2) for Q, K, P, V across all layers of Llama2 (7B) (Touvron et al. 2023) and Unidi fusser. It shows that quantizing Q, K to INTS performs higher accuracy than using E4M3 and E5M2. Second, Matmul using INT8 is two times faster than using FP8 in many commonly used GPUs, e.g., RTX4090 and 3090.
Quantization granularity for P, V: We propose use thy in per-block and (V) in per- channel for three reasons. (1) Per-channel quantization for P and per-token quantization for V are not viable because dequantization requires scale factors of the outer axis. (2) = exp(S: rowmax(S:)), where S; is the Matmul result of block of Q and KT, the max value in each row of P 1. Hence, we can assign a single static scale s boy to a block P. whose accuracy equals per-token quantization. (3) Per-channel quantization can address the channel-wised outlier of V.
Data type of P, V: We choose INT8 for 1/r(P) P) and shv (V) because Matmul using INT8 is two times faster than using FP8 in some commonly used GPUs, and although the accuracy using &p (P) and (V) in INT8 is worse than E4M3 and E5M2, the average accuracy is similar (See Table 2
Accuracy metrics. We use three metries assess the accuracy of quantized attention output O' compared to attention output in full- full-precision O: First, we flatten O' and O into vectors in the shape of 1 n. Then, Cosine Sim= 000 VR Relative LI= 0'1001 [ RMSE- (O - 0')
Table 4: Average accuracy using different
Table 5: Worst accuracy using different accumu-
accumulators across all layers of real models.
lators across all layers of real models.
Accum.
Cos Sim
Relative L1 J
RMSE | 2.94e-3
Accum
Cos Sim f
Relative L1
RMSE
FP32
99.98%
0.0156
FP32
99.84%
0.0511
4.229e-3
FP16
99.98%
0.0156
2.94c-3
FP16
99.84%
0.05
4.229e-3
4.4
FP16 ACCUMULATOR: MUCH MORE ACCURATE AND EFFICIENT SOLUTION
The above solution for shr(P (P) and (V) has one problem, that is, the accuracy using INT8 is very poor in some model layers. Table 3 shows the worst accuracy using different data types for Q,K.P, V across all layers of Llama2 and Unidiffused
It shows that INT8 pP()) and (V) bring an unacceptable error. In response, we propose a very accurate and also efficient solu- tion. Specifically, we propose to use FP16 as the data type of Matmul PV with an FP16 accumulator.
The benefit of such a solution is obvious. First, in the context of some commonly used GPUs, e.g., RTX4090 and 3090, the speed of Matmul in FP16 with an FP16 accumulator is 2x faster than that with an FP32 accumulator. Moreover, using FP16 accumulators can save more register resources than using FP32 accumulators, accelerating the computation speed. Second, Table 3 shows that using FP16 for P. V. is much more accurate than using all the other 8-bit data types. Moreover, using FP16 accumulators incurs no accuracy loss than using FP32 accumulators. Specifically, Ta- ble and Ssshow the average and worst accuracy using FP16 or FP32 accumulators all layers of Llama2 and Unidi fidsfr.urrr showing that there is no accuracy loss of using the FP16 accumulator.
Table 6: Four kernel implementations of SageAttention.
Kernel SAGEAttn-T
(Q), th (x<(K) (K)
(p(P)
ty (V)
per-token, INT8
FP16, FP16 Accumulator
FP16, FP16 Accumulator
SAGEAttn-B (Algorithm SAGEAttn-VT (Figure 5 aj)
plo-block, INT8
FP16, FP16 Accumulator
FP16. FP16 Accumulator
per-token, INT8
per-block, INT8
per-channel, INT8 per-channel, INT8
SAGEAttn-V3
per-block, INT8
per-block, INT8
4.5 ADAPTIVE QUANTIZATION
Based on the discussion in Section and 4.4
we implement four attention kernels (See Ta-
ble (6) based on two sets of choices: (1) Using (Q) and (K) in per-token per-block.

---

Published as a conference paper at ICLR 2025
SAGEAttn-VT
Loop for K. V block of Q
SAGEAttn-B
Loop for K, too get block of
Global Memory
K KINTS
VINTS
Global Memory
K INT+
VpP16
On Chip
Sk
FP32
On Chip
Sk
FP32
QINTS
9q
SM1
S
QINTS 8q
SM1
S
x
SM1
SM2
P OnlineSof (mma(()
SM2
P SnlineSoftmazil
FP32
p = (P (27))8
p P.o(FFP6)
0
O
Accum. FP16)
(a) SugeAttention (per-token quantize Q.K: INT8 V)
(b) SageAttention (per-block quantize Q.K: FPI6 V)
Figure 5: Workflow of SageAttention.
Algorithm 1: Implementation of SAGEAttn-B.
Input: Matrices (FPP16), (FP16), V(FP16) E Rxx block size by, bkv-
Preprocessing: K = K (M)an(K):/1 Subtracting the mean value across tokens
Quantization: (8Q. Q) (Q((V/d)) (Ōk, R)
INT8 per-block quant Divide Q into Tm N/h, blocks ( and divide K, V into T, N/hav blacks (R,) and (V]
for [1, T... do: // Outer ooo is paralle
SMs (stream processors)
Load Q, and [t] into SM
for j in [1, Tn. do
Load Kj. Vj. and 6x []] into the SM;
S Matmml(Q,,, KJ) &g [i] x 8x]
mf = max(ma row.maxxx)))). () Pi- = exp(S - mm)). 12 = =d -m! rowsum (
O diag(c"! ! -m? -- Matmul(P) 10(PP1)) (o(F11)),
VG, Accum.tty FFI))
O !!?"
Write O:
return 0 (Oi):
(2) Using (P((P) and vy (V) in INT8 retaining P, V in FP16 with FP16 accumulator. SAGEAttn-B is accurate enough for all models and can achieve a 2x speedup (See Figure 6 and 7 However, SAGEAttn-V3 is also accurate for some layers in model and faster lit- tle (about 4%) than SAGEAttn-B. Therefore, we use various inputs to test the cosine similarity of sAGEAttn-VS for each layer of model. Then, we will select AnGEAttn-ve for those layers where SAGEAItn-VEl cosine similarity is bigger than 99.8% (the worst similarity of SAGEAttn-B), and the other layers are left for SAGEAttn-3.
4.6 FUSION TRICKS AND PERFORMANCE ANALYSIS
Fusion Tricks. To reduce the overhead of quantization, we fuse the quantization process with the operator preceding the attention layer. For instance, we fuse quantization within the ROPE (Rotary Position Embedding) (Su et al. 202 layer. Specifically, before the ROPE result (A) is written from shared memory into global memory, perform 6 À (A)). Subsequently. the 8 Å are written into global memory. Additionally, we also fuse the coefficient (1/vd) of QK" into the quantization process rather than leaving it in the attention layer. Specifically, we multiply Q by (1/vd) on chip before quantizating Q.
Performance Analysis. We will take SAGEAttn-R as an example to discuss the acceleration fects on actual hardware: (1) Matmul acceleration. Utilizing INT8 matrix multiplication units on current mainstream hardware can achieve 2-4x throughput. While FP16 accumulators do not of- fer throughput improvements on most compute cards, on-edge accelerators, such as the RTX4090, can still achieve 2x improvement over FP32 accumulators. (2) Quantization overhead. Quanti- zation and dequantization are considered the main overhead in current quantization methods (Lin) cet al. 2024). The computational overhead can not be avoided, but through fusing the quantization of Q, K with ROPE, we avoid the IO overhead of quantization. (3) Cache and registers. Currently,

---

Published as a conference paper at ICLR 2025
mainstream accelerators need store data in cache (such SharedMemory) during computation. Using 8-bit data for calculations can reduce the usage of the general cache, and using fp16 accu- mulators can also reduce the usage of accumulation registers. (4) Dram access. Using 8-bit data can halve the tensors transfer overhead from DRAM to the compute units. Although quantization introduces additional FP32 scales, these scales be considered negligible compared to the tensors.
5 EXPERIMENTS
Main results. The speed of SageAttention is approximately 2.1 x faster than FlashAttention Furthermore, SageAt ention achieves an average real speedup of 2.83x compared to the original attention in various models, with negligible loss in end-to-end metrics.
5.1 EXPERIMENTAL SETUP
Models. We validate the effectiveness of SageAttent ion across a diverse set of representative models from the fields of language, image, and video generation. Specifically, conduct experi- ments on five models: Lmama2 (7B) (Touvron et al. 2023) for text2text, CogvideoX (Yang et al. 2024 for text2video, Unidi fuusr (Bao al. 2023) and UltraP aatralllel (Ren al. (024) for
text2image, TIMM (Wightman 2019) for image classification, and Llaval. 6
(Liu et al. 2024a)
for visual question answering.
2 is evaluated on three zero-shot tasks: WikiText (Merity et al. 2022) to as- sess the model's s prediction confidence, LAMBADA (Paperno et al. 2016) evaluate contextual un- derstanding, and MMLU (Hendryeks et al.) 2020) for measuring knowledge across various sub- jects. CogvideoX evaluated using the open-sora (Zheng et al.) 2024c) prompt sets. Both UltraPixee and Unidi ffuser are assessed on the COCO annotations (Lin et al. (014), fea- turing (prompt, image) pairs. TIMM is evaluated on on three image datasets: ImageNet (Deng et al. 2009), Imagehet-Sketch (Sketch) (Wang et al., 2019), and r) (Hendrycks et al., 202T). Llaval. 6 evaluated on three datasets: TextVQA (Singh et al. 2019). POPE eL. 2023b), and VQAv2 (Goyal et al. 2017).
Datasets.
(ImageNet-
Metrics. For Llama2, we use perplexity (ppl.) (Jelinek et al. 1977) for WikiText, and Accuracy (Acc.) for LAMBADA and MMLU. For CogvideoX, folowing (Zhao al.) 2024a), we evaluate the quality of generated videos on five metrics: CLIPSIM and CLIP-Temp (CLIP-T) (Liu et al., 2024b) to measure the text-video alignment; (VQA-a) and (VQA-t) assess the video aesthetic and technical quality, respectively; and Flow-score (FScore) for temporal consistency (Wu et al.) 2023). For UltraPixel and mnidiffused generated images are compared with the images in the COCO annotations dataset in three aspects: FID (Heusel et al. 2017) and sFID) (Salimans) al. 2016) for fidelity evaluation, Clipscore (CLIP) (Hessel al.) (202)) 1) for text-image alignment, and ImwgeR (IR) [Xu et al.) 2024) for human preference. For TIMM and Llaval. 6, we use Accuracy.
RTX4090,
600
FlashAtti2
SageAttn Sagellttn
Sagelttin
500
Sageletn VT Sagelett
SapeAttn
Speed (TOPS)
xtfrmees
SagsArtmn
xftrmers
Sagelttr T
400
200
1K
2K
4K
8K
16K
32K
1K
2K
8K Sequence Longth
4K
16k
32K
Sequence Length
Figure 6: Speed comparison between SageAttention and baselines (RTX4090, headdim=64).
5.2 SPEED AND ACCURACY OF ATTENTION KERNELS
Speed. We conduct experiments to compare the Speed of SageAitent ion against baselines us- ing configurations with headdim=64 headdim=128, both with and without Causal Mask Vaswani (017). Specifically. Figure 6 and Figure 7 show the Speed of SageAttention and basclines across varying sequence lengths RTX4090. These results indicate that SageAttention
8

---

Published as a conference paper at ICLR 2025
KTX4090, (Heed dim 128, Causal Falsel
RTX4090, (Head dim 128. Causal - True)
600
Torch sformees
FlashWind SagsAtttn
Saguhlioul Sag ellttti88
SaguAitn-8
600
Tuech xtbrmers
FassiALin Sag ekett>
Sugudiinwr
SageALinvdi
Speed (TOPS)
Stge/ltt
400
400
200
200
1K
2K
4K
16K
32K
1K
4K
8K
16k
32K
Length
Sequence Length
Figure 7: Speed comparison between SageAr and baselines (RTX4090, headdim=
RTX3090, (Head dim 64, causal False)
RTX3090, IHead dim 64, causal True)
Tonch xformerk
FlashAzt22
Sag TT
Sagekttn-88
Torch xformers
FashMAt222 Sug aigeAttn
Sagelctn -T
SageAttn-d +B
Speed (TOPS)
Sagwhitn-T
aggeltinn
Segelete-B
100
100
2K
4K
16K
Length
Sequence Length
Figure 8: Speed comparison between SageAttention and baselines (RTX3090, headdim=64).
RTX3090, (Head dim 128, Causal
False)
RTX3090, (Head dim 128, Causal True)
Torch
Flash-Ultnd
Sageltiiw
SagwAllr-oB
Trech
FlastMiin2 clittn
Segwdettoe?
SageAlin-di
Speed (TOPS)
200
SagsAttn
Ssgeleth
100
100
8K
Sequence Length
Sequence Length
Figure 9: Speed comparison between SageAttention and baselines (RTX3090, headdim=128).
achieves peak of 341 TOPS and is 2x faster than FlashAttention2 and 2.9x faster than xform- ers on average. Figure 8 and Figure illustrate the results on RTX3090, showing similar speedup performance.
Accuracy. Table shows the numerical error four implementations of SageAttent ion com- pared with attention in full-precision. This experiment is conducted using a set of (Q, K, V) con- forming to a normal distribution. It shows the error of the four implementations is rather small. SAGEAt and SACEAttn-B achieve 100% cosine similarity and RMSE in the e-4 level.
Table 7: Real speedup of SageAttention RTX4090.
Model
Shape of Q, K,
Original attention
SageAttention
Speedup
CogvideoX
(2, 30, 17776, 64)
163.37 (FlashAtt22)
327.57
2.01x
Llama2
(4, 32, 1536, 128)
130.99 (FlashAttn2)
231.74
1.77x
UltraPixel
(2, 32, 7285, 64)
152.03 (FlashAttn2)
325.18
2.14x
Unidiffuser
(4, 24, 1105, 64)
105.68 (xformers)
246.93
2.34x
TIMM
(12, 64, 197, 64)
18.910 (Torch)
111.41
5.89x
5.3 END-TO-END PERFORMANCE
Speedup. We measure the real speed of SageAttention and the original attention on Unidi ffuser, fltraPixel, Cogvideos a.lama2 and TTMM on RTX4090. Table 7 shows that SageAttention outperforms original attention across all models. Specifically, SageAttention yields 2.83x speedup compared to the original attentions on average.
Metrics loss. We assessed the end-to-end metrics of various models using SageAttents on com- pared to using attention in full-precision. Detailed evaluation results are presented in Table g for Llama2, CogvideoX, Unidiffusez, UltraPixell and TIMM, respectively. The results in- dicate that SageAttent across all models. Specifically, on I. amlmm2, CogvideoX, U1 UltraPixel. and Unidi fuser,
successfully matches the performance of attention in full-precision

---

Published as a conference paper at ICLR 2025
Table 8: End-to-end metrics loss across text, image, and video generation models.
Model
Attention
WikiText (Ppl.) +
Lambda (Acc.) f MMLU (Acc.) f
Full-Precision
5.823
0.886
0.46
L.lama2
SageAttention
5.824
0.887
0.46
Model
Attention
CLIPSIM
CLIP-T 1
VQA-a
VQA-t f
FScore
Full-Precision
0.1837
0.9976
68.962
75,925 75.037
3.7684
CogvideoX
SageAttention
0.1836
0.9976
68.839
3.8339
Model
Attention
FID J
sFID
CLIP
IR f
Full-Precision
163.33
145.08
0.3152
0.1609
Unidi fusser
SageAttention
166.49
143.18
0.3154
0.1521
Full-Precision
179.78
141.35
0.3132
0.6169
CltraPixel
SageAttention
179.79
141.63
0.3131
0.6110
Model
Attention
ImageNet (Acc.) f Sketch (Acc.)
ImageNet-r (Acc.) f
Full-Precision
84.79%
45.32%
59.55%
TIMM
SageAttention
84.74%
45.78 %
60.32%
Model
Attention
TextVQA (Acc.) f
POPE (Acc.) 1
VQAv2 (Acc.)
Full-Precision
60.25%
86.45%
77.55%
Llaval.6 6
SageAttention
60.09 %
86.44 %
77.47%
SageAttention resulted in only minor average degradation of 0.2% compared to attention in full full-precision. Moreover, on TIMM, SageAt tent even surpasses attention in full cull-precision.
Table 9: Accuracy of SageAttention kernels.
Table 10: Overhead of smoothing K.
attention
Cos Sim
Relative L1 J RMSE L
Model
Smooth
TOPS 1
SAGEALLa-T 1.0
0.019
6.8e-4
X
327.57
CogvideoX
SAGEAtEn-B
1.0
0.021
7.3c-4
v
327.52
SAGEALLn-vT
99.9%
0,064
0.065
X
325.18
UltraPixes
SAGEAttn-vB
98.9%
0.138
0.067
v
324.56
Table 11: Benefit of adaptive quantization.
attention SAGEAttn-T
model
CLIPSIM t TOPS f
Model
MMLU t
TOPS 1
0.1827
292.17
0.46
208.59
Cogvideox
Llama2
SageAttention
0.1835
327.57
0.46
231.74
5.4 ABLATION STUDY
Overhead of smoothing K. Table I0 presents the overhead associated with smoothing K on the attention speed in real models. The results indicate minimal reduction, less than 0.2%.
Benefit of adaptive quantization. We analyzed the performance differences between using only sAGEAttn-T and employing an adaptive strategy (tageAttention) Table II presents the metries and average speed of attention on CogvideoX and Llama2. The results indicate that the adaptive strategy increases the speed of attention by 11.7% without any loss in metrics.
6 CONCLUSION AND FUTURE WORK
We introduce SageAttention on, an efficient and precise INT8 quantization method for attention. First, we propose a method to smooth matrix K, enhancing the accuracy with under 0.2% speed over- head. Second, we use FP16 accumulators in the Matmul of (P, V) to boost both accuracy and speed. Third, we use adaptive quantization further improve OPS by 12% without sacrificing accuracy. Our method surpasses FlashAttention2 and xformers by approximately 2.1x and 2.7x. respectively. Extensive testing confirms that our approach maintains end-to-end metrics across various models, including language, image, and video generation models.
10

---

Published as a conference paper at ICLR 2025
ACKNKOWLEGEMENT
The authors would like to thank Haofeng Huang for his valuable help with the implementation. This work was supported by the NSFC Project (No. 62376131), Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University. J.Z is also supported by the XPlorer Prize.
REFERENCES
Fan Bao, Shen Nie, Kaiwen Xue. Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: vit backbone for diffusion models. In CVPR, 2023.
Ganesh Bikshandi and Jay Shah. A case study in cuda kernel fusion: Implementing flashattention-2 on nvidia hopper architecture using the cutlass library. arXiv preprint arXiv:2312.11918, 2023.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309 12307, 2023.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Xiangxiang Chu, Zhi Tian, Yuging Wang, Bo Zhang, Haibing Ren, Xiaolin Wei, Huaxia Xia, and Chunhua Shen. Twins: Revisiting the design of spatial attention in vision transformers. Advances in neural information processing systems, 04:93355-9666 2021.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.0869/ 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 15:16344-16359) 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255, Icee, 2009.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
Tianyu Fu, Haofeng Huang, Xuefei Ning, Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, Guohao Dai, Huazhong Yang, and Yu Wang. Moa: Mixture of sparse attention for automatic large language model compression, 2024a.
Tianyu Fu, Tengxuan Liu, Qinghao Han, Guohao Dai, Shengen Yan, Huazhong Yang, Xuefei Ning, and Yu Wang. Framefusion: Combining similarity and importance for video token reduction on abs large visual language models, 2024b. URL httpps:///axvvv.o/ wttp:////a./...o../a//////// 2501. 01986
Yash Goyal, Tejas Khot, Douglas Summers-Star Dhruv Batra, and Devi Parikh. Making the v in vqa matter: Elevating the role of image understanding in visual question answering. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 9904-6913, 2017.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. 2020.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu. Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. The many faces of robustness: A critical analysis of out-of out-of-distribution generalization. ICCV, 2021.
11

---

Published as a conference paper at ICLR 2025
Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. 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.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Yuezhou Hu, Weiyu Huang, Zichen Liang, Chang Chen, Jintao Zhang, Jun Zhu, and Jianfei Chen. Identifying sensitive weights via post-quantization integral. 2025.
Weiyu Huang. Yuezhou Hu, Guohao Jian, Jun Zhu, and Jianfei Chen. Pruning large language mod- els with semi-structural adaptive sparse training, 2024. URL https:///axxi..... abs/ 2407 20584
Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-arithmeti--nnly inference. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2704-2713. 2018.
Fred Jelinek, Robert Mercer, Lalit R Bahl, and James K Baker. Perplexity-a measure of the difficulty of speech recognition tasks. The Journal of the Acoustical Society of America, 62(S1):
$63-S63, 1977.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for looglooonnnn context lims via dynamic sparse attention. arXiv preprint arXiv:2407 024900 2024.
Youhe Jiang, Ran Yan, Xiaozhe Yao. Yang Zhou, Beidi Chen, and Binhang Yuan. Hexgen: Gener- ative inference of large language model over heterogeneous environment. In Forty-first Interna- tional Conference on Machine Learning.
Youhe Jiang, Fangcheng Fu, Xiaozhe Yao, Guoliang He, Xupeng Miao, Ana Klimovic, Bin Cui, Binhang Yuan, and Eiko Yoneki. Demystifyngg cost-efficiency in Ilm serving over heterogeneous gpus. arXiv preprint vrXiv:2502.00722 2025a.
Youhe Jiang, Ran Yan, and Binhang Yuan. He Disaggregated generative inference of Ilms in heterogeneous environment. In International Conference on Learning Representations (ICLR), 2025b.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference ma- chine learning, pp. 5156-5165. PMLR, 2020.
Sehoon Kim, Amir Gholami, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. I-bert: Integer- only bert quantization. In International conference machine learning, pp. 5506-5518, PMLR, 2021.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xformers: A modular and hackable trans- former modelling library. 2022.
Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems, 36: 15:11516-1177, 2023a.
Bingrui Li, Wei Huang, Andi Han, Zhanpeng Zhou, Taiji Suzuki, Jun Zhu, and Jianfei Chen. On the optimization and generalization two-layer transformers with sign gradient descent. In Interna- fional Conference on Learning Representations (ICLR), 2025.
Kunchang Li, Yali Wang, Gao Peng, Guanglu Song. Yu Liu, Hongsheng Li, and Yu Qiao. Uniformer: Unified transformer for efficient spatial-temports representation learning. In International Con- ference Learning Representations.
12

---

Published as a conference paper at ICLR 2025
Yifan Li, Yifan Du, Kun Zhou, Jinpeng Wang, Wayne Xin Zhao, and Ji-Rong Wen. Evaluating object hallucination in large vision-language models. arXiv preprint hat::2305.10353 2023b.
Tsung-V Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740-755. Springer, 2014.
Yujun Lin, Haotian Tang. Shang Yang, Zhekai Zhang, Guangxuan Xiao, Chuang Gan, and Song
Han. Qserve: W4a8kv4 quantization and system co-design for efficient Ilm serving, 2024. URL
https:// /tapi////xvvv...//. org/ abs/ 2405. 04532
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transformers for near- infinite context. In The Twelfth International Conference on Learning Representations.
Haotian Liu, Chunyuan Li, Yuheng Li, Bo Li, Yuanhan Zhang. Sheng Shen, and Yong Jae Lee. Llava-next: Improved reasoning, ocr, and world knowledge, January 2024a. URI. (tRihttps::/
Slava-v..gituul wttp://.oo///gggg/g//
30- -lava-next)
Yaofang Liu, Xiaodong Cun, Xuebo Liu, Xintao Wang, Yong Zhang, Haoxin Chen, Yang Liu, Tieyong Zeng, Raymond Chan, and Ying Shan. Evalcrafter: Benchmarking and evaluating large video generation models. In Proceedings the IEEE/CVF Conference Computer Vision and 2024b.
Pattern Recognition, pp. 22139-2211
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEF/CVF international conference on computer vision, pp- 10012-10022, 2021.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations. 2022.
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint 1.05.v/805.020000 2018.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. In Proceedings the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). pp. 1524-1534, 2016.
Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang. Ze Liu, Yifan Xiong, Ziyue Yang, Bolin Ni, Jingcheng Hu, et al. Fp8-Im: Training fp8 large language models. arXiv preprint arXiv:2310.1833 2023.
PyTorch Contributors. Torch backend documentation. stable/baccendeee
http:://pytorch.orr//aot
www.ttt cuda
Jingjing Ren, Wenbo Li, Haoyu Chen, Renjing Pei, Bin Shao, Yong Guo, Long Peng, Fenglong Song, and Lei Zhu. Ultrapixel: Advancing ultra-bigh-resolluton image synthesis to new peaks. arXiv preprint Xiix:2407.00258 2024.
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08608, 2024.
Amanprect Singh, Vivck Natarajan, Meet Shah, Yu Jiang, Xinlei Chen, Dhruv Batra, Devi Parikh, and Mareus Rohrbach. Towards vqa models that read. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8317-8326, 2019.
Jianlin Su. Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. arXiv preprint 104.:2004009900 2021.
13

---

Published as a conference paper at ICLR 2025
Xiao Sun, Jungwook Choi, Chia-Yo Yu Chen, Naigang Wang, Swagath Venkataramani, Vijayalak- shmi Viji Srinivasan, Xiaodong Cui, Wei Zhang, and Kailash Gopalakrishnan. Hybrid 8-bit float- ing point (hfp8) training and inference for deep neural networks. Advances in neural information processing systems, 32, 2019.
Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. MAPL 2019, pp. 10-19, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450367196.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykow, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.092286 2023.
Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Shashanka Venkataramanan, Amir Ghodrati, Yuki M Asano, Fatih Porikli, and Amirhossein Habib- ian. Skip-attention: Improving vision transformers by paying less attention. arXiv preprint arXiv:2301.02240. 2023.
Haohan Wang, Songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global representa- tions by penalizing local predictive power. Advances in Neural Information Processing Systems, 32, 2019.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 04777 2020.
Ziteng Wang, Jianfei Chen, and Jun Zhu. Remoe: Fully differentiable mixture-of-experts with relu routing. In International Conference on Learning Representations (ICLR), 2025.
Ross Wightman. Pytorch image models.
pytorch- ieale---aae---- image-models 2019.
Haoning Wu, Erli Zhang. Liang Liao, Chaofeng Chen, Jingwen Hou, Annan Wang, Wenxiu Sun, Qiong Yan, and Weisi Lin. Exploring video quality assessment on user generated contents from aesthetic and technical perspectives. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 00014-201555 2023.
Haocheng Xi, Han Cai, Ligeng Zhu, Yao Lu, Kurt Keutzer, Jianfei Chen, and Song Han. Coat: Compressing optimizer states and activation for memory-efficient fp8 training. arXiv preprint arXiv:2410.19311 2024a.
Haocheng Xi, Yuxiang Chen, Kang Zhao, KAI JUN TEH, Jianfei Chen, and Jun Zhu. Jetfire: Efficient and accurate transformer pretraining with int8 data flow and per-block quantization. In Forty-first International Conference Machine Learning, 2024b.
Haocheng Xi, Shuo Yang. Yilong Zhao, Chenfeng Xu, Muyang Li, Xiuyu Li, Yujun Lin, Han Cai, Jintao Zhang, Dacheng Li, et al. Sparse videogen: Accelerating video diffusion transformers with spatial-temporal sparsity. arXiv preprint arXiv:2502.00777, 2025.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. Infilm: Training-free long-context extrapolation for ilms with efficient context memory. In First Workshop on Long- Context Foundation Models@ ICML 2024, 2024.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38099, PMLR, 2023a.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint 200. 17453, 2023b.
Jiazheng Xu, Xiao Liu, Yuchen WU, Yuxuan Tong, Qinkai Li. Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. Advances in Neural Information Processing Systems, 36, 2024.
14

---

Published as a conference paper at ICLR 2025
Zhuoyi Yang. Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang, Wenyi Hong, Xiaohan Zhang, Guanyu Feng, et al. Cogvideox: Text-to-video diffusion models with an expert transformer. arXiv preprint (rXx:2408.00677, 2024.
Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. Metaformer is actually what you need for vision. In Proceedings the conference on computer vision and pattern recognition, pp. 10819-10829, 2022.
the IEEE/CVF
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, and Jianfei Chen. Sageatten- tion2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization. arXiv preprint arXiv:2411. 100958, 2024.
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, and Jianfei Chen. Sageattention2:
Efficient attention with smoothing q and per-thread quantization. 2025a.
Jintao Zhang, Guoliang Li, and Jinyang Su. Sage: A framework of precise retrieval for rag. In 2025 IEEE 41th International Conference on Data Engineering (ICDE). IEEE, 2025b.
Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, and Jianfei Chen. Spargeattn: Accurate sparse attention accelerating any model inference. arXiv preprint vrXiv:2502.18137 2025c.
Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, and Jianfei Chen. Spargeattn: Training-free sparse attention accelerating any model inference. 2025d.
Pengle Zhang, Jia wei, Jintao Zhang, Jun Zhu, and Jianfei Chen. Accurate int8 training through dynamic block-level fallback. 2025e.
Tianchen Zhao, Tongcheng Fang, Enshu Liu, Rui Wan, Widyadewi Soedarmadji, Shiyao Li, Zinan Lin, Guohao Dai, Shengen Yan, Huazhong Yang, et al. Vidit-q: Efficient and accurate quantiza- tion of diffusion transformers for image and video generation. arXiv preprint 24000025400 2024a.
Tianchen Zhao, Xuefei Ning, Tongcheng Fang, Enshu Liu, Guyue Huang, Zinan Lin, Shengen Yan, Guohao Dai, and Yu Wang. Mixdq: Memory-efficient few-step text-to-image diffusion mod- els with metric-decoupled mixed precision quantization. In European Conference on Computer Vision, pp. 285-302. Springer, 2024b.
Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming- Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. arXiv preprint MrXi::2409.02908, 2024a.
Kaiwen Zheng, Guande He, Jianfei Chen, Fan Bao, and Jun Zhu. Diffusion bridge implicit models. arXiv preprint arXiv:2405. 5585,, 2024b.
Kaiwen Zheng, Guande He, Jianfci Chen, Fan Bao, and Jun Zhu. Elucidating the preconditioning in consistency distillation. arXiv preprint arXiv:2502.02222.. 2025.
Zangwei Zheng, Xiangyu Peng, Tianji Yang, Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all, March 2024c. URL httpst::: higthub.b... com/mpcaitechh/oen. Sora
15

---

Published as a conference paper at ICLR 2025
A EXPERIMENTAL DETAIL
A. ENVIRONMENT
We implemented our Attention kernels using OpenAI Triton (Tillet et al. 2019) and conducted experiments on Ubuntu 22.04 servers. Tests on the RTX 4090 utilized server with PCIE 5.0, 16-core Xeon(R) 6430 CPU, and 120GB DDR4 RAM, while the RTX3090 tests employed a server with 16-core Xeon(R) 8358P CPU and 80GB DDR4 RAM. To reproduce our results, experiments should conducted in the environment of torch 2.4. 2.4.0+culll 121, triton-nightly (version 20240816), python 3.11, and (gec, g++) in version 9.
A.2 HYPER-PARAMEN PARAMETERS FOR ATTENTION KERNELS
We use 128 for block size of Q, and 64 for a block size of K and V. The parameters Num_ Mar_S and Num Stt_St,ges, which represent the number of warp schedulers and the number of processing stages in our GPU kernels, respectively, are detailed in Table
Table 12: Hyper-parameters for our Attention Kernels.
HeadDim
Causal Mask
Num Warps
Num Stages
64
False
3
64
True
128
False
3
128
True
5
A.3 DETAILS OF DATASETS AND MODELS
We choose the first 256 annotations from the COCO 201 4val dataset as the prompt set for UltraPixed and Unid ridii didiifff ffuser image generation. We also used the corresponding 256 images of the 256 prompts as the ground truth images to calculate the FID and sFID. For Cogvi deoX, the model is trained long texts, 120 words. The specific model we used for TIMM vir base p 6 224.augreg2_.n222
applied open-sora prompt set, each consisting of more than
A.4 ADDITIONAL EXPERIMENTS
Table 13: Comparison of SageAttention with AWQ (W4A16) on Llama2.
SageAttention
AWQ AWQ+SageAtteettion
Perplexity Speedup of Linear Computation Speedup of Attention
5,4721
5.4729
5.5988
5,5998
0
0
2x
2x
Table 14: Comparison of Stgenttention with Q-diffusion (W8A8) on Unidiffuser.
FID|
sFID
CLIP
ImageReward1
Full Precision
163.33
145.08
31.52
0.1609
SageAttention
166.49
143.18
31.54
0.1521
O-diffusion (W8A8)
395.99
178.56
18.03
-2.273
A.5 COMPARISON WITH OTHER METHODS
There are some task- task-specific quantization methods, such as AWQ for LLMs, Q-diffusion for text- to-image, and ViDiT-Q for text-to-vides applications. SageAttention is orthogonal to them because those works are mainly used to quantize the linear layers. Second, AWQ is only used to compress the parameters of L.LMs with acceleration effect in computation. Q-diffusion has not reported its
16

---

Published as a conference paper at ICLR 2025
Table 15: Comparison of SageAttention with VIDIT-Q on CogvideoX.
CLIPSIM CLIP-TY
VQA-at
VQA-t
FScore1
End-to-end Speedupt
Full Precision
0.1837
0.9976
68.962
75.925
3.7684
SageAttention
0,1836
0,9976
68.839
75.037
3,8339
34.3% %
VIDIT-Q (W8A8)
0. 1888
0.9974
68.185
71.011
3.7342
22% (theoretical maximum)
acceleration effect in their paper and provided codes with acceleration effect in its official repository. ViDiT-Q has not provided the codes with acceleration effect in its official repository. Nonetheless, we compare SageAttention with those works as follows. (1) We compare the perplexity of Llama2- 7B on WikiText and the speedup in the prefilling stage. The results are shown in Table 13 compare StgeAttention with Q-diffusion (W8A8) on Unidiffuser. The results are shown Table 14 We compare SageAttention with VIDIT-Q on CogvideoX and the results are shown in Table 15 Since the official repository does not provide acceleration code, we estimate a theoretical maximum: the Linear layer accounts for 24% of Cogvideo's lateney, and W8A8 offers at most 4x speedup for the Linear layer, resulting a theoretical maximum end-to-end speedup of 100
We
140-24x % 22%.
A.6 SOME INSIGIITS
Table shows that the metric of Llama2 remains stable with quantization. The reason that the distribution of Q, K, and V in the attention of Llama2-7B is relatively uniform. As result, quan- tizing Q, K, and V to INT8 or FP8 does not significantly impact the accuracy of attention. This insight inspires the idea that better control over outlier activations in models could lead to more precise quantization results. As works like Fu et al. (2024a); Zhang et al.] (2024 2025a), we believe SageAttent ion can also be effectively applied to various applications related to Transformers. such as MOE systems (Wang et al., 2025). linear layer quantization (Hu et al. 2025 Zhang et al. 2025e Zhao al.] 2024a)) RAG systems (Zhang al., 2025b), training optimization (Li et al. 2023a 2025: Huang et al. 2024 Xi et al. 2024a), heterogeneous GPU systems (Jiang et al. 2022a
Jiang et al. 2025b), and diffusion models (Zheng al. 2024a b 2025
Fu al. 2024b Zhao al.
2024b Xi et al. 2025 Zhang et al. 2025c
Table 16: SageAttention based on Torch Attention.
Sequence Length
Torch Attention
SageAttention based on Torch Attention
1024
46
48
2048
42
55
4096
55
87
8192
OOM
00M
B IMPLEMENTATION BASED ON TORCH ATTENTION
FlashAttention is the statc-of-the-art and most commonly used standard attention; another com- monly used attention is Torch attention (PyTorch Contributors). We report the implementation speeds based on Torch in Table 16
Table 17: Numerical error of Q K using different type of quantization.
Data Type
Cosine Sim
Relative LI
INT8 99.54%
0.084
B4MB
92.83%
0.342
E5M2
77.95%
0.681
B.1 ADDITIONAL. PRECISION COMPARISION
Table 17 shows the precision of Q K using per-token quantization in different data types com- pared to Q K in full precision. This experiment is conducted using Q, K from the 24th layer of
17

---

Published as a conference paper at ICLR 2025
Table 18: Error of quantized attention with without smoothed K.
Quantization Type
Smoothed
Cosine Sim t
Relative +
RMSE +
Without
62.24%
1.187
0.294
Per-token (SAGRAttn-T)
With
99.47 %
0.045
0.031
Without
30.60%
1.286
0,464
Per-block (SAGEAttn-B)
With
99.31 %
0.072
0.035
Without
41.40%
1.554
0.399
Per-tensor
With
98.06 %
0.126
0.059
FlashAttention-3 (quantized version)
26.76%
2.5354
0.5378
Full-Precision
(b) SAGE Attention
Figure 10: An image generation example of UltraPixel.
(a) Full-Precision
(b) SAGE Attention
Figure 11: A video generation example of Open-Sora.
Table 19: Comparison of real speedup on RTX3090.
Model
Shape Q, K, V
Original Attention
SAGEAttention
Speedup
CogvideoX
(2, 30, 17776, 64)
71.57 (FlashAttn2)
129.87
1.81x
Llama2
(4, 32, 1536, 128)
56.54 (FlashAutn2)
108.91
1.93x
UltraPixel
(2, 32, 7285, 64)
65.86 (FlashAttn2)
131.74
2.00x
ffuser
(4, 24, 1105, 64)
47.64 (xformers)
108.91
2.29x
TIMM
(12, 64, 197, 64)
12.33 (Torch)
66.34
5.38x
ffuserr.... shows that quantizing Q, K to INT8 performs higher precision than using E4M3 and E5M2.
18

---

Published as a conference paper at ICLR 2025
Prompt1
Prompt2
Prompt3
Prompt4
Full Precision
SageAttention
Figure 12: More image generation examples of UltraPixel, where promptl="Two dogs are looking up while they stand near the toilet in the bathroom", prompt2="A gray bicycle is locked some metal doors", prompt3="AM An image of car driving on the highway", and prompt4="A A cat on the lid of toilet looking perturbed"".
Prompt1
Prompt2
Prompt3
Prompt4
Full Precision
SageAttention
Figure 13: More image generation examples of Unidiffuser, where "Beauti= view of the Himalayas", prompt2="An An elephant under the sea", pronpt3="Fnglish Country Garden Design", and prompt4="An An old red electric rail train in Durango, Colorado".
Prompt1
Prompt2
Prompt3
Prompt4
Prompt5
Full Precision
SageAttention
Figure 14: More image generation examples of CogvideoX. For more details about the prompts and the full videos. refer to hattp:://anonynnn http:///ws...oooo.cccccce/////// wtwp:////eencceeeeeeeaaeeeeaaaea examples 3344433ADDAA mcl
video
Table 18 shows the precision of different quantization methods with and without smoothing K various models. The results demonstrate that smoothing K offers significant benefits of precision.
B.2 VISUALIZED RESULTS
Figure shows the high-resolution images (2560x1536 generated by UltraPixel using Atten- tion full precision and SageAt ion. It can be seen that SageAttent ion matches the full precision in the high quality and highly detailed images. Figure [| shows the videos (720x 1280) generated by Open-Sora (Zhngg et al. 2024c) in different precisions. SageAttention yields identically the same video the full precision one.
19

---

Published as a conference paper at ICLR 2025
Figure 12 Figure 13 and Figure 14 Unidiffuser, and CogvideoX.
show more visualized comparison results UltraPixel.
B.3 REAL SPEEDUP ON RTX3090
We further measure the real speed of SageAttention and the original Attention on Unidiffuser, UltraPixel, Cogvideox, Llama2 and TIMM on RTX3090. Table 7 shows that SageAttention outperforms original attention across all models. Specifically, SageAt on yields 2.7x speedup compared to the original Attentions on average.
20


Reviews:
[{"review_id": "n5icgrPNML", "reviewer": "Reviewer_UTAx", "summary": "This paper introduces a new attention quantization method for speeding up transformer inference. To accelerate attention, the authors propose using INT8 quantization instead of FP8 for faster matrix multiplication on GPUs, along with a method for smoothing the K matrix to improve accuracy. Instead of quantizing the P and V matrices, they maintain them in FP16 and use a low-precision FP16 accumulator for faster multiplication without accuracy loss.  Finally, they offer different speed-accuracy trade-offs and a layer-wise selection method for optimal performance. Extensive experiments have been done on different transformer models for text, image and video generation tasks. The results show that the proposed SageAttention speeds up the FlashAttention2 and xformers by more than 2 times without losing any performance.", "questions": "I am not an expert in this field and do not have additional questions at this stage. I might have more questions at the discussion phase.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The contributions of this work are well motivated. \n\n2. The proposed method seems to be novel although I am not an expert in this field. \n\n3. The experiments are quite extensive, covering two different GPUs (RTX4090 and RTX3090), representative models for language, image, and video generation, and a wide range of datasets. \n\n4. The results are quite impressive, showing more than two times speedup without performance degradation.", "weaknesses": "1. Some design choices seem to be decided by the specific hardwares that are evaluated RTX4090 and 3090 (L271). Are those design choices also compatible with other GPUs like A100 and H100? \n\n2. Table 7 shows that different model/task has different speedup. How is the speedup related to the specific transformer architecture, model size, and complexity of the task?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "bxwvBEGmF1", "reviewer": "Reviewer_tyTT", "summary": "This papers conducts in depth analysis of viability to quantize LLMs/Diffusion models into INT8 frameworks. It also proposes smoothing methods to alleviate the outlier pains in the QKV projection process, demonstrating viable tradeoffs. Comparisons to relevant work is strong.", "questions": "As above in weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Paper is well written.\n- Experiments are thorough.\n- Problem is challenging.", "weaknesses": "- Full comparison to strong SOTA methods such as Flash attention 3, though slightly mentioned in the introduction and in Table 14, is not deeply explored. \n- Only targeted 4090/3000 series GPUs - it would be recommended to be tested on stronger GPUs at server level that is facing the strongest limitations.\n- It would be great to test across VLMs too.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "vnVVtrliV4", "reviewer": "Reviewer_wwcu", "summary": "This paper presents SageAttention, an innovative quantization method designed to optimize the computational efficiency of attention mechanisms in transformer models by employing an 8-bit integer (INT8) quantization. It employs the FlashAttention-wise quantization and matrix smoothing with FP16 for Matmul. With all these proposed components,  the method demonstrates improved computational performance and maintains accuracy compared to existing solutions like FlashAttention2 and xformers.", "questions": "Could the authors elaborate on the \"Llama\" column of Table 1? Specifically, why does the number remain stable even with quantization? Understanding this aspect could provide valuable insights into the robustness of the proposed method.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ The paper is well-written, with a logical structure and organization that facilitates understanding.\n\n+ SageAttention shows competitive performance, outperforming FlashAttention2 and xformers by approximately 2.1x and 2.7x, respectively.\n\n+ The method exhibits almost no end-to-end metrics loss across a variety of models, including large language models (LLMs), text-to-image (T2I), and text-to-video (T2V).\n\n+ The discovery of channel-wise consistency, as illustrated in Figure 4, is particularly noteworthy and adds depth to the research.", "weaknesses": "- The method relies heavily on FlashAttention, which may weaken its technical contribution and originality. What will the performance be if it does not employ the FlashAttention as the basis?\n\n- The reported superiority over FlashAttention3 appears to be quite marginal, raising questions about the significance of the improvements.\n\n- Another major weakness of this paper is that it does not compare SageAttention with other task-specific quantization methods, such as AWQ [1] for LLMs, Q-diffusion [2] for text-to-image, and ViDiT-Q [3] for text-to-video applications, which could provide a more comprehensive evaluation of its performance.\n\n[1] AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration\n\n[2] Q-Diffusion: Quantizing Diffusion Models\n\n[3] ViDiT-Q: Efficient and Accurate Quantization of Diffusion Transformers for Image and Video Generation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "2UkwPUZJcT", "reviewer": "Reviewer_xY6z", "summary": "This work explores post-training quantization of the attention mechanism. It claims that it's not well explored in the literature and the existing works focus either on quantized training or post-training quantization of non-attention layers (of course, among other attention acceleration methods, like linear attention, algorithmic accelerations, etc.). The work observed that the Keys matrix has similar values per-token, and proposed to smooth it by applying a per-token bias. It explores optimal quantization strategies and combines it with flash attention algorithmic techniques to obtain the best performance. SageAttention achieves very high reconstruction quality and does not seem to lose much performance on downstream metrics.", "questions": "- Is there a final attention variant among the proposed ones which performs the best uniformly across all tasks?\n- L154: \"The expression diag(l_i^j)^{-1}\" would produce inf-s since l_i^j is initialized as 0. Or how is O_i^j initialized?\n- L151: What is \"online softmax\" (i dont see it to be defined explicitly)? How does it differ from the standard softmax?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The performance boosts seem to be impressive compared to \n- The work carries a valuable observation that the keys matrix has huge correlation per-token. It can be valuable even outside of \n- I find the evaluation section to be very thorough, including many domains (text2image, text2video, LLMs). I especially appreciate text2video (where the sequences are the longest) and performance exploration on second-tier GPUs (e.g. 3090), which are more suitable for production use cases.\n- I am not an expert in quantization, but I checked some prior works and didn't find similar approaches. This makes me consider the paper as novel.\n- The paper provides the source code: I have not tried running it but it is valuable.", "weaknesses": "- I think that writing should be improved. A table I would like to see the most is the end-to-end model speed improvement and metrics degradation after integrating the proposed attention mechanism. First, for many models accelerating the attention op itself does not necessarily lead to much improved overall speed when the sequence length is short and it's MLPs who carry the main burden. Second, it's not entirely clear when looking at various tables which variant is being used at the end (and why there are so many of them throughout the paper instead of being just in some restricted ablation section). Finally, there are quite many typos and grammatical errors, e.g.:\n    - L084: \"matrice\" => \"matrices\"\n    - L138 (and elsewhere): \"IO\" => \"I/O\"\n    - L144: \"it compute\" => \"it computes\"\n    - L142=144: missing line or equation\n    - L151: \"The σ()\" => \"The\" should be ommited or it should be \"The function σ()\".\n    - L184: \"First, We\" => \"First, we\"\n    - L264-L266: ~3 typos in 3 lines\n    - etc.\n- There are no qualitatives provided for image/video generation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)