SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration (2410.02367v8)

Abstract: 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.

• The paper introduces SageAttention as an 8-bit quantization method that accelerates transformer attention while preserving output accuracy.
• It optimizes Q, K, P, and V matrices using INT8 multiplication, smoothing techniques, and FP16 accumulators for minimal precision loss.
• Experimental results demonstrate speed improvements of 2.1x–2.7x and an average throughput of 340 TOPS on the RTX 4090.

SageAttention: Accurate 8-bit Attention for Plug-and-Play Inference Acceleration

The paper introduces SageAttention, a method designed to efficiently accelerate the attention mechanism within transformer models through 8-bit quantization, without sacrificing accuracy. As the core of transformer architectures, attention processes with a computational complexity of $O(N^2)$, often becoming a bottleneck in scenarios with long sequence lengths. Despite the existence of numerous quantization methods for model acceleration, they primarily focus on linear layers, leaving attention largely unoptimized. SageAttention seeks to address this gap through a quantization approach that not only improves processing speed but also ensures minimal degradation in output quality.

Methodology

SageAttention employs several strategies to achieve efficient quantization:

1. Quantization of Attention Components: Unlike previous methods that apply quantization primarily to linear layers, SageAttention quantizes $Q, K, P, V$ matrices directly in attention. It formulates attention using dynamic quantization, allowing for the plug-and-play integration into existing models without further training.
2. INT8 Data Type: The choice to use INT8 for matrix multiplication derives from its superior accuracy over FP8 formats and its throughput advantages in GPUs such as the RTX 4090. This decision is bolstered by evidence showing that INT8 delivers higher precision in the computation of attention matrices compared to alternative low-precision representations.
3. Smoothing $K$: A key challenge in quantizing attention is the presence of channel-wise outliers in matrix $K$. SageAttention applies a smoothing technique, subtracting the mean from $K$, to mitigate these outliers without impacting the final attention scores. This technique improves accuracy with negligible computational overhead.
4. Use of FP16 Accumulators: To maintain high precision in the presence of quantization, SageAttention retains $P$ and $V$ in FP16 and uses FP16 accumulators for their multiplication. This approach significantly enhances accuracy while still delivering a twofold speed improvement over traditional FP32 accumulations.
5. Adaptive Quantization Strategy: By implementing varying levels of quantization granularity and evaluating on multiple input scenarios, SageAttention adaptively selects the optimal quantization strategy to maximize computational speed while maintaining accuracy across different model layers.

Experimental Evaluation

The empirical analysis demonstrates that SageAttention achieves substantial speed improvements without compromising performance metrics. Key outcomes include:

• SageAttention outperforms existing methods such as FlashAttention2 and xformers by approximately 2.1x and 2.7x, respectively.
• Across diverse applications, including text, image, and video generation models, SageAttention maintains end-to-end accuracy metrics comparable to full-precision attention.
• Real-world evaluations reveal an average throughput of 340 TOPS on the RTX 4090, suggesting that SageAttention approaches the theoretical limits of INT8 processing efficiency.

Implications and Future Directions

The development of SageAttention carries significant implications for practitioners and researchers looking to optimize deployment of transformer models in resource-constrained environments. The combination of speed and accuracy enhancements means that SageAttention can be incorporated into a wide range of applications, facilitating more efficient large-scale deployments.

Looking forward, the paper suggests future work to adapt SageAttention for newer architectures, such as Nvidia's Hopper, potentially unlocking further efficiency gains. This continued evolution and expansion could redefine performance benchmarks for transformer-based applications.

SageAttention sets a robust precedent for achieving high performance through quantization and could inspire further research into quantizing other computational components within neural networks—thereby offering potential for even broader impacts in the field.