SageAttention: Atención precisa de 8 bits para aceleración de inferencia listo para usar

SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration

October 3, 2024

Autores: Jintao Zhang, Jia wei, Pengle Zhang, Jun Zhu, Jianfei Chen

cs.AI

Resumen

La arquitectura del transformador predomina en varios modelos. Como el núcleo del transformador, la atención tiene una complejidad computacional de O(N^2), en comparación con O(N) para transformaciones lineales. Al manejar longitudes de secuencia largas, la atención se convierte en el componente principal que consume tiempo. Aunque la cuantificación ha demostrado ser un método efectivo para acelerar la inferencia del modelo, los métodos de cuantificación existentes se centran principalmente en optimizar la capa lineal. En respuesta, primero analizamos detalladamente la viabilidad de la cuantificación en la atención. A continuación, proponemos SageAttention, un método de cuantificación altamente eficiente y preciso para la atención. Las OPS (operaciones por segundo) de nuestro enfoque superan a FlashAttention2 y xformers aproximadamente 2.1 veces y 2.7 veces, respectivamente. SageAttention también logra un rendimiento de precisión superior a FlashAttention3. Experimentos exhaustivos confirman que nuestro enfoque incurre en casi ninguna pérdida de métricas de extremo a extremo en diversos modelos, incluidos los de procesamiento de lenguaje grande, generación de imágenes y generación de video.

English

The transformer architecture predominates across various models. As the heart of the transformer, attention has a computational complexity of O(N^2), compared to O(N) for linear transformations. When handling large sequence lengths, attention becomes the primary time-consuming component. Although quantization has proven to be an effective method for accelerating model inference, existing quantization methods primarily focus on optimizing the linear layer. In response, we first analyze the feasibility of quantization in attention detailedly. Following that, we propose SageAttention, a highly efficient and accurate quantization method for attention. The OPS (operations per second) of our approach outperforms FlashAttention2 and xformers by about 2.1 times and 2.7 times, respectively. SageAttention also achieves superior accuracy performance over FlashAttention3. Comprehensive experiments confirm that our approach incurs almost no end-to-end metrics loss across diverse models, including those for large language processing, image generation, and video generation.