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)に比べてO(N^2)の計算量を持ちます。大規模なシーケンス長を扱う際、アテンションは主要な時間を要する要素となります。量子化はモデルの推論を加速する効果的な手法であることが証明されていますが、既存の量子化手法は主に線形層の最適化に焦点を当てています。このため、我々はまず、アテンションにおける量子化の実現可能性を詳細に分析します。その後、我々はSageAttentionという、アテンションのための高効率かつ正確な量子化手法を提案します。当該手法のOPS（1秒あたりの演算回数）は、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.