Relatório Técnico SageAttention2: Atenção Precisa de 4 Bits para Aceleração de Inferência Plug-and-Play

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

November 17, 2024

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

cs.AI

Resumo

Embora a quantização para camadas lineares seja amplamente utilizada, sua aplicação para acelerar o processo de atenção ainda é limitada. O SageAttention utiliza multiplicação de matriz de 8 bits, multiplicação de matriz de 16 bits com acumulador de 16 bits e métodos de aprimoramento de precisão, implementando um kernel preciso e com aumento de velocidade de 2x em comparação com o FlashAttention2. Para aprimorar ainda mais a eficiência do cálculo de atenção mantendo a precisão, propomos o SageAttention2, que utiliza uma multiplicação de matriz de 4 bits (Matmul) significativamente mais rápida juntamente com técnicas adicionais de aprimoramento de precisão. Primeiramente, propomos quantizar as matrizes (Q, K) para INT4 em uma granularidade de nível de warp e quantizar as matrizes (widetilde P, V) para FP8. Em segundo lugar, propomos um método para suavizar Q e V, aprimorando a precisão da atenção com INT4 QK e FP8 PV. Em terceiro lugar, analisamos a precisão da quantização ao longo dos passos de tempo e camadas, e propomos um método de quantização adaptativa para garantir as métricas de ponta a ponta em vários modelos. As operações por segundo (OPS) do SageAttention2 superam o FlashAttention2 e os xformers em cerca de 3x e 5x no RTX4090, respectivamente. Experimentos abrangentes confirmam que nossa abordagem incorre em perda negligenciável de métricas de ponta a ponta em diversos modelos, incluindo aqueles para processamento de linguagem extensa, geração de imagens e geração de vídeos. Os códigos estão disponíveis em 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.