SageAttention: Точное внимание 8 бит для ускорения подключения и использования.

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

October 3, 2024

Авторы: Jintao Zhang, Jia wei, Pengle Zhang, Jun Zhu, Jianfei Chen

cs.AI

Аннотация

Архитектура трансформера преобладает в различных моделях. В качестве сердца трансформера, внимание имеет вычислительную сложность O(N^2), по сравнению с O(N) для линейных преобразований. При работе с большими длинами последовательностей внимание становится основным компонентом, требующим времени. Хотя квантование доказало свою эффективность в ускорении вывода модели, существующие методы квантования в основном сосредоточены на оптимизации линейного слоя. В ответ на это мы первоначально анализируем возможность квантования внимания подробно. Затем мы предлагаем SageAttention - высокоэффективный и точный метод квантования для внимания. OPS (операции в секунду) нашего подхода превосходит FlashAttention2 и xformers примерно в 2,1 и 2,7 раза соответственно. SageAttention также достигает более высокой точности по сравнению с FlashAttention3. Обширные эксперименты подтверждают, что наш подход практически не приводит к потере метрик от начала до конца через различные модели, включая модели для обработки больших языков, генерации изображений и генерации видео.

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.