Rapport technique SageAttention2 : Attention précise sur 4 bits pour accélération de l'inférence plug-and-play.

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

Résumé

Bien que la quantification pour les couches linéaires soit largement utilisée, son application pour accélérer le processus d'attention reste limitée. SageAttention utilise une multiplication de matrices sur 8 bits, une multiplication de matrices sur 16 bits avec un accumulateur sur 16 bits, et des méthodes d'amélioration de la précision, mettant en œuvre un noyau précis et deux fois plus rapide par rapport à FlashAttention2. Pour améliorer davantage l'efficacité du calcul d'attention tout en maintenant la précision, nous proposons SageAttention2, qui utilise une multiplication de matrices sur 4 bits (Matmul) significativement plus rapide aux côtés de techniques supplémentaires d'amélioration de la précision. Tout d'abord, nous proposons de quantifier les matrices (Q, K) en INT4 avec une granularité au niveau de la vague et de quantifier les matrices (widetilde P, V) en FP8. Deuxièmement, nous proposons une méthode pour lisser Q et V, améliorant la précision de l'attention avec INT4 QK et FP8 PV. Troisièmement, nous analysons la précision de la quantification à travers les pas de temps et les couches, puis proposons une méthode de quantification adaptative pour garantir les métriques de bout en bout sur divers modèles. Les opérations par seconde (OPS) de SageAttention2 surpassent FlashAttention2 et xformers d'environ 3x et 5x sur RTX4090, respectivement. Des expériences approfondies confirment que notre approche entraîne une perte négligeable de métriques de bout en bout à travers divers modèles, y compris ceux pour le traitement de langage, la génération d'images et la génération de vidéos. Les codes sont disponibles sur 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.