SageAttention2++:SageAttention2 的更高效实现
机器学习
2025-06-09 v3 人工智能
硬件体系结构
计算机视觉与模式识别
摘要
注意力的效率至关重要,因为其时间复杂度随序列长度呈二次增长。SageAttention2 通过利用量化加速注意力中的矩阵乘法(Matmul)来解决这一问题。为了进一步加速 SageAttention2,我们提出使用在 FP16 中累加的更快的 FP8 Matmul 指令。该指令比 SageAttention2 中使用的 FP8 Matmul 快 2 倍。我们的实验表明,SageAttention2++ 在保持与 SageAttention2 相同注意力精度的同时,比 FlashAttention 实现了 3.9 倍的加速。这意味着 SageAttention2++ 有效地加速了包括语言、图像和视频生成在内的各种模型,且端到端指标损失可忽略不计。代码将在 https://github.com/thu-ml/SageAttention 上提供。
引用
@article{arxiv.2505.21136,
title = {SageAttention2++: A More Efficient Implementation of SageAttention2},
author = {Jintao Zhang and Xiaoming Xu and Jia Wei and Haofeng Huang and Pengle Zhang and Chendong Xiang and Jun Zhu and Jianfei Chen},
journal= {arXiv preprint arXiv:2505.21136},
year = {2025}
}