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

## Abstract

Abstract The transformer architecture predominates across various models. As the heart of the transformer, attention has a computational complexity of O ⁢ ( N 2 ) 𝑂 superscript 𝑁 2 O(N^{2}) italic_O ( italic_N start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ) , compared to O ⁢ ( N ) 𝑂 𝑁 O(N) italic_O ( italic_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 codes are available at https://github.com/thu-ml/SageAttention .

## 1 Introduction

Attention is the fundamental component of transformers (Vaswani, ), and efficiently computing attention is crucial for transformer-based applications.
Moreover, there is a recent trend in processing longer sequences, which further strengthens the need for faster attention. In tasks like video generation (Yang et al., ) and language model prefilling (Dubey et al., ), the sequence length can easily go up to 8K$\sim$ 128K. Due to its quadratic complexity, the cost of attention dominates all other operations in such scenarios, as illustrated in Figure 3.

Figure: Figure 2: Latency of attention.
Refer to caption: x2.png

Quantization is an effective strategy for enhancing neural networks’ computational and memory efficiency by reducing the numerical precision. There are abundant works on accelerating training (Sun et al., ; Xi et al., ; Peng et al., ) and inference (Jacob et al., ; Xiao et al., ) with low-precision numerical formats such as FP8, INT8, or INT4. However, existing works primarily focused on 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 attention.
Moreover, many quantization methods require extra training, and the cost can be prohibitive for large-scale models.
While FlashAttention3 (Shah et al., ) was released recently and offers an FP8 version tailored for 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 1.

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) operations: $QK$ and $PV$. Direct 8-bit quantization and dequantization of the matrices $(Q,K,P,V)$ in attention will result in significantly degraded performance across various models. For example, the text-to-image model Unidiffuser (Bao et al., ) will generate a completely blurry image with both INT8 and FlashAttention3’s FP8 implementation (See Figure 3), and Llama2 only achieves a random-guessing-level accuracy of 25.5% on the MMLU dataset with INT8 attention. After investigating deeply, we identified two primary challenges:
(C1) The matrix $K$ exhibits a significant channel-wise outlier, leading to substantial accuracy loss during quantization.
(C2) Simply quantizing $(P,V)$ into INT8 does not consistently ensure the accuracy of $PV$ across various scenarios.

In this paper, we propose SageAttention, a quantization method to accelerate attention while *preserving accuracy*. SageAttention is easy-to-use. As a post-training quantization method, it can be *used in a plug-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 in attention to INT8 rather than FP8. This decision is 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 matrice $(Q,K)$ is more precise than FP8 in attention (See Table 3).
To address (C1), we propose a method to smooth the $K$ matrix. This method significantly enhances accuracy with a negligible time overhead ($<$0.2%).
To address (C2), alternative to quantizing $(P,V)$ to 8-bit, we propose a more accurate yet efficient method for the Matmul $PV$: we maintain $(P,V)$ in FP16 and use a low-precision 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.

We offer a high-performance implementation of SageAttention on RTX4090 and 3090 GPUs using Triton (Tillet et al., ). Our implementation contains a fused kernel combining ROPE with quantization and a fast self-attention kernel inspired by FlashAttention-style tiling. The implementation utilizes the fast INT8 *mma(u8.u8.s32)* and FP16-with-FP16-accumulator *mma(f16.f16.f16)* instructions of Nvidia Tensor Core.
Our kernel is about $2.1\times$ and $2.7\times$ faster than FlashAttention2 and xformers, respectively.
Notably, it achieves 340 TOPS on RTX4090 at 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 FlashAttention3, which is exclusive to the much more powerful and expensive Hopper GPUs.
We extensively evaluate the end-to-end metrics of our approach on state-of-the-art image/video generation, image classification, and language models. On all tasks, SageAttention can be directly adopted in a plug-and-play manner with negligible loss in model performance, while offering more than 2$\times$ speedup than FlashAttention2 and xformers.

## 2 Related Work

