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

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.

한 줄 요약

SageAttention은 8비트 정밀도로 주의 메커니즘을 정확하게 양자화하여, 다양한 모델에서 빠른 추론 속도와 거의 무시할 수 있는 정확도 손실을 달성한다.

핵심 기여도

SageAttention은 주의 메커니즘의 8비트 양자화를 통해 FlashAttention2 대비 약 2.1배, xformers 대비 2.7배 빠른 연산 성능을 달성.
기존 양자화 연구는 주로 선형층에 집중했지만, SageAttention은 주의 메커니즘에 특화된 정확도 유지 전략을 제안.
K 행렬의 채널별 이상치 문제를 해결하기 위한 스무딩 기법을 도입, 정확도 손실 0.2% 미만으로 유지.
(P, V) 행렬 연산에서 FP16 누산기를 사용해 정확도와 속도를 동시에 향상.

핵심 아이디어

기존 양자화 연구는 주로 선형층에만 적용되었으며, 주의 메커니즘은 고정밀도 연산으로 남아 있어 성능 향상이 제한적이었다. SageAttention은 주의 메커니즘 내 Q, K, P, V 행렬을 8비트 정밀도로 정확하게 양자화하는 방법을 제안한다. 특히, K 행렬의 채널별 이상치 문제를 해결하기 위해 스무딩 기법을 도입하고, (P, V) 행렬의 Matmul 연산에서는 FP16 누산기를 사용해 정확도를 유지하면서도 속도를 향상시킨다. 이는 기존 FP8 기반 방식(예: FlashAttention3)보다 더 높은 정확도와 유사한 성능을 제공한다.

기술적 접근법

**양자화 대상**: Q, K, P, V 행렬 중 Q와 K는 8비트로, P와 V는 FP16으로 유지.
**K 행렬 스무딩**: 채널별 이상치를 감소시켜 정확도 손실을 0.2% 미만으로 유지.
**(P, V) Matmul 최적화**: FP16 누산기를 사용하여 정확도 유지와 동시에 연산 속도 향상.
**Triton 기반 커널 구현**: RTX 4090 및 3090 GPU에서 INT8 mma(u8.u8.s32) 및 FP16 mma(f16.f16.f16) 명령어 활용.
**적응형 양자화 전략**: 각 레이어별 최적의 주의 구현 선택으로 성능 향상 (11.7% 속도 향상).

주요 결과

SageAttention은 FlashAttention2 대비 약 2.1배, xformers 대비 2.7배 빠른 연산 성능을 달성.
Llama2 모델에서 MMLU 데이터셋에서 25.5% 정확도를 달성 (INT8 주의 사용 시), 기존 FP8 기반 방식보다 높은 정확도.
Unidiffuser 모델에서 INT8 및 FlashAttention3의 FP8 구현은 완전히 흐릿한 이미지 생성, SageAttention은 정확도 유지.
RTX 4090에서 headdim=64일 때 340 TOPS 성능 달성, FlashAttention2 대비 2배 이상 빠름.

의의 및 한계

SageAttention은 주의 메커니즘의 8비트 양자화를 통해 높은 정확도와 빠른 추론 속도를 동시에 달성하며, 기존 연구에서 간과되었던 주의 메커니즘 최적화의 가능성을 열었다. 특히, 다양한 모델(언어, 이미지, 영상 생성)에서 plug-and-play 방식으로 적용 가능하다는 점에서 실용적 가치가 높다. 그러나 현재는 RTX 4090 및 3090 GPU에 최적화되어 있으며, 다른 하드웨어 아키텍처에서의 성능은 추가 연구가 필요하다. 또한, 스무딩 기법의 효과는 특정 모델에 따라 달라질 수 있으므로, 보편적인 적용 가능성도 검증이 필요하다.

실용적 활용

SageAttention은 대규모 언어 모델, 이미지 생성 모델, 영상 생성 모델 등에서 추론 속도를 2배 이상 향상시키며 정확도를 유지할 수 있어, 클라우드 인프라, 모바일 기기, 실시간 생성 시스템 등에 적용 가능하다. 특히, 기존 모델을 수정하지 않고도 plug-and-play 방식으로 사용할 수 있어, 빠른 도입과 배포가 가능하다.