arXiv:2411.10958v4 宣告类型: replace-cross 摘要：尽管量化已经广泛应用于线性层，但其在加速注意力过程中的应用仍然有限。为了在不牺牲精度的前提下进一步提高注意力计算的效率，同时比SageAttention更快，我们提出了SageAttention2，它利用了显著更快的4位矩阵乘法（Matmul）以及额外的精度增强技术。首先，我们提出将矩阵$(Q, K)$按硬件友好的线程级粒度量化为INT4，将矩阵$(\widetilde{P}, V)$量化为FP8。其次，我们提出了一种平滑$Q$的方法，以增强INT4 $QK^\top$的准确性。第三，我们提出了$\widetilde{P}V$的两级累加策略，以增强FP8 $\widetilde{P}V$的准确性。SageAttention2在RTX4090上的每秒操作数（OPS）分别超过FlashAttention2和xformers约3倍和4.5倍。此外，SageAttention2在Hopper GPU上的速度与FlashAttention3(fp8)相当，但在准确性方面要高得多。全面的实验表明，我们的方法在不同模型中（包括语言、图像和视频生成模型）引起的端到端指标损失可以忽略不计。我们的代码可在 https://github.com/thu-ml/SageAttention 获取。