Informe Técnico SageAttention2: Atención precisa de 4 bits para aceleración de inferencia 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

Resumen

Aunque la cuantificación para capas lineales se ha utilizado ampliamente, su aplicación para acelerar el proceso de atención sigue siendo limitada. SageAttention utiliza multiplicación de matrices de 8 bits, multiplicación de matrices de 16 bits con acumulador de 16 bits y métodos que mejoran la precisión, implementando un núcleo preciso y con una aceleración 2 veces mayor en comparación con FlashAttention2. Para mejorar aún más la eficiencia del cálculo de atención manteniendo la precisión, proponemos SageAttention2, que utiliza una multiplicación de matrices de 4 bits significativamente más rápida junto con técnicas adicionales que mejoran la precisión. En primer lugar, proponemos cuantificar las matrices (Q, K) a INT4 en una granularidad a nivel de warp y cuantificar las matrices (widetilde P, V) a FP8. En segundo lugar, proponemos un método para suavizar Q y V, mejorando la precisión de la atención con INT4 QK y FP8 PV. En tercer lugar, analizamos la precisión de la cuantificación a lo largo de pasos de tiempo y capas, y luego proponemos un método de cuantificación adaptativa para garantizar las métricas de extremo a extremo en diversos modelos. Las operaciones por segundo (OPS) de SageAttention2 superan a FlashAttention2 y xformers en aproximadamente 3 veces y 5 veces en RTX4090, respectivamente. Experimentos exhaustivos confirman que nuestro enfoque incurre en una pérdida de métricas de extremo a extremo negligente en diversos modelos, incluidos los de procesamiento de lenguaje grande, generación de imágenes y generación de video. Los códigos están disponibles en 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.