SageAttention: Atenção de 8 bits precisa para Aceleração de Inferência Plug-and-Play

SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration

October 3, 2024

Autores: Jintao Zhang, Jia wei, Pengle Zhang, Jun Zhu, Jianfei Chen

cs.AI

Resumo

A arquitetura transformer predomina em vários modelos. Como o cerne do transformer, a atenção possui uma complexidade computacional de O(N^2), comparada a O(N) para transformações lineares. Ao lidar com comprimentos de sequência longos, a atenção se torna o componente principal que consome tempo. Embora a quantização tenha se mostrado um método eficaz para acelerar a inferência do modelo, os métodos de quantização existentes focam principalmente na otimização da camada linear. Em resposta, primeiro analisamos detalhadamente a viabilidade da quantização na atenção. Em seguida, propomos o SageAttention, um método de quantização altamente eficiente e preciso para atenção. As OPS (operações por segundo) de nossa abordagem superam o FlashAttention2 e os xformers em cerca de 2,1 vezes e 2,7 vezes, respectivamente. O SageAttention também alcança um desempenho de precisão superior ao FlashAttention3. Experimentos abrangentes confirmam que nossa abordagem incorre em quase nenhuma perda de métricas de ponta a ponta em diversos modelos, incluindo aqueles para processamento de linguagem extensa, geração de imagens e geração de vídeos.

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.