We categorize efficient Attention works into three groups:
(1) Sparse Attention. This strategy only selects parts of a sequence from a given context for processing with standard Attention. Implementations like Swin transformer (Liu et al., ), Twins (Chu et al., ), UniFormer (), Attentionsinks (Xiao et al., ), InfLLM (Xiao et al., ), LongLora (Chen et al., ), Minference (Jiang et al., ), and SkipAttention (Venkataramanan et al., ) show promise. However, these methods’ limitations are 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., ), Performer (Choromanski et al., ), MetaFormer (Yu et al., ), and LinearAttention (Katharopoulos et al., ), which lower the time complexity of Attention from $O(N^{2})$ into $O(N)$. These methods excel in specific scenarios 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., ) platform accelerates Attention with customizable blocks and dedicated CUDA kernels. FlashAttention (Dao et al., ) proposes tiling to reduce the memory reads/writes between GPU global memory and on-chip SRAM for significant speedups. FlashAttention2 (Dao, ) refine the parallelism and warps partition of FlashAttention. Bikshandi & Shah () further optimize FlashAttention2 by kernel fusion. FlashAttention3 (Shah et al., ) is proposed for Hopper architecture. However, FlashAttention3 is exclusive to the Hopper GPU architecture, and the accuracy of its quantization version is significantly lower than our method (See Table 1). RingAttention () scales FlashAttention across multiple GPUs. I-bert (Kim et al., ) quantizes all tensors in a transformer block into INT8 but is restricted to 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 reviewing 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=QK^{\top}/\sqrt{d},~{}P=\sigma(S),~{}O=PV$, where $\sigma(S)_{ij}=\exp(S_{ij})/\sum_{k}\exp(S_{ik})$ is the softmax operation.
The matrices $Q$, $K$, and $V$ each have dimensions $N\times d$, while the matrix $S$, $P$ are $N\times N$. While $d$ is typically small, e.g., 64 or 128, $N$ can be thousands if not millions. Therefore, the $N\times N$ matrices $(S,P)$ are much larger than $(Q,K,V)$, and a naive implementation suffers from the huge amount of global memory IO for $(S,P)$ reads/writes. FlashAttention (Dao, ) proposes to tile $Q$, $K$, and $V$ from the token dimension into blocks $\{Q_{i}\},\{K_{i}\},\{V_{i}\}$ with block size of $b_{q}$, $b_{kv}$, $b_{kv}$, respectively. Then, to avoid the memory IO for $(S,P)$, it use online softmax (Milakov & Gimelshein, ) to progressively compute each block of $O$, i.e., $O_{i}$:

First, for each block of $\{K_{i}\},\{V_{i}\}$, it compute the following equations iteratively:

$$ $\displaystyle S^{j}_{i}=Q_{i}K_{j}^{\top}/\sqrt{d},$ $\displaystyle~{}~{}~{}(m^{j}_{i},\widetilde{P}^{j}_{i})=\tilde{\sigma}(m^{j-1} _{i},S^{j}_{i}),$ (1) $\displaystyle l_{i}^{j}=\exp(m^{j}_{i}-m^{j-1}_{i})l_{i}^{j-1}+\mathrm{rowsum} (\widetilde{P}^{j}_{i}),$ $\displaystyle~{}~{}O^{j}_{i}=\mathrm{diag}\left(\exp(m^{j}_{i}-m^{j-1}_{i}) \right)O^{j-1}_{i}+\widetilde{P}^{j}_{i}V_{j}$ (2) $$

Where $m_{i}^{j}$ and $l_{i}^{j}$ are $b_{q}\times 1$ vectors, which are initialized to $-\infty$ and 0 0 respectively. The $\tilde{\sigma}()$ is an online softmax operator: $m^{j}_{i}=\max\{m^{j-1}_{i},\mathrm{rowmax}(S^{j}_{i})\},~{}\widetilde{P}_{j}^
{i}=\exp(S^{j}_{i}-m^{j}_{i})$.

Finally, the output $O_{i}$ can be computed by $O_{i}=\mathrm{diag}(l_{i}^{j})^{-1}O_{i}^{j}$.

### 3.2 Dynamic Quantization

A matrix multiplication $C=AB$ can be accelerated with quantization as:

