SageAttention2: Fast Quantized Attention

• SageAttention2 is an efficient attention mechanism that employs low-bit quantization to accelerate QK^T and PV computations in standard attention models.
• It utilizes INT4/INT8 for Q/K and FP8 for P/V with FP16 accumulation, achieving up to 3.9× faster throughput compared to FlashAttention2.
• Empirical benchmarks across language, image, and video tasks demonstrate significant speedups while maintaining nearly identical accuracy to full-precision models.

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 (Zhang et al., 27 May 2025).

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 $PV$0 as FP8 (E4M3), performing $PV$1 multiplication with accumulation in FP32 registers using the mma.f32.f8.f8.f32 TensorCore instruction. SageAttention2++ replaces this with FP8 × FP8 $PV$2 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 $PV$3 FP32 accumulation.

Range narrowing is required because FP16 accumulators can overflow ($PV$4). To control this, $PV$5 and $PV$6 are bounded such that $PV$7, with $PV$8, $PV$9 chosen so that delayed FP32 buffering remains in range: $QK^T$0. 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 $QK^T$1 with specific schemes:

• $QK^T$2 quantization: INT4/INT8 block-wise. Given $QK^T$3,

$QK^T$4

Reconstruction: $QK^T$5

• $QK^T$6 quantization: FP8 (E4M3) per block. Softmax-attention $QK^T$7,

$QK^T$8

with $QK^T$9.

• $N \times d \times d \times N \to N \times N$0 quantization: FP8 (E4M3) per-channel,

$N \times d \times d \times N \to N \times N$1

with $N \times d \times d \times N \to N \times N$2.

• $N \times d \times d \times N \to N \times N$3 accumulation: FP16 via TensorCore,

$N \times d \times d \times N \to N \times N$4

Together:

$N \times d \times d \times N \to N \times N$5

3. Computational Complexity and Throughput

Full-precision attention operates at $N \times d \times d \times N \to N \times N$6 for $N \times d \times d \times N \to N \times N$7 and $N \times d \times d \times N \to N \times N$8 for $N \times d \times d \times N \to N \times N$9 per head. SageAttention2++ maintains these asymptotic bounds but improves the constant factor. Defining $PV$0 as FP16 MatMul throughput and $PV$1 as FP8$PV$2FP8$PV$3FP16 throughput:

$PV$4

for $PV$5. For overall end-to-end kernel time $PV$6:

$PV$7

$PV$8

with $PV$9, 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 $N \times d \times d \times D \to N \times D$0, FP8 $N \times d \times d \times D \to N \times D$1): $N \times d \times d \times D \to N \times D$23.9× vs. FlashAttention2
• SageAttention2++(8+8) (INT8 $N \times d \times d \times D \to N \times D$3, FP8 $N \times d \times d \times D \to N \times D$4): $N \times d \times d \times D \to N \times D$53.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$N \times d \times d \times D \to N \times D$6video	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$N \times d \times d \times D \to N \times D$7image	Full-prec vs SageAttn2++	$N \times d \times d \times D \to N \times D$80.5 FID, $N \times d \times d \times D \to N \times D$90.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 $D$0 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 $D$1 blocks; tiling follows FlashAttention conventions.

FP8 matmuls introduce two implementation caveats:

• FP16 accumulation range ($D$2) 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$D$3FP8$D$4FP16 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 (Zhang et al., 27 May 2025).