SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization

要約

線形層の量子化は広く使用されていますが、アテンションプロセスを加速するためのその応用は依然として限られています。
精度を維持しながら SageAttend と比較してアテンション計算の効率をさらに高めるために、追加の精度向上技術と並行して大幅に高速な 4 ビット行列乗算 (Matmul) を利用する SageAttendant2 を提案します。
まず、ハードウェアに適したスレッドレベルの粒度で行列 $(Q, K)$ を INT4 に量子化し、行列 $(\widetilde P, V)$ を FP8 に量子化することを提案します。
次に、$Q$ を平滑化し、INT4 $QK$ の精度を高める方法を提案します。
第三に、FP8 $\widetilde PV$ の精度を高めるために $PV$ に FP32 Matmul バッファを使用することを提案します。
SageAttendant2 の 1 秒あたりの操作数 (OPS) は、RTX4090 上で FlashAttendant2 と xformers をそれぞれ約 3 倍と 5 倍上回っています。
包括的な実験により、私たちのアプローチでは、大規模な言語処理、画像生成、ビデオ生成などのさまざまなモデルにわたって、エンドツーエンドのメトリクス損失が無視できる程度であることが確認されています。
コードは https://github.com/thu-ml/SageAttend で入手できます。

要約(オリジナル)

Although quantization for linear layers has been widely used, its application to accelerate the attention process remains limited. To further enhance the efficiency of attention computation compared to SageAttention 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 hardware-friendly thread-level granularity and quantize matrixes $(\widetilde P, V)$ to FP8. Second, we propose a method to smooth $Q$, enhancing the accuracy of INT4 $QK$. Third, we propose to use an FP32 Matmul buffer for $PV$ to enhance the accuracy of FP8 $\widetilde PV$. 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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google