SageAttention2通过四比特量化和去噪优化，实现更快更准的注意力计算。
SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
- 在线程级粒度对Q、K进行INT4量化，V用FP8，提升计算效率。
- 通过平滑Q矩阵，显著提高INT4 QKᵀ的精度，减少误差累积。
- 采用双层累加策略处理VP，保障FP8计算的准确性，适合多模态生成任务。
尽管线性层的量化已被广泛使用，但其在加速注意力计算方面的应用仍有限。为在保持精度的同时进一步提升注意力计算效率，我们提出SageAttention2，结合更快的4比特矩阵乘法（Matmul）及多项精度增强技术。首先，在硬件友好的线程级粒度上将矩阵(Q, K)量化为INT4，将(˜P, V)量化为FP8；其次，提出一种平滑Q的方法，提升INT4 QKᵀ的准确性；第三，设计两层累加策略处理˜PV，增强FP8 ˜PV的精度。在RTX4090上，SageAttention2的每秒操作数（OPS）分别比FlashAttention2和xformers高出约3倍和4.5倍。此外，其速度接近Hopper GPU上的FlashAttention3(fp8)，同时提供更高精度。大量实验表明，该方法在语言、图像与视频生成等多样模型中仅带来可忽略的端到端指标损失。代码已开源：https://github.com/thu-ml/SageAttention。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。