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
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.