SageAttention2++: SageAttention2のより効率的な実装

SageAttention2++: A More Efficient Implementation of SageAttention2

May 27, 2025

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

cs.AI

要旨

アテンションの効率性は極めて重要です。なぜなら、その時間計算量はシーケンス長に対して二次関数的に増加するからです。SageAttention2は、アテンションにおける行列積（Matmul）を加速するために量子化を活用することでこの課題に対処しています。SageAttention2をさらに高速化するために、FP8 MatmulをFP16で累積する高速な命令を活用することを提案します。この命令は、SageAttention2で使用されているFP8 Matmulよりも2倍高速です。実験結果によると、SageAttention2++はFlashAttentionと比較して3.9倍の高速化を達成しつつ、SageAttention2と同等のアテンション精度を維持しています。これは、SageAttention2++が言語、画像、動画生成を含む様々なモデルを効果的に加速し、エンドツーエンドのメトリクスの損失を無視できるレベルに抑えられることを意味します。コードはhttps://github.com/thu-ml/SageAttentionで公開予定です。

English

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.