$$ $\displaystyle(\delta_{A},\hat{A})=\psi(A),~{}~{}~{}~{}(\delta_{B},\hat{B})= \psi(B),~{}~{}~{}~{}\hat{C}=\hat{A}\hat{B},~{}~{}~{}~{}C=\psi^{-1}_{\delta_{A} \delta_{B}}(\hat{C})$ (3) $$

Here, $\psi$ is a *quantizer* which converts a high-precision (e.g., FP32) matrix $A$ to a low-precision format $\hat{A}$ (e.g., INT8 or FP8) with a *scale* $\delta_{A}$, and $\psi^{-1}$ is a *dequantizer* to convert back to high-precision. We should have $\psi^{-1}_{\delta_{A}}(\hat{A})\approx A$. The actual matrix multiplication $\hat{A}\hat{B}$ 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 compute the scale as the maximum absolute value of the entire tensor, scales the elements to the maximum representable range of INT8 [-127, +127], and then cast to INT8 with rounding: $\hat{A}=\lceil A/\delta_{A}\rfloor,\delta_{A}=\mathrm{max}(|A|)/127$.
Likewise, *per-token quantizer* assigns a scale factor for each token of a tensor: $\hat{A}[i,:]=\lceil A[i,:]/\delta_{A}\rfloor,\delta_{A}[i,:]=\mathrm{max}(|A[i
,:]|)/127$.
Also, *per-channel quantizer* assigns a scale factor for each channel of the tensor, i.e., along the channel dimension: $A[:,i]=\lceil A[:,i]/\delta_{A}\rfloor,\delta_{A}=\max(|A[:,i]|)/{127}$.
Based on the tiling approach of FlashAttention, we can apply per-block quantization correspondingly. *per-block quantizer* asigns a scale factor for every $b=m-n$ tokens: $\hat{A}[m:n,:]=\lceil A[m:n,:]/\delta_{A}\rfloor,\delta_{A}=\mathrm{max}(|A[m:
n,:]|)/127$.
Dequantization simply involves a element-wise scaling: $\psi_{\delta_{A}}^{-1}(\hat{A})=\delta_{A}\hat{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 are not natively trained with quantized attention, SageAttention is designed to be plug-and-play.

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, we formulate the quantized attention as follows.

$$ $\displaystyle\textbf{Quantization: }{\color[rgb]{0,0,1}\definecolor[named]{ pgfstrokecolor}{rgb}{0,0,1}{(\delta_{Q},\hat{Q})}}={\color[rgb]{0,0,1} \definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\psi_{Q}}}(Q/\sqrt{d}),~{}~{}{ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{(\delta_{K}, \hat{K})}}={\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{ \phi_{K}}}(K),~{}~{}{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb }{0,0,1}{(\delta_{P},\hat{P})}}={\color[rgb]{0,0,1}\definecolor[named]{ pgfstrokecolor}{rgb}{0,0,1}{\psi_{P}}}(\widetilde{P}),~{}~{}{\color[rgb]{0,0,1 }\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{(\delta_{V},\hat{V})}}={ \color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\psi_{V}}}(V)$ (4) $\displaystyle\textbf{Attention: }~{}~{}S={\color[rgb]{0,0,1}\definecolor[named ]{pgfstrokecolor}{rgb}{0,0,1}{\psi^{-1}_{\delta_{Q}\delta_{K}}(\hat{Q}\hat{K}^ {\top})}},~{}~{}(m^{\prime},P)=\tilde{\sigma}(m,S),~{}~{}O=\mathrm{diag}\left( \exp(m^{\prime}-m)\right)O+{\color[rgb]{0,0,1}\definecolor[named]{ pgfstrokecolor}{rgb}{0,0,1}{\psi^{-1}_{\delta_{P}\delta_{V}}(\hat{P}\hat{V})}}$ (5) $$

$\phi_{K}$ is a transformation to obtain quantized $K$, which we shall discuss in subsequent sections.
For simplicity, we omit all superscripts and subscripts, but the matrices used in attention are still tiles, and the computation is still organized as FlashAttention described in Section 3.1. Compared to the original full-precision version, as shown in Eq. 4, 5, SageAttention adds quantizers to $Q,K,P,V$ and dequantizers to the product to accelerate both Matmuls of $QK$ and $PV$. Online softmax is left in full-precision.

