SageAttention2++: Uma Implementação Mais Eficiente do SageAttention2

SageAttention2++: A More Efficient Implementation of SageAttention2

May 27, 2025

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

cs.AI

Resumo

A eficiência da atenção é crucial porque sua complexidade temporal cresce quadraticamente com o comprimento da sequência. O SageAttention2 aborda isso utilizando quantização para acelerar as multiplicações de matrizes (Matmul) na atenção. Para acelerar ainda mais o SageAttention2, propomos utilizar a instrução mais rápida de Matmul FP8 acumulada em FP16. Essa instrução é 2x mais rápida que a Matmul FP8 usada no SageAttention2. Nossos experimentos mostram que o SageAttention2++ alcança um ganho de velocidade de 3,9x em relação ao FlashAttention, mantendo a mesma precisão de atenção do SageAttention2. Isso significa que o SageAttention2++ acelera efetivamente diversos modelos, incluindo aqueles para geração de linguagem, imagens e vídeos, com perda insignificante nas métricas de ponta a ponta. O código estará disponível em 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.