⏶43
SageAttention2++: SageAttention2 的更高效实现
发表
由 
Jintao Zhang 提交
Jintao Zhang 提交作者: Jintao Zhang, Xiaoming Xu, Jia Wei, Haofeng Huang, Pengle Zhang, Chendong Xiang, Jun Zhu, Jianfei Chen
Jintao Zhang, Xiaoming Xu, Jia Wei, Haofeng Huang, Pengle Zhang, Chendong Xiang, Jun Zhu, Jianfei Chen摘要
注意力的效率至关重要，因为其时间复杂度随序列长度呈二次方增长。SageAttention2 通过利用量化来加速注意力机制中的矩阵乘法（Matmul）来解决这一问题。为了进一步加速 SageAttention2，我们建议利用 FP8 Matmul 在 FP16 中累积的更快指令。该指令比 SageAttention2 中使用的 FP8 Matmul 快 2 倍。我们的实验表明，SageAttention2++ 比 FlashAttention 提速 3.9 倍，同时保持与 SageAttention2 相同的注意力精度。这意味着 SageAttention2++ 有效地加速了各种模型，包括用于语言、图像和视频生成的模型，而端到端指标损失可以忽略不计。代码将发布在 https://github.com/thu-ml/SageAttention。
SageAttention2++ 相较于 FlashAttention 实现了 3.9 倍的速度提升，同时保持了与 SageAttention2 相同的注意力准确性。代码将在 https://github.com/thu-ml/SageAttention 开源。