Figure: Figure 4: Typical examples of data distribution of (Q, K, V).
Refer to caption: x4.png

**Table 1: End-to-end metrics comparison of different quantization methods.**
| Quantization<br>(Q, K) | Smoothing<br>K | Llama<br>WikiText<br>$\downarrow$ | CogVideo<br>(Fscore)<br>$\uparrow$ | Unidiffuser<br>(FID)<br>$\downarrow$ | UltraPixel<br>(FID)<br>$\downarrow$ | TIMM<br>ImageNet<br>$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| Full-Precision | - | 5.823 | 3.768 | 163.33 | 179.78 | 84.79% |
| Per-token | ✗ | 5.824 | 1.924 | 221.18 | 193.36 | 84.21% |
| ✓ | 5.824 | 3.734 | 166.52 | 179.79 | 84.74% |  |
| Per-block | ✗ | 5.825 | 2.014 | 229.08 | 195.67 | 84.18% |
| ✓ | 5.824 | 3.718 | 166.93 | 179.98 | 84.76% |  |
| Per-tensor | ✗ | 5.826 | 1.902 | 267.06 | 196.26 | 84.12% |
| ✓ | 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 a large error. Particularly, quantizing $Q,K$ to INT8 yields completely blurry image/video in text-to-image/video tasks.
As shown in Figure 4, we visualize two typical groups of $Q,K,V$ from a text-to-image model Unidiffuser (Bao et al., ) and a text-to-video model CogvideoX (Yang et al., ). Notably, $K$ exhibits distinct channel-wised outlier. 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^{\top}$. Moreover, the previous smoothing technique proposed for linear layers  (Xiao et al., ) cannot be applied since $Q$ is also heavily affected by outliers. Fortunately, the channel outliers of $K$ have a pattern: Each token’s key is actually a *large bias shared by all tokens*, plus a small token-wise signal. Therefore, the outlier is not from large variation across tokens, but simply the large bias.
Based on this observation, we propose to smooth the matrix $K$ by a transform $\gamma$, which subtracts averaged $K$ across all tokens:

$$ $\displaystyle\gamma(K)=K-\mathrm{mean}(K)$ (6) $$

where $\mathrm{mean}(K)=\frac{1}{N}\sum_{t=1}^{N}K[t,:]$ is the average key, with a shape $1\times d$.
Note that such a transformation does not change the attention score $P$, because for any query $q$, we have $\sigma(q(K-\mathrm{mean}(K)^{\top}))=\sigma(qK^{\top}-q\cdot\mathrm{mean}(K))=
\sigma(qK^{\top})$. Finally, the transformation from full-precision $K$ to quantized $\hat{K}$ can be written as $\phi_{K}(K)=\psi_{K}\circ\gamma$, where $\psi_{K}$ is a quantizer. In other words, a full-precision $K$ is substracted with the mean, before eventually being quantized.

Table 1 presents end-to-end metrics for different quantization methods with and without smoothing K on various models. The results demonstrate that smoothing K offers significant benefits of accuracy. Moreover, the speed overhead of smoothing $K$ for attention is less than 0.2% (See Table10).

**Table 2: Average accuracy using different data types across all layers of real models.**
| $Q,K$ | $\widetilde{P},V$ | Cos Sim $\uparrow$ | Relative L1 $\downarrow$ | RMSE $\downarrow$ |
| --- | --- | --- | --- | --- |
|  | E4M3 | 99.94% | 0.0345 | 3.53e-3 |
| INT8 | E5M2 | 99.81% | 0.0572 | 6.11e-3 |
|  | INT8 | 99.70% | 0.1035 | 6.82e-3 |
| E4M3 | E4M3 | 99.81% | 0.0607 | 5.93e-3 |
| E5M2 | 99.68% | 0.0769 | 7.72e-3 |  |
| INT8 | 99.58% | 0.1199 | 8.31e-3 |  |
| E5M2 | E4M3 | 99.37% | 0.1107 | 1.09e-2 |
| E5M2 | 99.22% | 0.1213 | 1.20e-2 |  |
| INT8 | 99.13% | 0.1583 | 1.24e-2 |  |

