SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866914068682506240 |
|---|---|
| author | Zhang, Jintao Wei, Jia Huang, Haofeng Zhang, Pengle Zhu, Jun Chen, Jianfei |
| author_facet | Zhang, Jintao Wei, Jia Huang, Haofeng Zhang, Pengle Zhu, Jun Chen, Jianfei |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2410_02367 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration Zhang, Jintao Wei, Jia Huang, Haofeng Zhang, Pengle Zhu, Jun Chen, Jianfei Machine Learning 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. |
| title | SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2410.02367 |