4比特量化三倍加速不掉点！清华即插即用的SageAttention迎来升级

从8-Bit到4-Bit。

原标题：4比特量化三倍加速不掉点！清华即插即用的SageAttention迎来升级
文章来源：机器之心
内容字数：6979字

清华大学陈键飞团队提出SageAttention2：4-Bit即插即用注意力机制，实现3-4.5倍推理加速

本文总结了清华大学陈键飞团队最新提出的SageAttention2论文要点。该工作在之前的SageAttention基础上，进一步将注意力机制中的Q、K矩阵量化到INT4，实现了更高速的推理速度，并在多个大型模型上保持了端到端的精度。

1. 背景与挑战

随着大型模型序列长度的增加，注意力机制(Attention)的计算开销成为瓶颈。虽然线性层的低比特量化已较为成熟，但注意力模块大多仍使用高精度(FP16或FP32)运算。SageAttention率先将Attention中的QK^T量化到INT8，取得了显著加速效果。然而，INT8的矩阵乘法速度仍不及INT4，且FP16的乘法累加器加速仅在特定显卡上有效。因此，将注意力机制量化到INT4成为进一步提升效率的关键，但也面临巨大挑战：直接量化到INT4会导致精度严重下降。

2. SageAttention2的技术方案

为了解决INT4量化带来的精度损失问题，SageAttention2提出了以下技术方案：

1. 对Q和K进行平滑处理： 在对K进行平滑处理的基础上，新增对Q进行平滑处理(Q – mean(Q))，并补偿到最终结果中，有效降低了量化误差。

2. Per-thread量化： 将Q、K矩阵的量化粒度细化到per-thread级别，提高了4-Bit QK^T乘法的精度，且不增加额外开销。

3. FP32寄存器累加FP8 PV乘法结果： 使用FP32寄存器累加FlashAttention分块粒度的PV的FP22乘法结果，避免了累积误差。

4. 使用E4M3格式的FP8： 实验表明，E4M3格式的FP8精度最高，接近FP16。

5. 可选的V矩阵平滑： 对V矩阵进行平滑处理，进一步提升PV矩阵乘法的精度。

3. 实验结果

SageAttention2在速度和精度上都取得了显著提升：

• 速度： 相比FlashAttention2和xformers分别实现了3倍和4.5倍的推理加速，在不同显卡上均有不同程度的加速。

• 精度： 在视频、图像、文本生成等多种大型模型上保持了端到端的精度，例如在CogvideoX-1.5-5B模型上实现了1.8倍的端到端加速，且视频质量无损。

4. 总结

SageAttention2通过一系列创新技术，成功地将注意力机制量化到INT4，实现了显著的推理加速，并在多个大型模型上保持了精度。其即插即用的特性也方便了在实际应用中的部署，为大型模型的效率提升提供了有力支持。该工作已开源，方便开发者使用。

联系作者

文章来源：机器之心
作者微信：
作者简介：专业的人工智能媒体和产业服务平台