### 4.3 Quantization for Q, K, P, V

Quantization granularity for $Q,K$: $\psi_{Q}(Q)$ and $\psi_{K}(K)$ can in granularity of per-token, per-block or per-tensor. This is because per-channel quantization is not feasible, since the scale factors of the inner axis of $QK^{T}$ cannot used to do dequantization (Xiao et al., ).

Data type of $Q,K$: We choose INT8 for $\psi_{Q}(Q)$ and $\psi_{K}(K)$ for two reasons. First, Table 3 shows the average accuracy using different data types (INT8, E4M3, E5M2) for $Q,K,\widetilde{P},V$ across all layers of Llama2 (7B) (Touvron et al., ) and Unidiffuser. It shows that quantizing $Q,K$ to INT8 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 $\widetilde{P},V$: We propose to use $\psi_{P}(\widetilde{P})$ in per-block and $\psi_{V}(V)$ in per-channel for three reasons. First, per-channel quantization for $\widetilde{P}$ and per-token quantization for $V$ are not viable because dequantization requires scale factors of outer axis. Second, $\widetilde{P}=\mathrm{exp}(S_{i}-\mathrm{rowmax}(S_{i}))$, where $S_{i}$ is the Matmul result of a block of $Q$ and $K^{T}$, the max value in each row of $\widetilde{P}$ is 1. Hence, we can assign a single static scale $s=\frac{1}{127}$ to a block $\widetilde{P}$, whose accuracy equals per-token quantization. Third, per-channel quantization can address the channel-wised outlier of $V$.

Data type of $\widetilde{P},V$: We choose INT8 for $\psi_{P}(\widetilde{P})$ and $\psi_{V}(V)$ because Matmul using INT8 is two times faster than using FP8 in some commonly used GPUs, and although the accuracy using $\psi_{P}(\widetilde{P})$ and $\psi_{V}(V)$ in INT8 is worse than E4M3 and E5M2, the average accuracy is similar (See Table 3).

Accuracy metrics. We use three metrics to assess the accuracy of quantized attention output $O^{\prime}$ compared to attention output in full-precision $O$: First, we flatten $O^{\prime}$ and $O$ into vectors in the shape of $1\times n$. Then, Cosine Sim$=\sum OO^{\prime}/\smash{\sqrt{\sum O^{2}}}\smash{\sqrt{\sum O^{\prime 2}}}$, Relative L1$=\sum|O-O^{\prime}|/\sum|O|$, RMSE$=\sqrt{(1/n)\sum(O-O^{\prime})^{2}}$.

**Table 4: Average accuracy using different accumulators across all layers of real models.**
| Accum. | Cos Sim $\uparrow$ | Relative L1 $\downarrow$ | RMSE $\downarrow$ |
| --- | --- | --- | --- |
| FP32 | 99.98% | 0.0156 | 2.94e-3 |
| FP16 | 99.98% | 0.0156 | 2.94e-3 |

### 4.4 FP16 accumulator: much more accurate and efficient solution

