SageAttention2++: Una implementación más eficiente de 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

Resumen

La eficiencia de la atención es crucial porque su complejidad temporal crece cuadráticamente con la longitud de la secuencia. SageAttention2 aborda esto utilizando cuantización para acelerar las multiplicaciones de matrices (Matmul) en la atención. Para acelerar aún más SageAttention2, proponemos utilizar la instrucción más rápida de Matmul FP8 acumulada en FP16. Esta instrucción es 2x más rápida que la Matmul FP8 utilizada en SageAttention2. Nuestros experimentos muestran que SageAttention2++ logra una aceleración de 3.9x sobre FlashAttention mientras mantiene la misma precisión en la atención que SageAttention2. Esto significa que SageAttention2++ acelera eficazmente varios modelos, incluyendo aquellos para generación de lenguaje, imágenes y video, con una pérdida insignificante en las métricas de extremo a extremo. El código estará disponible en 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.