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

초록

Transformer 아키텍처가 다양한 모델에서 우세합니다. Transformer의 핵심인 어텐션은 선형 변환의 O(N)에 비해 O(N^2)의 계산 복잡도를 가지고 있습니다. 큰 시퀀스 길이를 처리할 때 어텐션이 주요한 시간 소모 구성 요소가 됩니다. 양자화는 모델 추론 가속화를 위한 효과적인 방법으로 입증되었지만, 기존의 양자화 방법은 주로 선형 레이어의 최적화에 초점을 맞추고 있습니다. 이에 우리는 먼저 어텐션에서 양자화의 실행 가능성을 상세히 분석합니다. 그 후에 우리는 SageAttention이라는 어텐션을 위한 매우 효율적이고 정확한 양자화 방법을 제안합니다. 우리 방법의 OPS(초당 연산 횟수)는 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.