SageAttention: Attenzione precisa a 8 bit per l'accelerazione dell'inferenza plug-and-play

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

October 3, 2024

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

cs.AI

Abstract

L'architettura del transformer prevale tra vari modelli. Essendo il cuore del transformer, l'attenzione ha una complessità computazionale di O(N^2), rispetto a O(N) per le trasformazioni lineari. Quando si gestiscono lunghezze di sequenza elevate, l'attenzione diventa il componente principale che richiede più tempo. Sebbene la quantizzazione si sia dimostrata un metodo efficace per accelerare l'inferenza del modello, i metodi di quantizzazione esistenti si concentrano principalmente sull'ottimizzazione dello strato lineare. In risposta, analizziamo prima dettagliatamente la fattibilità della quantizzazione nell'attenzione. Successivamente, proponiamo SageAttention, un metodo di quantizzazione altamente efficiente e accurato per l'attenzione. Le OPS (operazioni al secondo) del nostro approccio superano di circa 2,1 volte FlashAttention2 e di 2,7 volte xformers. SageAttention ottiene anche prestazioni di precisione superiori rispetto a FlashAttention3. Esperimenti completi confermano che il nostro approccio comporta quasi nessuna perdita di metriche end-to-end tra vari modelli, inclusi quelli per l'elaborazione del linguaggio, la generazione di immagini e la generazione di video.

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.