SageAttention2
“SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization”由Jintao Zhang等人撰写。文章提出SageAttention2，通过线程级INT4量化、Q矩阵平滑、两级累加策略等技术，在提升注意力计算效率的同时保持精度，在多种模型上取得优异性能。
简介
- 研究背景：随着序列长度增加，注意力机制的二次时间复杂度使其高效实现变得关键。现有优化方法各有局限，如线性和稀疏注意力方法适用范围有限，常用的注意力方法如FlashAttention、xformers和SageAttention等虽有不错性能，但SageAttention存在INT8矩阵乘法速度慢和特定GPU加速受限的问题。
- 相关工作
- FlashAttention：将注意力计算中的矩阵按token维度分块并行计算，降低计算复杂度，提升计算效率。
- 量化：通过将高精度矩阵转换为低精度格式加速矩阵乘法，不同量化器在数值格式和粒度上有差异。
- SageAttention：基于FlashAttention的分块策略，将Q、K量化为INT8，对K进行预处理以保持精度，对(\tilde{P})、V使用FP16并降低累加器精度加速计算，但存在局限性。
- SageAttention2方法
- 平滑Q：由于INT4数值范围有限，存在异常值影响量化精度。通过减去Q每个块的均值平滑Q，结合对K的平滑，将(QK^{\top})计算分解，分预处理和注意力两个阶段，提升INT4量化精度。
- INT4线程级量化：在SageAttention的基础上，提出线程级量化，根据GPU线程和矩阵内存布局，以更细粒度进行量化，避免额外去量化开销，提升精度。
- (\tilde{P}V)的FP8量化：鉴于(\tilde{P})的分布特点，将(\tilde{P})、V量化为FP8（E4M3），采用静态量化和按通道量化，在保持精度的同时利用GPU张量核心加速计算。
- 针对FP22累加器的FP32 MMA缓冲区：因实际CUDA实现中FP8矩阵乘法累加器为FP22导致精度损失，采用两级累加策略，用FP32缓冲区累加FP22值，还提出可选的平滑V技术提升精度。
- 实验
- 实验设置：在多种语言、图像和视频生成模型上进行实验，对比SageAttention2与多种基线方法，使用不同数据集和指标评估。
- 内核速度和精度：SageAttention2在RTX4090上比FlashAttention2和xformers快约3倍和4.5倍，在Hopper GPU上与FlashAttention3(fp8)速度相当但精度更高，在CogvideoX模型上精度优于其他基线方法。
- 端到端性能：SageAttention2在多种模型上保持端到端指标，可视化结果显示其生成的图像和视频质量高，且能显著加速模型，如在CogvideoX (1.5 - 5B)上实现1.8倍加速且无指标损失。
- 消融实验：线程级量化、平滑Q和两级累加技术的开销分别为0.35%、3.7%和0%，平滑V可提升精度，但在部分模型中无明显效果。
- 研究结论：SageAttention2是一种高效且准确的量化注意力机制，通过创新量化方法和精度提升技术，在速度和精度上优于多种现有方法，在不同类型模型中保持端到端性能，为加速注意力计算提供有效方案。
作用
SageAttention2 是清华大学陈键飞团队提出的高效注意力计算框架，其核心作用是通过低比特量化与硬件优化技术，显著提升注意力计算效率，同时保持模型精度。以下是其具体作用与技术实现：
SageAttention2是一种全新的即插即用注意力模块，其作用主要包括以下几个方面：
- 加速推理速度：采用4 - Bit量化技术，在多种硬件平台上实现了显著的推理加速。例如，在RTX4090上较FlashAttention2推理速度提升三倍，在A100上提升至1.6倍，在L20、L40、L40S上可以实现2倍的加速，为多样化环境中的AI模型部署提供了可能。
- 保持模型精度：通过对Q、K矩阵进行平滑处理，以及引入Per - thread量化方法等技术手段，克服了低比特量化常见的精度损失问题，在多种大型模型应用中保持了端到端的精度表现，确保了模型的多样性与稳定性。
- 支持多种应用场景：有助于促进AI绘画、视频生成、文本生成等多种应用场景的落地。以开源视频生成模型CogvideoX - 1.5 - 5B为例，采用SageAttention2后，其端到端的推理速度提升达1.8倍，且在视频生成效果上无损失。