SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization

Motivation

• Attention’s quadratic time complexity makes efficient implementation crucial as sequence lengths increase, while linear and sparse methods are suitable only for a limited range of models and tasks.Hardware-oriented methods such as FlashAttention, xformers, and SageAttention retain computations across the sequence.Source: §1 · p. 1

Challenges

• Per-tensor INT4 quantization makes CogvideoX videos blurry, while Llama3 reaches only 25% MMLU accuracy; limited INT4 range and FP8 accumulator precision are identified challenges.The paper describes the FP8 tensor-core accumulator as FP22, with 1 sign bit, 8 exponent bits, and 13 mantissa bits.Source: §1 · pp. 1–2

Performance

• CogvideoX (1.5-5B) generation is 1.8x faster with SageAttention2 than FlashAttention2, with no video quality loss.The figure compares videos generated using the two methods on RTX4090.Source: Figure 1 · p. 2

Our approach

• The method quantizes Q and K to INT4 per thread, smooths Q, and uses FP8 for P̃ and V with two-level accumulation for P̃V.It also subtracts Q’s channel mean and adds the corresponding Q̄mK term after QKᵀ; smoothing V is optional.Source: §1 · p. 2

Performance

• SageAttention2 is about 3x faster than FlashAttention2 and 4.5x faster than xformers, while maintaining end-to-end metrics across language, image, and video generation models.On Hopper GPUs, it matches FlashAttention3(fp8) speed with higher accuracy.Source: §1 · pp. 2–3

2.1. FlashAttention

• FlashAttention tiles Q, K, and V by token blocks and computes output tiles in parallel, iteratively loading K and V blocks and updating outputs with online softmax.Attention’s two matrix multiplications have O(N^2d) complexity.Source: §2.1 · p. 3

2.2. Quantization

• Quantization approximates C = AB by multiplying low-precision quantized matrices and scaling the result; per-tensor INT4 uses the tensor maximum absolute value to set its scale.Low-precision matrix multiplication is usually multiple times faster than higher-precision multiplication on modern GPUs.Source: §2.2 · p. 3

2.3. SageAttention

• SageAttention quantizes Q and K to INT8 per block, smooths K by subtracting its token-wise mean, and retains P̃ and V in FP16.It uses an FP16 accumulator for P̃V; this accelerates multiplication on RTX4090 but not on L20, L40, or H100.Source: §2.3 · p. 3

3.1. Smooth Q

• Smoothing Q+K significantly improves end-to-end accuracy on Llama3.1 and CogvideoX, ranking ahead of smoothing Q, smoothing K, and other baselines.The reported effectiveness order is smoothing Q+K > smoothing Q > smoothing K > other baselines.Source: §3.1 · p. 4

3.2. INT4 Per-thread Quantization

• Per-thread INT4 quantization achieves accuracy very close to per-token quantization and significantly outperforms coarser granularities, with almost no speed degradation.Per-token quantization introduces noticeable overhead because it reduces hardware efficiency.Source: §3.2 · p. 5

3.3. FP8 quantization for ˜PV

• E4M3 FP8 accuracy is very close to FP16 and superior to E5M2 and INT8; FP8 tensor-core Matmul is twice as fast as FP16.The comparison uses real tensors across CogvideoX layers with a 32-bit accumulator.Source: §3.3 · pp. 5–6

3.4. FP32 MMA Buffer for FP22 Accumulator

• The two-level strategy accumulates each FP22 block result in an FP32 buffer, limiting low-precision accumulation to a small token block.The authors report being the first to investigate FP22 accumulator effects and apply this strategy to attention.Source: §3.4 · p. 6
• Smoothing V is optional and benefits accuracy mainly when V has channel-wise bias, which is absent in some models such as Llama3.1.The technique is not used in the main experiments.Source: §3.4 · p. 6

4. Experiment

• SageAttention2-8b matches FlashAttention3(fp8) speed on Hopper GPUs and is much more accurate, while SageAttention2 maintains end-to-end metrics across language, image, and video generation models.Source: §4 · p. 6

4.2. Speed and Accuracy of Kernels

• SageAttn2-4b and SageAttn2-8b are approximately 3x and 2.7x faster than FlashAttention2, and 4.5x and 4x faster than xformers, respectively.The comparison covers varying sequence lengths on RTX4090.Source: Figure 5 · p. 7, §4.2 · p. 9

4.3. End-to-end Performance

• SageAttn2-4b outperforms all baselines while maintaining most end-to-end accuracy across models; SageAttn2-8b incurs almost no metrics loss.Table 2 reports metrics across text, image, and video generation models.Source: Table 2 · p. 7, §4.3 · p. 9

4.1. Setup

• The experiments cover ten representative models spanning text-to-text, text-to-video, text-to-image, and image classification; the Hopper-oriented 8-bit variant omits Q smoothing.Two CUDA attention kernels are implemented, and the baselines include SmoothAttn, HadmdAttn, SageAttention, and FlashAttn3(fp8).Source: §4.1 · p. 8, Table 3 · p. 8

4.2. Speed and Accuracy of Kernels

• SageAttn2-4b has superior average accuracy to other baselines across CogvideoX layers with INT4 Q, K and FP8 P, V.The reported comparison is summarized in Tables 4 and 17.Source: Table 4 · p. 9, §4.2 · p. 9

