8-bit量化让注意力层提速两倍多，精度几乎不降。
SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration
- 设计8比特量化方法，精准优化注意力计算
- 速度比FlashAttention2快2.1倍，比xformers快2.7倍
- 适合需要高速推理的大模型应用，如文本、图像生成
Transformer架构在各类模型中占主导地位。作为其核心的注意力机制具有$O(N^2)$的计算复杂度，远高于线性变换的$O(N)$。在处理长序列时，注意力成为主要耗时环节。尽管量化已被证明是加速推理的有效手段，但现有方法主要聚焦于优化线性层。为此，我们首次深入分析了注意力量化可行性，并提出SageAttention——一种高效且高精度的注意力量化方法。该方法在性能上相比FlashAttention2和xformers分别提升约2.1倍和2.7倍。SageAttention在精度上优于FlashAttention3。全面实验表明，该方法在多种模型（包括大语言模型、图像生成与视频生成）上几乎无端到端指标损失。代码已开源：https://github.com/thu-ml/SageAttention。
原文摘要 · Abstract (English)
The transformer architecture predominates across various models. As the heart of the transformer, attention has a computational complexity of $O(N^2)$, compared to $O(N)$ for linear transformations. When handling large sequence lengths, attention becomes the primary time-consuming component. Although quantization has proven to be an effective method for accelerating model inference, existing quantization methods primarily focus on optimizing the linear layer. In response, we first analyze the feasibility of quantization in attention detailedly. Following that, we propose SageAttention, a highly efficient and accurate quantization method for attention. The OPS (operations per second) of our approach outperforms FlashAttention2 and xformers by about 2.1 times and 2.7 times, respectively. SageAttention also achieves superior accuracy performance over FlashAttention3. Comprehensive experiments confirm that our approach incurs almost no end-to-end metrics loss across diverse models, including those for large language processing, image generation, and video generation. The codes are available at https://github.com/thu-ml/SageAttention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。