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

要約

トランスアーキテクチャは、さまざまなモデルで優勢です。
トランスの心臓の心臓は、線形変換の場合、$ o（n）$と比較して、$ o（n^2）$の計算の複雑さを持っています。
大きなシーケンスの長さを処理するとき、注意が主要な時間式コンポーネントになります。
量子化はモデル推論を加速するための効果的な方法であることが証明されていますが、既存の量子化方法は主に線形層の最適化に焦点を当てています。
これに応じて、最初に注意の量子化の実現可能性を詳細に分析します。
それに続いて、私たちは、注意のための非常に効率的で正確な量子化方法であるSageattentionを提案します。
アプローチのOPS（1秒あたりの操作）は、Flashattention2とXFormersをそれぞれ約2.1倍および2.7倍上回ります。
Sageattentionはまた、フラッシュ放送よりも優れた精度パフォーマンスを達成します3。
包括的な実験では、私たちのアプローチが、大規模な言語処理、画像生成、ビデオ生成を含む多様なモデル全体で、エンドツーエンドのメトリックの損失をほとんど発生しないことを確認しています。
コードはhttps://github.com/thu-ml/sageattentionで入手できます。

要約(オリジナル)

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. The codes are available at https://github.com/thu-ml/SageAttention.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google