The above solution for $\psi_{P}(\widetilde{P})$ and $\psi_{V}(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,\widetilde{P},V$ across all layers
of Llama2 and Unidiffuser. It shows that INT8 $\psi_{P}(\widetilde{P})$ and $\psi_{V}(V)$ bring an unacceptable error.
In response, we propose a very accurate and also efficient solution. Specifically, we propose to use FP16 as the data type of Matmul $\widetilde{P}V$ with a 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 a FP16 accumulator is 2x faster than that with a 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 $\widetilde{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, Table 5 and 5 show the average and worst accuracy using FP16 or FP32 accumulators on all layers of Llama2 and Unidiffuser, showing that there is no accuracy loss of using the FP16 accumulator.

**Table 6: Four kernel implementations of SageAttention.**
| Kernel | $\psi_{Q}(Q)$, $\psi_{K}(K)$ | $\psi_{P}(P)$ | $\psi_{V}(V)$ |
| --- | --- | --- | --- |
| SAGEAttn-T | per-token, INT8 | FP16, FP16 Accumulator | FP16, FP16 Accumulator |
| SAGEAttn-B (Algorithm 1) | per-block, INT8 | FP16, FP16 Accumulator | FP16, FP16 Accumulator |
| SAGEAttn-vT (Figure 5(a)) | per-token, INT8 | per-block, INT8 | per-channel, INT8 |
| SAGEAttn-vB | per-block, INT8 | per-block, INT8 | per-channel, INT8 |

Figure: Figure 5: Workflow of SageAttention.
Refer to caption: x5.png

Figure: Algorithm 1 Implementation of SAGEAttn-B.

### 4.5 Adaptive Quantization

Based on the discussion in Section 4.3 and  4.4, we implement four attention kernels (See Table 6) based on two sets of choices: (1) Using $\psi_{Q}(Q)$ and $\psi_{K}(K)$ in per-token or per-block. (2) Using $\psi_{P}(\widetilde{P})$ and $\psi_{V}(V)$ in INT8 or retaining $\widetilde{P},V$ in FP16 with a FP16 accumulator.

The speed of these kernels is in the order ($\texttt{SAGEAttn-vB}>\texttt{SAGEAttn-vT}>\texttt{SAGEAttn-B}>\texttt{SAGEAttn
-T}$), but the accuracy order is opposite. Generally, SAGEAttn-B is accurate enough for all models (see Table 1) and can achieve a 2x speedup (See Figure 6 and 7). However, SAGEAttn-vB and SAGEAttn-vT are also accurate for some layers in a model. Therefore, we use various inputs to test the cosine similarity of SAGEAttn-B and SAGEAttn-T for each layer of a model. Then, we will select SAGEAttn-vB or SAGEAttn-vT if their cosine similarity is bigger than 99.8% (the worst similarity of SAGEAttn-B).

### 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., ) layer. Specifically, before the ROPE result ($A$) is written from shared memory into global memory, we perform $\delta_{A},\hat{A}=\psi(A)$. Subsequently, the $\delta_{A},\hat{A}$ are written into global memory. Additionally, we also fuse the coefficient ($1/{\sqrt{d}}$) of $QK^{T}$ into the quantization process rather than leaving it in the attention layer. Specifically, we multiply $Q$ by ($1/{\sqrt{d}}$) on chip before quantizating $Q$.

Performance Analysis. We will take SAGEAttn-B as an example to discuss the acceleration effects on actual hardware:
(1) Matmul acceleration. Utilizing INT8 matrix multiplication units on current mainstream hardware can achieve 2-4$\times$ throughput. While FP16 accumulators do not offer throughput improvements on most compute cards, on-edge accelerators, such as the RTX4090, can still achieve a 2x improvement over FP32 accumulators.
(2) Quantization overhead.
Quantization and dequantization are considered the main overhead in current quantization methods (Lin et al., ). 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, mainstream accelerators need to store data in cache (such as SharedMemory) during computation. Using 8-bit data for calculations can reduce the usage of the general cache, and using fp16 accumulators 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 can be considered negligible compared to the tensors.

## 5 Experiments

Main results. The speed of SageAttention is approximately $\textbf{2.1}\times$ faster than FlashAttention-2. Furthermore, SageAttention achieves an average real speedup of $\textbf{2.83}\times$ 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 SageAttention across a diverse set of representative models from the fields of language, image, and video generation.
Specifically, we conduct experiments on five models: Llama2 (7B) (Touvron et al., ) for text2text, CogvideoX (Yang et al., ) for text2video, Unidiffuser (Bao et al., ) and UltraPixel (Ren et al., ) for text2image, TIMM (Wightman, ) for image classification.

Datasets.
Llama2 is evaluated on three zero-shot tasks: WikiText (Merity et al., ) to assess the model’s prediction confidence, LAMBADA (Paperno et al., ) evaluate contextual understanding, and MMLU (Hendrycks et al., ) for measuring knowledge across various subjects.
CogvideoX is evaluated using the open-sora (Zheng et al., ) prompt sets.
Both UltraPixel and Unidiffuser are assessed on the COCO annotations (Lin et al., ), featuring (prompt, image) pairs.
TIMM is evaluated on on three image datasets: ImageNet (Deng et al., ), ImageNet-Sketch (Sketch) (Wang et al., ), and ImageNet-Rendition (ImageNet-r) (Hendrycks et al., ).

Metrics.
For Llama2, we use perplexity (ppl.) (Jelinek et al., ) for WikiText, and Accuracy (Acc.) for LAMBADA and MMLU.
For CogvideoX, folowing (Zhao et al., ), we evaluate the quality of generated videos on five metrics: CLIPSIM and CLIP-Temp (CLIP-T) (Liu et al., ) 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., ).
For UltraPixel and Unidiffuser, generated images are compared with the images in COCO annotations dataset in three aspects: FID (Heusel et al., ) and sFID (Salimans et al., ) for fidelity evaluation, Clipscore (CLIP) (Hessel et al., ) for text-image alignment, and ImageReward (IR) (Xu et al., ) for human preference.
For TIMM, we use Accuracy.

Figure: Figure 6: Speed comparison between SageAttention and baselines (RTX4090, headdim=64).
Refer to caption: x6.png

Figure: Figure 7: Speed comparison between SageAttention and baselines (RTX4090, headdim=128).
Refer to caption: x7.png

Figure: Figure 8: Speed comparison between SageAttention and baselines (RTX3090, headdim=64).
Refer to caption: x8.png

Figure: Figure 9: Speed comparison between SageAttention and baselines (RTX3090, headdim=128).
Refer to caption: x9.png

### 5.2 Speed and Accuracy of attention Kernels

Speed. We conduct experiments to compare the Speed of SageAttention against baselines using configurations with headdim=64 or headdim=128, both with and without Causal Mask Vaswani (). Specifically, Figure 6 and Figure 7 show the Speed of SageAttention and baselines across varying sequence lengths on RTX4090. These results indicate that SageAttention achieves a peak of 341 TOPS and is 2x faster than FlashAttention2 and 2.9x faster than xformers on average.
Figure 8 and Figure 9 illustrate the results on RTX3090, showing a similar speedup performance.

Accuracy.
Table 10 shows the numerical error of four implementations of SageAttention compared with attention in full-precision. This experiment is conducted using a set of (Q, K, V) conforming to a normal distribution. We find that the error of the four implementations is rather small. Notably, SAGEAttn-T and SAGEAttn-B achieve 100% cosine similarity and RMSE in e-4 level.

**Table 7: Real speedup (TOPS) of SageAttention on RTX4090.**
| Model | Shape of Q, K, V | Original attention | SageAttention | Speedup |
| --- | --- | --- | --- | --- |
| CogvideoX | (2, 30, 17776, 64) | 163.37 (FlashAttn2) | 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 |

**Table 8: End-to-end metrics loss across text, image, and video generation models.**
| Model | attention | WikiText (Ppl.) $\downarrow$ | Lambda (Acc.) $\uparrow$ | MMLU (Acc.) $\uparrow$ |
| --- | --- | --- | --- | --- |
| Llama2 | Full-Precision | 5.823 | 0.886 | 0.46 |
| SageAttention | 5.824 | 0.887 | 0.46 |  |

### 5.3 End-to-end Performance

Speedup. We measure the real speed of SageAttention and the original attention on Unidiffuser, UltraPixel, CogvideoX, Llama2 and TIMM 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 SageAttention compared to using attention in full-precision.
Detailed evaluation results are presented in Table 8 for Llama2, CogvideoX, Unidiffuser, UltraPixel, and TIMM, respectively. The results indicate that SageAttention successfully matches the performance of attention in full-precision across all models. Specifically, on Llama2, CogvideoX, UltraPixel, and Unidiffuser, SageAttention resulted in only a minor average degradation of 0.2% compared to attention in full-precision. Moreover, on TIMM, SageAttention even surpasses attention in full-precision.

**Table 9: Accuracy of SageAttention kernels.**
| attention | Cos Sim $\uparrow$ | Relative L1 $\downarrow$ | RMSE $\downarrow$ |
| --- | --- | --- | --- |
| SAGEAttn-T | 1.0 | 0.019 | 6.8e-4 |
| SAGEAttn-B | 1.0 | 0.021 | 7.3e-4 |
| SAGEAttn-vT | 99.9% | 0.064 | 0.065 |
| SAGEAttn-vB | 98.9% | 0.138 | 0.067 |

**Table 11: Benefit of adaptive quantization.**
| attention | model | CLIPSIM $\uparrow$ | TOPS $\uparrow$ | Model | MMLU $\uparrow$ | TOPS $\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| SAGEAttn-T | CogvideoX | 0.1827 | 292.17 | Llama2 | 0.46 | 208.59 |
| SageAttention | 0.1835 | 327.57 | 0.46 | 231.74 |  |  |

### 5.4 Ablation Study

Overhead of smoothing K. Table 10 presents the overhead associated with smoothing K on the attention speed in real models. The results indicate a 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 (SageAttention). Table 11 presents the metrics 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, 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 overhead. Second, we use FP16 accumulators in the Matmul of (P, V) to boost both accuracy and speed. Third, we use adaptive quantization to 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.

Future Work. We leave the implementation in Hopper architecture for future work.

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
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.
soundness: 4
strengths: 1. The contributions of this work are well motivated. 

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

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

4. The results are quite impressive, showing more than two times speedup without performance degradation.
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.
weaknesses: 1. Some design choices seem to be decided by the specific hardwares that are evaluated RTX4090 and 3090 (L271). Are those design choices also compatible with other GPUs like A100 and H100? 

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

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: As above in weakness.
soundness: 3
strengths: - Paper is well written.
- Experiments are thorough.
- Problem is challenging.
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.
weaknesses: - Full comparison to strong SOTA methods such as Flash attention 3, though slightly mentioned in the introduction and in Table 14, is not deeply explored. 
- Only targeted 4090/3000 series GPUs - it would be recommended to be tested on stronger GPUs at server level that is facing the strongest limitations.
- It would be great to test across VLMs too.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 3
strengths: + The paper is well-written, with a logical structure and organization that facilitates understanding.

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

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

+ The discovery of channel-wise consistency, as illustrated in Figure 4, is particularly noteworthy and adds depth to the research.
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.
weaknesses: - The method relies heavily on FlashAttention, which may weaken its technical contribution and originality. What will the performance be if it does not employ the FlashAttention as the basis?

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

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

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

[2] Q-Diffusion: Quantizing Diffusion Models

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

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 2
questions: - Is there a final attention variant among the proposed ones which performs the best uniformly across all tasks?
- L154: "The expression diag(l_i^j)^{-1}" would produce inf-s since l_i^j is initialized as 0. Or how is O_i^j initialized?
- L151: What is "online softmax" (i dont see it to be defined explicitly)? How does it differ from the standard softmax?
soundness: 3
strengths: - The performance boosts seem to be impressive compared to 
- The work carries a valuable observation that the keys matrix has huge correlation per-token. It can be valuable even outside of 
- I find the evaluation section to be very thorough, including many domains (text2image, text2video, LLMs). I especially appreciate text2video (where the sequences are the longest) and performance exploration on second-tier GPUs (e.g. 3090), which are more suitable for production use cases.
- I am not an expert in quantization, but I checked some prior works and didn't find similar approaches. This makes me consider the paper as novel.
- The paper provides the source code: I have not tried running it but it is valuable.
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.
weaknesses: - I think that writing should be improved. A table I would like to see the most is the end-to-end model speed improvement and metrics degradation after integrating the proposed attention mechanism. First, for many models accelerating the attention op itself does not necessarily lead to much improved overall speed when the sequence length is short and it's MLPs who carry the main burden. Second, it's not entirely clear when looking at various tables which variant is being used at the end (and why there are so many of them throughout the paper instead of being just in some restricted ablation section). Finally, there are quite many typos and grammatical errors, e.g.:
    - L084: "matrice" => "matrices"
    - L138 (and elsewhere): "IO" => "I/O"
    - L144: "it compute" => "it computes"
    - L142=144: missing line or equation
    - L151: "The σ()" => "The" should be ommited or it should be "The function σ()".
    - L184: "First, We" => "First, we"
    - L264-L266: ~3 typos in 3 lines
    - etc.
- There are no qualitatives provided for image/video generation.

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

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

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

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

Decision: Accept (Poster)