4.3. End-to-end Performance

• A 1.8x speedup in CogvideoX (1.5-5B) with no metrics loss is reported for SageAttn2-8b; SageAttn2-4b further accelerates models with a little metrics loss.The comparison concerns generation latency for models with long sequence lengths.Source: Table 8 · p. 9, §4.3 · p. 9
• Across HunyuanVideo, Mochi, and CogvideoX examples, SageAttn2-8b shows no visible differences from full-precision attention, while SageAttn2-4b has minor differences and is much better than baselines.Source: Figure 6 · p. 8, Figure 7 · p. 8, §4.3 · p. 9

4.4. Ablation Study

• Per-thread quantization, Q smoothing, and two-level accumulation add 0.35%, 3.7%, and 0% overhead on kernel speed, respectively.The overhead comparison is reported in Table 18.Source: §4.4 · p. 9

• SageAttention2 quantizes Q, K at thread-level granularity to INT4 and P, V to FP8.Source: §5 · p. 9
• Q smoothing and two-level accumulation are proposed to improve the accuracy of INT4 QK^T and FP8 PV, respectively.Source: §5 · p. 9
• SageAttention2 is approximately 3x faster than FlashAttention2 and 4.5x faster than xformers, and matches FlashAttention3(fp8) speed on Hopper GPUs with higher accuracy.The reported end-to-end metrics are maintained across language, image, and video generation models.Source: §5 · p. 9

• The authors aim to advance machine learning and do not identify potential societal consequences requiring specific emphasis.Source: Impact Statement · p. 10

• Figure 8 shows CogvideoX (2B) comparison examples generated from open-sora prompt-set prompts.Source: Figure 8 · p. 15

A.2. Additional Kernel Speed Comparison

• Speed comparisons cover RTX4090, L20, H100, and H20 at head dimensions 64 and 128, with and without causal masking.The passage identifies Figures 10–16 as the comparisons.Source: §A.2 · p. 15
• Table 9 summarizes the performance gains of attention methods against baselines on modern GPUs.The supplied passage does not give the table's individual speedup values.Source: §A.2 · p. 15, Table 9 · p. 15
• Table 10 presents a CogvideoX attention-accuracy example with and without smoothing V.Its caption does not state which setting is more accurate.Source: Table 10 · p. 15

A.3. Smoothing V

• This centering addresses positive or negative channel-wise biases in V, which can make ePV values large when multiplied by eP rows ranging from 0 to 1.CogvideoX is given as an example of channel values ranging from 8 to 9.Source: §A.3 · p. 15
• Subtracting each V column's mean moves values toward zero, where FP22 represents them more precisely.The mean is retained in FP16 and added back to the output because each row of eP sums to 1.Source: §A.3 · p. 15
• Figure 9 compares SageAttn2-8b with FlashAttention3 on Mochi and HunyuanVideo.The caption does not report the comparison outcome.Source: Figure 9 · p. 16

A.4. Experiment of Smoothing V

• On real CogvideoX tensors, smoothing V improves SageAttention2 accuracy with INT4 Q,K and FP8 eP,V, and is generally effective for diffusion models.Table 10 presents the with/without-smoothing accuracy comparison.Source: §A.4 · p. 16

A.5. Theoretical Analysis of Smoothing

• Under Gaussian activation and i.i.d. token assumptions, smoothing concentrates values toward zero and reduces absolute maxima, which lowers quantization error.The analysis relates smaller group maxima to smaller round-to-nearest quantization error.Source: §A.5 · pp. 16, 17, 19

A.6. Per-Thread Quantization Formulation

• Per-thread quantization groups query and key tokens so each thread uses one dequantization scale, avoiding per-token dequantization overhead.For a 128x64 block, this gives 32 Q groups and 4 K groups, with 32× and 4× finer granularity than per-block quantization, respectively.Source: §A.6 · p. 20

A.7. Datasets and Metrics in Experiments

• Evaluation covers language, video, image-generation, and image-classification tasks, using task-specific end-to-end metrics and quantized-output accuracy measures.The accuracy measures compare quantized attention output O′ against full-precision O.Source: §A.7 · p. 20

A.8. Kernel Benchmark Setup

• Kernel benchmarks use batch size 4, 32 heads, head dimensions 64 or 128, optional causal masking, and Gaussian floating-point or range-uniform integer inputs.The integer input ranges are [-128, 127] for INT8 and [-8, 7] for INT4.Source: §A.8 · p. 21

A.9. Additional Experiments and Analysis

• Tables 11–13 report end-to-end metrics for Llama2, CogvideoX, and an image-classification model.Their captions identify the evaluated models but do not supply metric values.Source: Table 11 · p. 21, Table 12 · p. 21, Table 13 · p. 21
• On Llama-3-262k, SageAttention2-8b maintains model performance while FlashAttention3(fp8) degrades; SageAttention2 also consistently outperforms baselines on Qwen2-Audio ASR.The long-context evaluation reaches 262k tokens on H100; the audio evaluation uses Librispeech and WER.Source: Table 14 · p. 21, §A.9 · p. 23, Table 20 · p. 23