SageAttention : Attention précise sur 8 bits pour l'accélération de l'inférence plug-and-play

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

Résumé

L'architecture transformer prédomine à travers divers modèles. En tant que cœur du transformer, l'attention a une complexité computationnelle de O(N^2), comparée à O(N) pour les transformations linéaires. Lorsqu'il s'agit de traiter de longues séquences, l'attention devient le composant principal consommateur de temps. Bien que la quantification se soit avérée être une méthode efficace pour accélérer l'inférence des modèles, les méthodes de quantification existantes se concentrent principalement sur l'optimisation de la couche linéaire. En réponse, nous analysons d'abord en détail la faisabilité de la quantification dans l'attention. Ensuite, nous proposons SageAttention, une méthode de quantification hautement efficace et précise pour l'attention. Le nombre d'opérations par seconde (OPS) de notre approche surpasse FlashAttention2 et xformers d'environ 2,1 fois et 2,7 fois respectivement. SageAttention atteint également des performances de précision supérieures à FlashAttention3. Des expériences approfondies confirment que notre approche n'entraîne presque aucune perte de métriques de bout en bout à travers divers modèles, y compris ceux pour le traitement de langage, la génération d'images et la génération de vidéos.

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.