Technischer Bericht zu SageAttention2: Präzise 4-Bit-Aufmerksamkeit für beschleunigte Plug-and-Play-Inferenz.

SageAttention2 Technical Report: Accurate 4 Bit Attention for Plug-and-play Inference Acceleration

November 17, 2024

Autoren: Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, Jianfei Chen

cs.AI

Zusammenfassung

Obwohl die Quantisierung für lineare Schichten weit verbreitet ist, bleibt ihre Anwendung zur Beschleunigung des Aufmerksamkeitsprozesses begrenzt. SageAttention nutzt 8-Bit-Matrixmultiplikation, 16-Bit-Matrixmultiplikation mit 16-Bit-Akkumulator und präzisionssteigernde Methoden, um einen genauen und um den Faktor 2 beschleunigten Kernel im Vergleich zu FlashAttention2 zu implementieren. Um die Effizienz der Aufmerksamkeitsberechnung weiter zu steigern und gleichzeitig die Präzision beizubehalten, schlagen wir SageAttention2 vor, das signifikant schnellere 4-Bit-Matrixmultiplikation (Matmul) zusammen mit zusätzlichen präzisionssteigernden Techniken verwendet. Zunächst schlagen wir vor, Matrizen (Q, K) auf INT4-Ebene zu quantisieren und Matrizen (widetilde P, V) auf FP8-Ebene zu quantisieren. Zweitens schlagen wir eine Methode vor, um Q und V zu glätten und die Genauigkeit der Aufmerksamkeit mit INT4 QK und FP8 PV zu verbessern. Drittens analysieren wir die Quantisierungsgenauigkeit über Zeitstufen und Schichten und schlagen dann eine adaptive Quantisierungsmethode vor, um die End-to-End-Metriken über verschiedene Modelle sicherzustellen. Die Operationen pro Sekunde (OPS) von SageAttention2 übertreffen FlashAttention2 und xformers auf RTX4090 jeweils um etwa das 3- bzw. 5-fache. Umfassende Experimente bestätigen, dass unser Ansatz vernachlässigbare End-to-End-Metrikenverluste über verschiedene Modelle hinweg verursacht, einschließlich solcher für die Verarbeitung großer Sprachen, Bild- und Videogenerierung. Der Code ist verfügbar unter https://github.com/thu-ml/SageAttention.

English

Although quantization for linear layers has been widely used, its application to accelerate the attention process remains limited. SageAttention utilizes 8-bit matrix multiplication, 16-bit matrix multiplication with 16-bit accumulator, and precision-enhancing methods, implementing an accurate and 2x speedup kernel compared to FlashAttention2. To further enhance the efficiency of attention computation while maintaining precision, we propose SageAttention2, which utilizes significantly faster 4-bit matrix multiplication (Matmul) alongside additional precision-enhancing techniques. First, we propose to quantize matrixes (Q, K) to INT4 in a warp-level granularity and quantize matrixes (widetilde P, V) to FP8. Second, we propose a method to smooth Q and V, enhancing the accuracy of attention with INT4 QK and FP8 PV. Third, we analyze the quantization accuracy across timesteps and layers, then propose an adaptive quantization method to ensure the end-to-end metrics over various models. The operations per second (OPS) of SageAttention2 surpass FlashAttention2 and xformers by about 3x and 5x on RTX4090, respectively. Comprehensive experiments confirm that our approach incurs negligible end-to-end metrics loss across diverse models, including those for large language processing, image generation, and video generation. The codes are available at https://github.com/thu-ml/SageAttention.