Технический отчет SageAttention2: Точное внимание на 4 бита для ускорения вывода "включи и играй".

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

November 17, 2024

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

cs.AI

Аннотация

Хотя квантование для линейных слоев широко используется, его применение для ускорения процесса внимания остается ограниченным. SageAttention использует умножение матриц 8-битным способом, умножение матриц 16 бит с 16-битным аккумулятором и методы увеличения точности, реализуя точное и ускоренное ядро в 2 раза по сравнению с FlashAttention2. Для дальнейшего увеличения эффективности вычислений внимания при сохранении точности мы предлагаем SageAttention2, который использует значительно более быстрое умножение матриц 4 бит (Matmul) вместе с дополнительными методами увеличения точности. Во-первых, мы предлагаем квантовать матрицы (Q, K) в INT4 с гранулярностью на уровне warp и квантовать матрицы (widetilde P, V) в FP8. Во-вторых, мы предлагаем метод сглаживания Q и V, улучшая точность внимания с INT4 QK и FP8 PV. В-третьих, мы анализируем точность квантования по временным шагам и слоям, затем предлагаем адаптивный метод квантования для обеспечения метрик от начала до конца по различным моделям. Операции в секунду (OPS) SageAttention2 превосходят FlashAttention2 и xformers примерно в 3 и 5 раз соответственно на RTX4090. Обширные эксперименты подтверждают, что наш подход вызывает незначительные потери метрик от начала до конца по различным моделям, включая модели для обработки больших объемов языковых данных, генерации изображений и видео. Коды доступны по адресу 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.