PaperlayerView exact revision on arXiv

Source-linked AI summary

SageAttention2++: A More Efficient Implementation of SageAttention2

Jintao Zhang, Xiaoming Xu, Jia Wei, Haofeng Huang, Pengle Zhang, Chendong Xiang, Jun Zhu, Jianfei Chen

arXiv:2505.21136v3cs.LGcs.AIcs.ARcs.CV

TL;DR

Attention’s quadratic cost makes efficient long-sequence computation important. SageAttention2++ accelerates SageAttention2 by using FP8 PV Matmul accumulated in FP16, with narrowed quantization ranges to keep values representable. It achieves 3.9× speedup over FlashAttention while matching SageAttention2’s attention accuracy and incurring negligible end-to-end metrics loss across diverse models.

  • Problem

    Attention’s quadratic time complexity makes efficient implementations important for real-world applications with long sequences.

    Source: §1 · p. 1

  • Method

    SageAttention2++ accelerates SageAttention2 by using FP8 Matmul with an FP16 accumulator for PV and narrowing P and V quantization ranges to fit FP16’s representable range.

    Source: §1 · p. 1, §3 · p. 2

  • Results

    3.9× speedup over FlashAttention while maintaining the same attention accuracy as SageAttention2, with negligible end-to-end metrics loss across diverse models.

    Source: §1 · p. 1, §5 · p. 5

  • Takeaways & Limitations

    SageAttention2++ accelerates language, image, and video generation models with negligible end-to-end metrics loss.

    Source: §5 · p. 5

Abstract

from arXiv · show

The efficiency of attention is critical because its time complexity grows quadratically with sequence length. SageAttention2 addresses this by utilizing quantization to accelerate matrix multiplications (Matmul) in attention. To further accelerate SageAttention2, we propose to utilize the faster instruction of FP8 Matmul accumulated in FP16. The instruction is 2x faster than the FP8 Matmul used in SageAttention2. Our experiments show that SageAttention2++ achieves a 3.9x speedup over FlashAttention while maintaining the same attention accuracy as SageAttention2. This means SageAttention2++ effectively accelerates various models, including those for language, image, and video generation, with negligible end-to-end metrics loss. The code will be available at https://github.com/thu-ml/SageAttention.

1. Introduction

Attention’s quadratic time complexity motivates efficient implementations for long sequences, while linear and sparse approaches can have limited generality across models and tasks. SageAttention2++ accelerates SageAttention2 with FP8 matrix multiplication accumulated in FP16, achieving 3.9× speedup over FlashAttention while matching SageAttention2’s attention accuracy.

  • Attention’s quadratic time complexity motivates efficient long-sequence implementations; linear and sparse approaches can have limited generality across models and tasks.Hardware-optimized attention instead maintains full-sequence computation.Source: §1 · p. 1
  • SageAttention2++ uses the faster FP8 Matmul instruction with an FP16 accumulator for the attention PV matrix multiplication.Direct use risks exceeding FP16’s representable range, so the method narrows P and V quantization ranges.Source: §1 · p. 1
  • 3.9× speedup over FlashAttention, while matching SageAttention2’s attention accuracy, with negligible end-to-end metrics loss across evaluated text, image, and video generation models.The paper reports comprehensive evaluations and plug-and-play acceleration across these model domains.Source: §1 · p. 1

2. Preliminary

SageAttention2 is a FlashAttention-based quantization method that uses low-bit Tensor Core operations to accelerate attention matrix multiplications. The preliminary section also contrasts Matmul accumulator types and their speedups.

2. Preliminary

  • Table 1 summarizes Matmul instruction speedups relative to FP16 matrix multiplication with an FP32 accumulator.Source: Table 1 · p. 1

2.1. SageAttention2

  • SageAttention2 quantizes Q and K to INT4/INT8, eP and V to FP8, and rescales the quantized values to compute PV.Q, K, and eP use per-block scale factors, while V uses per-channel scaling.Source: §2.1 · pp. 1–2

2.2. Data Type of Accumulator for Matmul

  • FP8 Matmul with an FP16 accumulator is 4× faster than FP16, whereas FP8 Matmul with an FP32 accumulator is only 2× faster.Source: §2.2 · p. 2

3. SageAttention2++

SageAttention2++ retains SageAttention2’s workflow but uses FP8 PV Matmul with an FP16 accumulator. It narrows P and V quantization ranges to stay within FP16’s range and delays FP32 conversion to reduce overhead.

