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ビット累算器、精度向上手法を活用し、FlashAttention2 に比べて正確で2倍の高速化カーネルを実装しています。精度を維持しつつ注意計算の効率をさらに向上させるために、SageAttention2 を提案します。これは、より高速な4ビット行列乗算（Matmul）と追加の精度向上技術を活用しています。まず、行列（Q、K）をワープレベルの粒度で INT4 に量子化し、行列（widetilde P、V）を FP8 に量子化することを提案します。次に、Q と V を滑らかにする方法を提案し、INT4 QK と FP8 PV による注意の精度を向上させます。さらに、各タイムステップとレイヤーでの量子化精度を分析し、さまざまなモデルにわたるエンドツーエンドのメトリクスを確保するための適応型量子化手法を提案します。SageAttention2 の1秒あたりの演算回数（OPS）は、RTX4090 上で FlashAttention2 や xformers をそれぞれ約3倍、5倍上回っています。包括的な実験により、大規模言語処理、画像生成、ビデオ生成など多様なモデルにわたって、我々の手法がほとんどエンドツーエンドのメトリクス損失をもたらさないことが確認されました。コードは 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.