Rapporto Tecnico SageAttention2: Attivazione Precisa a 4 Bit per
Accelerazione dell'Inferenza Plug-and-PlaySageAttention2 Technical Report: Accurate 4 Bit Attention for
Plug-and-play Inference Acceleration
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.