---
title: 'SageAttention2: Fast Quantized Attention'
url: https://www.emergentmind.com/topics/sageattention2-f8c247cf-f64b-41bc-9fa6-0e53baaaf06c
type: topic
---

# SageAttention2: Fast Quantized Attention

SageAttention2 is an efficient attention mechanism designed to address the quadratic time complexity of standard attention operations, specifically by leveraging low-bit quantization to accelerate the critical matrix multiplications within attention kernels. The method targets both the $QK^T$ computation for similarity scores and the $PV$ computation for weighted value aggregation, quantizing on-the-fly and optimizing for NVIDIA TensorCore architecture. SageAttention2++ introduces further optimizations, achieving substantial kernel throughput improvements by utilizing FP8 multiplication with FP16 accumulation, minimizing range overflow, and reducing data conversion overhead. Empirical results demonstrate that SageAttention2++ attains up to a 3.9× speedup over FlashAttention2, with negligible degradation in end-to-end metrics for language, image, and video generation tasks [2505.21136].

## 1. Core Algorithmic Modifications

SageAttention2++ modifies the matrix multiplication paradigm used in SageAttention2. In typical attention implementations (e.g., FlashAttention-style), computational expense per head is dominated by two operations:

- $QK^T$: Shapes $N \times d \times d \times N \to N \times N$
- $PV$: Shapes $N \times d \times d \times D \to N \times D$ (where $D$ is head-value dimension)

SageAttention2 quantizes $Q$ and $K$ as INT4/INT8 and $P$ and $V$ as FP8 (E4M3), performing $PV$ multiplication with accumulation in FP32 registers using the mma.f32.f8.f8.f32 TensorCore instruction. SageAttention2++ replaces this with FP8 × FP8 $\rightarrow$ FP16 accumulation via mma.f16.f8.f8.f16, supported by Ada-generation GPUs such as RTX 4090 and RTX 5090. This change yields approximately 4× faster throughput compared to native FP16 MatMul and 2× faster than FP8 $\rightarrow$ FP32 accumulation.

Range narrowing is required because FP16 accumulators can overflow (${|\text{FP16}_{\mathrm{max}}|} \sim 65,504$). To control this, $p$ and $v$ are bounded such that $32 \cdot p \cdot v \leq 65,504$, with $P_r = 112$, $V_r = 4.5$ chosen so that delayed FP32 buffering remains in range: $P_r \cdot V_r \leq 2047/2$. Delayed FP32 buffering reduces PTX conversion cost by accumulating two FP16 results before conversion.

## 2. Quantization Methodology and Formulas

SageAttention2++ employs block-wise quantization for $Q, K, P, V$ with specific schemes:

- **$Q, K$ quantization:** INT4/INT8 block-wise. Given $Q_i \in \mathbb{R}^{B \times d}$,
  $$
  \delta_Q = \max_i |Q_i| / S_Q, \quad \hat{Q}_i = \operatorname{round}(Q_i / \delta_Q), \quad S_Q = 15 \text{ for INT4},\ 127 \text{ for INT8}
  $$
  Reconstruction: $Q_i \approx \hat{Q}_i \cdot \delta_Q$

- **$P$ quantization:** FP8 (E4M3) per block. Softmax-attention $P = \mathrm{softmax}(QK^T)$,
  $$
  \delta_P = \max |P| / P_r, \quad \hat{P} = \operatorname{round}(P / \delta_P)
  $$
  with $P_r = 112$.

- **$V$ quantization:** FP8 (E4M3) per-channel,
  $$
  \delta_V(j) = \max_{\text{row}} |V(:, j)| / V_r,\quad \hat{V}(:, j) = \operatorname{round}(V(:, j) / \delta_V(j))
  $$
  with $V_r = 4.5$.

- **$PV$ accumulation:** FP16 via TensorCore,
  $$
  \text{out}_{\text{FP16}} = \text{mma.f16.f8.f8.f16}(\hat{P}, \hat{V}),\quad O = \text{out}_{\text{FP16}} \cdot \delta_P \cdot \delta_V
  $$
  Together:
  $$
  \hat{P} = \lfloor P / \delta_P \rceil, \quad \hat{V} = \lfloor V / \delta_V \rceil, \quad O = (\langle \hat{P}, \hat{V} \rangle_{\mathrm{FP16}\ \text{acc.}})\cdot\delta_P\delta_V
  $$

## 3. Computational Complexity and Throughput

