SageAttention: Nauwkeurige 8-bits Aandacht voor Plug-and-play Inferentie Versnelling

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

October 3, 2024

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

cs.AI

Samenvatting

De transformer-architectuur domineert over verschillende modellen. Als het hart van de transformer heeft aandacht een computationele complexiteit van O(N^2), in vergelijking met O(N) voor lineaire transformaties. Bij het verwerken van lange sequentielengtes wordt aandacht de belangrijkste tijdrovende component. Hoewel kwantisatie effectief is gebleken voor het versnellen van modelinferentie, richten bestaande kwantisatiemethoden zich voornamelijk op het optimaliseren van de lineaire laag. Als reactie daarop analyseren we eerst gedetailleerd de haalbaarheid van kwantisatie in aandacht. Vervolgens stellen we SageAttention voor, een zeer efficiënte en nauwkeurige kwantisatiemethode voor aandacht. De OPS (bewerkingen per seconde) van onze aanpak presteert beter dan FlashAttention2 en xformers met respectievelijk ongeveer 2,1 keer en 2,7 keer. SageAttention behaalt ook superieure nauwkeurigheidsprestaties ten opzichte van FlashAttention3. Uitgebreide experimenten bevestigen dat onze aanpak vrijwel geen verlies van end-to-end metrieken met zich meebrengt over diverse modellen, waaronder die voor grote taalverwerking, beeldgeneratie en videogeneratie.

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.