3. SageAttention2++

  • SageAttention2++ keeps SageAttention2’s smoothing and QK⊤ quantization workflow, but uses FP8 Matmul with an FP16 accumulator for PV.Source: §3 · p. 2

3.1. Narrowing the FP8 Quantization Range

  • FP8 values for P and V can make accumulated PV exceed FP16’s representable range, so the method narrows their quantization ranges.The cited example |p| ≤224 and |v| ≤9 satisfies the range condition.Source: §3.1 · p. 2

3.2. Delayed FP32 Buffering

  • Accumulating two consecutive MMA results in FP16 before FP32 conversion halves conversion overhead while preserving the FP16 range constraint.Source: §3.2 · p. 2

3.1. Narrowing the FP8 Quantization Range

  • The selected scales Pr = 224 and Vr = 4.5 yield negligible attention-accuracy error among the feasible pairs and were chosen for optimal performance.Source: §3.2 · p. 2

4. Experiment

SageAttention2++ achieves up to 3.9× speedup over FlashAttention2 with negligible impact on end-to-end metrics across diverse model architectures. Experiments compare kernel speeds and end-to-end performance across text, image, and video-generation models.

4. Experiment

  • Up to 3.9× speedup over FlashAttention2, with negligible impact on end-to-end metrics across diverse model architectures.The evaluation spans representative language, image, and video-generation models.Source: §4 · p. 2

4.1. Setup

  • Evaluation spans text, image, and video-generation models, comparing INT8 and INT4 Q,K variants against FlashAttention2 and SageAttention baselines.The paper reports a CUDA implementation.Source: §4.1 · pp. 2–3

4.2. Speed of Kernels

  • On RTX4090, SageAttn2++(4+8) and SageAttn2++(8+8) are approximately 3.9× and 3.0× faster than FlashAttention2, respectively.The comparisons also cover headdim=64 and RTX5090 configurations.Source: Figure 1 · p. 3, Figure 2 · p. 3, Figure 3 · p. 3

4.3. End-to-end Performance

  • SageAttn2++(8+8) matches SageAttention2’s end-to-end metrics with almost no loss, while SageAttn2++(4+8) brings a little metrics loss.Table 3 covers end-to-end metrics across text, image, and video-generation models.Source: Table 3 · p. 4, §4.3 · p. 5
  • Visible comparison examples are presented in Figures 5, 7, and 6.Figure 5 is identified as an example of using SageAttention2++.Source: Figure 5 · p. 4, §4.3 · p. 5

5. Conclusion

SageAttention2++ uses a faster FP8 Matmul instruction with FP16 accumulation for the PV multiplication. It achieves a 3.9× speedup over FlashAttention while maintaining SageAttention2’s attention accuracy.

  • SageAttention2++ uses FP8 Matmul accumulated in FP16 for the PV multiplication.Source: §5 · p. 5
  • 3.9× speedup over FlashAttention while maintaining the same attention accuracy as SageAttention2.Source: §5 · p. 5
  • The method accelerates language, image, and video-generation models with negligible end-to-end metrics loss.Source: §5 · p. 5

A.1. Visible Comparison Examples

Figure 7 presents visible examples of using SageAttention2++ on image generation.

  • Figure 7 presents visible examples of SageAttention2++ on image generation.Source: Figure 7 · p. 9

A.2. Datasets and Metrics in Experiments

The experiments use task-specific datasets and end-to-end metrics for text, video, and image generation. Quantized attention output is also assessed against full-precision output using three accuracy metrics.

  • Text-to-text evaluation uses WikiText, LAMBADA, and Needle-in-A-Haystack; video uses open-sora prompt sets, and image uses COCO annotations.Source: §A.2 · p. 9
  • Text-to-text metrics are perplexity for WikiText and accuracy for LAMBADA and NIAH.Source: §A.2 · p. 9
  • Video-generation metrics cover text-video alignment, video aesthetic and technical quality, and temporal consistency.The named metrics are CLIPSIM and CLIP-Temp, VQA-a and VQA-t, and Flow-score, respectively.Source: §A.2 · p. 9
  • Image-generation evaluation uses FID and sFID for fidelity, CLIPscore for text-image alignment, and ImageReward for human preference.Source: §A.2 · p. 9
  • Quantized attention output O′ is compared with full-precision O using cosine similarity, relative L1 distance, and RMSE.Source: §A.2 · p. 9

Related papers

Loading 2505.21136v3…