Full-precision attention operates at $O(N^2 d)$ for $QK^T$ and $O(NdD)$ for $PV$ per head. SageAttention2++ maintains these asymptotic bounds but improves the constant factor. Defining $F_{\mathrm{FP16}}$ as FP16 MatMul throughput and $F_{\mathrm{F8F16}}$ as FP8$\times$FP8$\rightarrow$FP16 throughput:
$$
\text{Speedup} \approx F_{\mathrm{F8F16}} / F_{\mathrm{FP16}} \approx 4\times
$$
for $PV$. For overall end-to-end kernel time $T$:
$$
T_{\mathrm{full}} \approx \frac{C_{QK} N^2 d}{F_{\mathrm{FP16}}} + \frac{C_{PV} N d D}{F_{\mathrm{FP16}}}
$$
$$
T_{\mathrm{SA2++}} \approx \frac{C_{QK} N^2 d}{F_{\mathrm{INT4}}} + \frac{C_{PV} N d D}{F_{\mathrm{F8F16}}}
$$
with $F_{\mathrm{INT4}} \gtrsim 2\times F_{\mathrm{FP16}}$, leading to observed kernel speedups of 3–3.9× over FlashAttention2.

## 4. Empirical Benchmarks

Tests on NVIDIA RTX 4090 and 5090 (Ada Lovelace) with head-dimensions 64 and 128 and sequence lengths up to 8k revealed peak kernel speedups:
- SageAttention2++(4+8) (INT4 $QK$, FP8 $PV$): $\sim$3.9× vs. FlashAttention2
- SageAttention2++(8+8) (INT8 $QK$, FP8 $PV$): $\sim$3.0×

Consistent gains were observed for both causal and non-causal masks. End-to-end metrics for representative models:

| Model                                | Attention Variant           | Perplexity or Metric (Delta vs FP32) |
|--------------------------------------|----------------------------|--------------------------------------|
| Llama3.1(8B) (language)              | Full-prec                  | Ppl 6.013                            |
|                                      | SageAttn2                  | Ppl 6.019                            |
|                                      | SageAttn2++(8+8)           | Ppl 6.020                            |
| CogvideoX(2B) text$\to$video         | Full-prec                  | CLIPSim 0.179 / FScore 4.974         |
|                                      | SageAttn2(8+8)             | CLIPSim 0.178 / FScore 4.899         |
|                                      | SageAttn2++(8+8)           | CLIPSim 0.179 / FScore 4.386         |
| Flux/StableDiffusion3.5 text$\to$image| Full-prec vs SageAttn2++   | $\lesssim$0.5 FID, $\lesssim$0.02 sFID|

Across language, image, and video models, SageAttention2++(8+8) matches SageAttention2's metrics, while (4+8) variant incurs only slight degradation.

## 5. Architectural Integration and Implementation

SageAttention2++ functions as a drop-in replacement for the $PV$ kernel in FlashAttention-style fused kernels. Two new CUDA kernels (for the two quantization modes) are required, invoked in lieu of torch.flash_attn(). Hardware compatibility mandates support for mma.f16.f8.f8.f16 instructions (Ada Lovelace or later GPUs) and sufficient shared memory for quantized $P, V$ blocks; tiling follows FlashAttention conventions.

FP8 matmuls introduce two implementation caveats:
- FP16 accumulation range ($\sim \pm 6.55\times 10^4$) requires range-narrowing and block-wise scale factors.
- Delayed FP32 buffering maintains efficiency but needs careful PTX scheduling.

## 6. Conclusions and Future Directions

SageAttention2++ demonstrates that employing FP16-accumulating FP8 MatMul, together with well-designed quantization bounds and FP16 buffering, can yield up to 4× kernel throughput improvements with negligible accuracy loss. Evaluations on tasks spanning language modeling, image generation, and video synthesis reveal up to a 3.9× reduction in kernel latency relative to FlashAttention2 while maintaining nearly identical performance metrics as SageAttention2.

Areas proposed for future investigation include:
- Lower-bit accumulation, such as FP4$\to$FP8$\to$FP16 chains
- Dynamic range adaptation per token/block for improved quantization efficiency
- Integration with sparse or linear attention paradigms
- Exploiting new hardware instructions (e.g., Hopper asynchronous FP8)

The reference implementation is slated for release at https://github.com/thu-ml/SageAttention [2505.21136].

Source: https://www.emergentmind.com/topics/sageattention2-f8c247cf-f64b-41bc-9fa6-0e53baaaf06c