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

TL;DR

SageAttention2 introduces a highly efficient attention mechanism using 4-bit quantization and precision enhancement techniques, significantly outperforming existing methods in speed while maintaining high accuracy across diverse models.

Contribution

It proposes novel quantization and smoothing techniques for attention computation, achieving faster performance with minimal accuracy loss.

Findings

OPS surpasses FlashAttention2 and xformers by 3x and 4.5x on RTX4090

Matches FlashAttention3 speed on Hopper GPUs with higher accuracy

Negligible end-to-end metrics loss across language, image, and video models

Abstract

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 matrices $(Q, K)$ to INT4 in a hardware-friendly thread-level granularity and quantize matrices $(\widetilde P, V)$ to FP8. Second, we propose a method to smooth $Q$, enhancing the accuracy of INT4 $QK^\top$. Third, we propose a two-level accumulation strategy for $\widetilde PV$ to enhance the accuracy of FP8 $\widetilde PV$. The operations per second (OPS) of SageAttention2 surpass FlashAttention2 and xformers by about 3x and 4.5x on RTX4090, respectively. Moreover, SageAttention2 matches the speed of FlashAttention3(fp8) on the Hopper GPUs, while delivering much higher accuracy. Comprehensive experiments confirm that our approach incurs negligible end-to-end metrics loss across diverse models, including those for language, image, and video generation. The code is available at https://github.com/thu-ml/SageAttention.