SageAttention: Präzise 8-Bit Aufmerksamkeit für Plug-and-Play Inferenzbeschleunigung

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

October 3, 2024

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

cs.AI

Zusammenfassung

Die Transformer-Architektur dominiert über verschiedene Modelle hinweg. Als Herzstück des Transformers hat die Aufmerksamkeit eine Rechenkomplexität von O(N^2), im Vergleich zu O(N) für lineare Transformationen. Bei der Verarbeitung großer Sequenzlängen wird die Aufmerksamkeit zum hauptsächlichen zeitintensiven Bestandteil. Obwohl sich Quantisierung als effektive Methode zur Beschleunigung der Modellinferenz erwiesen hat, konzentrieren sich bestehende Quantisierungsmethoden hauptsächlich auf die Optimierung der linearen Schicht. Als Reaktion darauf analysieren wir zunächst ausführlich die Machbarkeit der Quantisierung in der Aufmerksamkeit. Anschließend schlagen wir SageAttention vor, eine äußerst effiziente und genaue Quantisierungsmethode für die Aufmerksamkeit. Die OPS (Operationen pro Sekunde) unseres Ansatzes übertrifft FlashAttention2 und xformers um etwa 2,1 bzw. 2,7 Mal. SageAttention erzielt auch eine überlegene Genauigkeitsleistung gegenüber FlashAttention3. Umfassende Experimente bestätigen, dass unser Ansatz fast keinen Verlust von End-to-End-Metriken bei verschiedenen Modellen verursacht, einschließlich solcher für die Verarbeitung großer Sprachen, Bildgenerierung und Videogenerierung.

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.