SageAttention2

Paper: https://arxiv.org/abs/2411.10958

https://github.com/thu-ml/SageAttention

“SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization”由Jintao Zhang等人撰写。文章提出SageAttention2，通过线程级INT4量化、Q矩阵平滑、两级累加策略等技术，在提升注意力计算效率的同时保持精度，在多种模型上取得优异性能。

0.简介

1. 研究背景：随着序列长度增加，注意力机制的二次时间复杂度使其高效实现变得关键。现有优化方法各有局限，如线性和稀疏注意力方法适用范围有限，常用的注意力方法如FlashAttention、xformers和SageAttention等虽有不错性能，但SageAttention存在INT8矩阵乘法速度慢和特定GPU加速受限的问题。
2. 相关工作
   • FlashAttention：将注意力计算中的矩阵按token维度分块并行计算，降低计算复杂度，提升计算效率。
   • 量化：通过将高精度矩阵转换为低精度格式加速矩阵乘法，不同量化器在数值格式和粒度上有差异。
   • SageAttention：基于FlashAttention的分块策略，将Q、K量化为INT8，对K进行预处理以保持精度，对$\tilde{P}$、V使用FP16并降低累加器精度加速计算，但存在局限性。
3. SageAttention2方法
   • 平滑Q：由于INT4数值范围有限，存在异常值影响量化精度。通过减去Q每个块的均值平滑Q，结合对K的平滑，将$Q{K}^{\top }$计算分解，分预处理和注意力两个阶段，提升INT4量化精度。
   • INT4线程级量化：在SageAttention的基础上，提出线程级量化，根据GPU线程和矩阵内存布局，以更细粒度进行量化，避免额外去量化开销，提升精度。
   • $\tilde{P}V$的FP8量化：鉴于$\tilde{P}$的分布特点，将$\tilde{P}$、V量化为FP8（E4M3），采用静态量化和按通