Technisch Rapport SageAttention2: Nauwkeurige 4-bits Aandacht voor Plug-and-play Inferentie Versnelling

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

November 17, 2024

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

cs.AI

Samenvatting

Hoewel kwantisatie voor lineaire lagen veel wordt toegepast, is de toepassing ervan om het aandachtsproces te versnellen beperkt gebleven. SageAttention maakt gebruik van 8-bits matrixvermenigvuldiging, 16-bits matrixvermenigvuldiging met 16-bits accumulator, en precisie-verhogende methoden, waarbij een nauwkeurige en 2x versnelde kernel wordt geïmplementeerd in vergelijking met FlashAttention2. Om de efficiëntie van aandachtsberekeningen verder te verbeteren terwijl de precisie behouden blijft, stellen we SageAttention2 voor, die aanzienlijk snellere 4-bits matrixvermenigvuldiging (Matmul) gebruikt naast extra precisie-verhogende technieken. Ten eerste stellen we voor om matrices (Q, K) te kwantiseren naar INT4 op warp-niveau en matrices (widetilde P, V) naar FP8 te kwantiseren. Ten tweede stellen we een methode voor om Q en V te verfijnen, wat de nauwkeurigheid van aandacht verbetert met INT4 QK en FP8 PV. Ten derde analyseren we de kwantisatienauwkeurigheid over tijdstappen en lagen, waarna we een adaptieve kwantisatiemethode voorstellen om de end-to-end metrieken over verschillende modellen te waarborgen. De bewerkingen per seconde (OPS) van SageAttention2 overtreffen FlashAttention2 en xformers met respectievelijk ongeveer 3x en 5x op RTX4090. Uitgebreide experimenten bevestigen dat onze aanpak verwaarloosbaar verlies aan end-to-end metrieken met zich meebrengt over diverse modellen, waaronder die voor grootschalige taalverwerking, beeldgeneratie en videogeneratie. De codes zijn beschikbaar op 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.