arXiv:2505.21136cs.LGcs.AI2025-05被引 28

用更快速的FP8矩阵乘法,让注意力机制快3.9倍。

SageAttention2++: A More Efficient Implementation of SageAttention2

  • 改用FP8 Matmul在FP16中累加,提升计算速度。
  • 相比FlashAttention提速3.9倍,精度与SageAttention2一致。
  • 适合需要高效推理的语言、图像、视频生成模型。

注意力机制的效率至关重要,因其时间复杂度随序列长度呈平方增长。SageAttention2通过量化加速注意力中的矩阵乘法(Matmul)。为进一步加速,本文提出利用更快的FP8 Matmul在FP16中累加的指令,该指令比SageAttention2中使用的FP8 Matmul快2倍。实验表明,SageAttention2++相比FlashAttention实现3.9倍加速,同时保持与SageAttention2相同的注意力精度。这意味着SageAttention2++能有效加速各类模型,包括语言、图像和视频生成模型,且端到端性能损失可忽略。代码将开源于https://github.com/thu-ml/SageAttention。

原文摘要 · Abstract (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.

注意力机制加速计算量化推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。