SageAttention：用于即插即用推理加速的准确 8 位注意力机制
机器学习
2025-10-02 v9
摘要
Transformer 架构在 various 模型中占据主导地位。作为 Transformer 的核心，注意力机制的计算复杂度为 ，而线性变换的复杂度为 。在处理大序列长度时，注意力成为主要的计算瓶颈。虽然量化已被证明是加速模型推理的有效方法，但现有量化方法主要聚焦于优化线性层。为此，我们首次详细分析了注意力中的量化可行性。随后，我们提出了 SageAttention，这是一种高效且准确的注意力量化方法。我们的方法的 OPS（每秒运算次数）分别比 FlashAttention2 和 xformers 高出约 2.1 倍和 2.7 倍。SageAttention 还在 FlashAttention3 方面实现了优越的准确性性能。全面的实验表明，我们的方法在包括大语言模型、图像生成和视频生成等 various 模型中几乎不会引入端到端指标损失。代码已发布于 https://github.com/thu-ml/SageAttention。
引用
@article{arxiv.2410.02367,
title = {SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration},
author = {Jintao Zhang and Jia Wei and Haofeng Huang and Pengle Zhang and Jun Zhu and Jianfei Chen},
journal= {arXiv preprint arXiv:2410.02367},
year = {2025}
}
备注
@inproceedings{zhang2025sageattention, title={SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration}, author={Zhang, Jintao and Wei, Jia and Zhang, Pengle and Zhu, Jun and Chen, Jianfei}, booktitle={International Conference on Learning Representations (ICLR)}, year={2025} }