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. Эта инструкция работает в 2 раза быстрее, чем FP8 Matmul, используемая в SageAttention2. Наши эксперименты показывают, что SageAttention2++ достигает ускорения в 3,9 раза по сравнению с FlashAttention, сохраняя при этом ту же точность внимания, что и 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.