SageAttention2++: A More Efficient Implementation of SageAttention2

TL;DR

The paper tackles the quadratic-time bottleneck of attention for long sequences by enhancing hardware-optimized quantized attention. It extends SageAttention2 to SageAttention2++ by using FP16-accumulated FP8 Matmul for the PV term, and by narrowing FP8 ranges and buffering to keep FP16 representable, achieving a $3.9x$ speedup over FlashAttention2 while preserving accuracy. The method is evaluated across language, image, and video generation models, showing consistent speedups with negligible end-to-end loss. This work provides a practical, plug-and-play approach to accelerate attention across diverse modalities, with code released for community use.

Abstract

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.

Figures (7)

• Figure 1: Speed comparison between SageAttention2++ and baselines (RTX4090, headdim=128).
• Figure 2: Speed comparison between SageAttention2++ and baselines (RTX4090, headdim=64).
• Figure 3: Speed comparison between SageAttention2++ and baselines (RTX5090, headdim=128).
• Figure 4: Speed comparison between SageAttention2++ and baselines (RTX5090, headdim=64).
• Figure 5: A visible example of using SageAttention2++.