Rapporto Tecnico SageAttention2: Attivazione Precisa a 4 Bit per Accelerazione dell'Inferenza Plug-and-Play

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

November 17, 2024

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

cs.AI

Abstract

Sebbene la quantizzazione per strati lineari sia stata ampiamente utilizzata, la sua applicazione per accelerare il processo di attenzione rimane limitata. SageAttention utilizza moltiplicazioni di matrici a 8 bit, moltiplicazioni di matrici a 16 bit con accumulatore a 16 bit e metodi di miglioramento della precisione, implementando un kernel accurato e con un aumento della velocità del 2x rispetto a FlashAttention2. Per migliorare ulteriormente l'efficienza del calcolo dell'attenzione mantenendo la precisione, proponiamo SageAttention2, che utilizza moltiplicazioni di matrici a 4 bit significativamente più veloci insieme a tecniche aggiuntive di miglioramento della precisione. In primo luogo, proponiamo di quantizzare le matrici (Q, K) in INT4 con una granularità a livello di warp e quantizzare le matrici (widetilde P, V) in FP8. In secondo luogo, proponiamo un metodo per levigare Q e V, migliorando la precisione dell'attenzione con INT4 QK e FP8 PV. In terzo luogo, analizziamo l'accuratezza della quantizzazione attraverso i passaggi temporali e i livelli, quindi proponiamo un metodo di quantizzazione adattiva per garantire le metriche end-to-end su vari modelli. Le operazioni al secondo (OPS) di SageAttention2 superano FlashAttention2 e xformers di circa 3x e 5x su RTX4090, rispettivamente. Esperimenti esaustivi confermano che il nostro approccio comporta una perdita trascurabile delle metriche end-to-end su modelli diversi, inclusi quelli per l'elaborazione del linguaggio, la generazione di immagini e la generazione di video. I codici sono disponibili su 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.