SageAttention2：通过彻底的离群值平滑和逐线程 INT4 量化实现高效注意力
摘要
尽管线性层的量化已被广泛使用，但其在加速注意力过程方面的应用仍然有限。为了在保持精度的同时进一步提高注意力计算相对于 SageAttention 的效率，我们提出了 SageAttention2，它利用了显著更快的 4 位矩阵乘法以及额外的精度增强技术。首先，我们提出以硬件友好的线程级粒度将矩阵 量化为 INT4，并将矩阵 量化为 FP8。其次，我们提出了一种平滑 的方法，以提高 INT4 的准确性。第三，我们提出了一种用于 的两级累加策略，以提高 FP8 的准确性。在 RTX4090 上，SageAttention2 的每秒操作数（OPS）分别比 FlashAttention2 和 xformers 高出约 3 倍和 4.5 倍。此外，SageAttention2 在 Hopper GPU 上匹配了 FlashAttention3(fp8) 的速度，同时提供了更高的精度。综合实验证实，我们的方法在各种模型中（包括用于语言、图像和视频生成的模型）造成的端到端指标损失可以忽略不计。代码可在 https://github.com/thu-ml/SageAttention 获取。
关键词
引用
@article{arxiv.2411.10958,
title = {SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization},
author = {Jintao Zhang and Haofeng Huang and Pengle Zhang and Jia Wei and Jun Zhu and Jianfei Chen},
journal= {arXiv preprint arXiv:2411.10958},
year = {2025}
}
备注
@inproceedings{zhang2024sageattention2, title={Sageattention2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization}, author={Zhang, Jintao and Huang, Haofeng and Zhang, Pengle and Wei, Jia and Zhu, Jun and Chen, Jianfei}, booktitle={International Conference on Machine Learning (ICML)}, year={2025} }