Blog Tags GitHub
Back to blog
attention-mechanism model-compression

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization

使用 INT4 线程级量化和 FP8 矩阵乘加速注意力计算，配合 Q+K 平滑与两级累加策略保持精度

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization

一、论文概述

项目内容
标题SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
作者Jintao Zhang*, Haofeng Huang*, Pengle Zhang, Jia Wei, Jun Zhu, Jianfei Chen (* 同等贡献)
机构Tsinghua University
论文arXiv:2411.10958
代码https://github.com/thu-ml/SageAttention
发表ICML 2025
发布2024-11-17 (v1), 2025-10-01 (v7)

二、核心思想

问题定义

量化在 Transformer 的线性层中已被广泛应用，但在注意力计算中的加速应用仍然有限。SageAttention 将 QK⊤QK^\top 量化为 INT8、P~V\widetilde{P}V 使用 FP16 累加器，存在两个局限：

  1. W1: INT8 Matmul 的速度仅为 INT4 的一半
  2. W2: FP16 累加器仅在 RTX 4090/3090 上提供加速，在其他 GPU（L20、H100 等）上无效

解决方案概述

SageAttention2 的核心方案是将 Q,KQ, K 量化为 INT4、P~,V\widetilde{P}, V 量化为 FP8 (E4M3)，并提出三项精度增强技术：

  1. Q+K 平滑：消除 Q 和 K 中的通道异常值，使 INT4 量化更准确
  2. 线程级 INT4 量化：基于 PTX MMA 指令内存布局，每个 GPU 线程对应一个量化尺度，精度接近 token 级量化但无额外开销
  3. 两级累加策略：发现 FP8 矩阵乘的 FP32 累加器实际只有 22 位有效精度（FP22），使用 FP32 缓冲区累积 FP22 结果以恢复精度

两种变体

变体Q,K 量化P~\widetilde{P}, V 量化适用平台
SageAttn2-4bINT4 per-threadFP8 per-block/channelRTX4090, L20, A100 等
SageAttn2-8bINT8 per-threadFP8 per-block/channelHopper (H100/H20) 等无 INT4 Tensor Core 的平台

三、技术架构

整体工作流

SageAttention2 工作流

流程步骤：① 平滑 Q, K, V → ② GEMV 计算 ΔS\Delta S → ③ 线程级量化 Q,K 和通道级量化 V → ④ 执行 SageAttention2 kernel → ⑤ 校正输出

核心公式

Q+K 平滑（Equation 2）：

γ(Qi)=Qi−qˉi,γ(Kj)=Kj−kˉ\gamma(Q_i) = Q_i - \bar{q}_i, \quad \gamma(K_j) = K_j - \bar{k}

其中 qˉi=mean(Qi)\bar{q}_i = \text{mean}(Q_i)、kˉ=mean(K)\bar{k} = \text{mean}(K) 是沿 token 维度的均值。

平滑后的 QK⊤QK^\top 分解：

Sij=QiKj⊤=γ(Qi)γ(Kj)⊤+qˉiγ(Kj)⊤+bS_{ij} = Q_i K_j^\top = \gamma(Q_i)\gamma(K_j)^\top + \bar{q}_i\gamma(K_j)^\top + b

其中 ΔSij=qˉiγ(Kj)⊤\Delta S_{ij} = \bar{q}_i\gamma(K_j)^\top 通过 GEMV 计算，bb 是 softmax 不变偏置。

线程级量化（Equation 8）：

基于 PTX mma.m16n8k64 指令的内存布局：

  • Qw[8k+i]Q_w[8k+i] 共享一个量化尺度
  • Kj[8k+2i]K_j[8k+2i] 和 Kj[8k+2i+1]K_j[8k+2i+1] 共享一个量化尺度

这使得每个 GPU 线程在反量化时仅使用单个尺度值，避免了 per-token 量化的向量点积开销。

FP8 量化 P~\widetilde{P}：

由于 P~ij=exp⁡(Sij−mij)∈[0,1]\widetilde{P}_{ij} = \exp(S_{ij} - m_{ij}) \in [0,1]，使用静态尺度 δP=1/448\delta_P = 1/448 量化到 E4M3。V 按通道量化以应对通道异常值。

FP22 累加器修复（两级累加）：

R_ij = Matmul_FP8(P_ij_quantized, V_j_quantized)  // FP22 精度
O_ij = diag(exp(m_{i,j-1} - m_{ij})) * O_{i,j-1} + R_ij  // FP32 累加

R_ij 在寄存器中以 FP22 精度计算（对小批量 bk=64b_k=64 足够），然后以 FP32 高精度累积到 O_ij。

可选 V 平滑：

当 V 存在通道偏差时：V=V−Vm→V = V - \overrightarrow{V_m}，最终输出 O=O+Vm→O = O + \overrightarrow{V_m}。因为 P~\widetilde{P} 每行和为 1，所以 P~Vm→=Vm→\widetilde{P}\overrightarrow{V_m} = \overrightarrow{V_m}。

模型组件

组件说明关键参数
Q+K Smoothing减去 Q 和 K 的 token 维度均值消除异常值，提升 INT4 精度
Per-thread Quantization基于 PTX MMA 内存布局的量化分组32× 和 4× 于 per-block 更细粒度
FP8 E4M3 for P~\widetilde{P}静态尺度量化 P~∈[0,1]\widetilde{P} \in [0,1]δP=1/448\delta_P = 1/448
Per-channel FP8 for V按通道量化 V 应对异常值每通道独立尺度
Two-level AccumulationFP22 MMA + FP32 缓冲修复 FP22 精度损失
Optional V Smoothing减去 V 的通道均值仅在 V 有通道偏差时有效

前向传播算法（Algorithm 1）

预处理: K = K - mean(K), (δ_V, V̂) = ψ_V(V)   // 按通道
分割 Q 为 T_m 个块 {Q_i}; 分割 K,V 为 T_n 个块 {K_j},{V_j}
for i = 1 to T_m do
    q̄_i = mean(Q_i), (δ_Q, Q̂_i) = ψ_Q(Q_i - q̄_i)  // 线程级
    for j in [1, T_n] do
        (δ_K, K̂_j) = ψ_K(K_j)  // 线程级
        S_ij = Matmul_INT4(Q̂_i, K̂_j) + GEMV(q̄_i, K̂_j^T)  // INT4 + FP8
        m_ij, P̃_ij, l_ij = OnlineSoftmax(...)
        O_ij(FP22) = Matmul_FP8(P̃_ij*448, V̂_j)
        O_ij(FP32) = diag(exp(...)) * O_{i,j-1}(FP32) + O_ij(FP22)
    end for
    O_i = O_{i,T_n}(FP32) / 448 * δ_V / l_i,T_n
end for

四、核心创新

创新点说明依据
Q+K 联合平滑SageAttention 仅平滑 K，本文发现 Q 也有显著异常值，联合平滑 CosSim 从 98% 提升到 99.46%Table 4, CogvideoX 全层平均
INT4 线程级量化首次将注意力量化到 INT4，利用 PTX MMA 内存布局实现 token 级精度无额外开销Table 6, CosSim 99.45% vs per-block 98.03%
FP8 用于 P~V\widetilde{P}V替代 SageAttention 的 FP16，Tensor Core 速度翻倍Table 7, E4M3 CosSim 99.44% vs FP16 99.45%
发现 FP22 累加器首次发现 NVIDIA mma.f32.f8.f8.f32 的 FP32 累加器实际只有 22 位有效精度实验验证：1 sign + 8 exp + 13 mantissa
两级累加策略用 FP32 缓冲区累积 FP22 结果，将误差限制在 block 范围内首次将两级累加应用于注意力计算

五、实验结果

端到端指标对比（Table 2）

文本生成 — Llama3.1 (8B)

方法WikiText↓Lambda↑MMLU↑Longbench↑
Full-Precision6.0130.8150.63549.40
SageAttention6.0170.8120.63449.55
SageAttn2-4b6.2560.7980.60748.79
SageAttn2-8b6.0190.8110.63449.59

SageAttn2-8b 几乎零损失，SageAttn2-4b 有微小降级但仍远优于其他基线。

视频生成 — CogvideoX (1.5-5B)

方法CLIPSIM↑CLIP-T↑VQA-a↑VQA-t↑FScore↑
Full-Precision0.17780.997970.23170.9282.507
FlashAttn3-fp80.15620.99186.5312.181✘
SageAttn2-4b0.17210.997857.72952.9892.884
SageAttn2-8b0.17750.998069.49274.4152.487

FlashAttn3(fp8) 在视频生成上完全失败（FScore 为 ✘），而 SageAttn2-8b 几乎无损。

图像生成 — Flux (schnell)

方法FID↓sFID↓CLIP↑IR↑
Full-Precision10.96016.64826.1801.009
SageAttention10.94416.64126.1711.008
SageAttn2-4b10.57717.49726.1410.998
SageAttn2-8b10.92716.72326.1751.009

SageAttn2-4b 在 FID 上甚至优于全精度！

内核速度对比（Figure 5）

内核速度对比

RTX4090 (head_dim=128):

  • SageAttn2-4b 比 FlashAttention2 快约 3×，比 xformers 快约 4.5×
  • SageAttn2-8b 比 FlashAttention2 快约 2.7×，比 xformers 快约 4×
  • 峰值 481 TOPS (RTX4090)

多 GPU 加速比汇总（Table 9）：

方法30904090A100L40L20H100H20
FlashAttention21.00×1.00×1.00×1.00×1.00×1.00×1.00×
SageAttention11.97×1.96×1.37×1.45×1.24×1.53×1.52×
SageAttention2✘2.93×✘2.60×2.46×2.61×3.12×

精度分析

不同平滑方法（Table 4）：

方法CosSim↑Relative L1↓RMSE↓
None80.04%0.39060.2223
Smooth K98.07%0.14930.0743
Smooth Q98.30%0.12500.0712
Smooth Q+K99.46%0.06480.0334

不同量化粒度（Table 6）：

方法CosSim↑Relative L1↓RMSE↓
Per-token99.45%0.06490.0335
Per-thread99.45%0.06220.0313
Per-block98.03%0.14920.0744
Per-tensor97.15%0.18000.0865

Per-thread 精度与 per-token 几乎相同，且远优于 per-block。

不同数据类型（Table 7）：

P~,V\widetilde{P},VCosSim↑Relative L1↓RMSE↓
INT877.05%0.56180.5044
E5M299.20%0.09050.0903
E4M399.44%0.06830.0347
FP1699.45%0.06490.0335

E4M3 精度与 FP16 非常接近。

端到端延迟（Table 8）

模型GPU原始SageAttn2-8bSageAttn2-4b
CogvideoX (2B)RTX409086s54s52s
CogvideoX (1.5-5B)RTX40901040s577s555s
HunyuanVideoL202221s1486s1435s
MochiL202336s1316s1190s
Llama3.1 (48K)RTX40909.2s5.7s5.6s
Llama3.1 (100K)L2039.9s25.4s23.2s

技术开销分析（Table 18）

技术TOPS开销
基础 Attention (INT4+FP8)284—
+ Per-thread quantization2830.35%
+ Two-level accumulation2830%
+ Smoothing Q2733.7%

三项技术的总开销仅约 4%，可忽略不计。

长上下文实验（Table 14, Figure 19）

在 Llama-3-262k (8B) 上，序列长度达 262K tokens：

  • SageAttention2 在 InfiniBench 所有子任务上保持全精度性能
  • FlashAttention3(fp8) 在 Retr.KV 任务上从 0.4→7.0 大幅下降

Needle-in-a-Haystack 结果

可视化对比

CogvideoX 对比 SageAttn2-8b 与 FlashAttention3 在 CogvideoX-1.5 上的对比 — FlashAttn3(fp8) 产生严重模糊。

Mochi/HunyuanVideo 对比 SageAttn2-8b 在 Mochi 和 HunyuanVideo 上无明显质量损失。

平滑效果可视化

Q 平滑前后量化分布 平滑后 INT4 量化范围被更均匀、充分地利用。

六、消融研究

各技术开销（Table 18, Table 19）

量化粒度TOPS相对 per-tensor 开销
Per-tensor286—
Per-block2840.7%
Per-thread2831.0%
Per-token2686.3%

Per-token 量化有 6.3% 的性能下降，而 per-thread 仅 1.0%，证明 per-thread 是精度和速度的最佳平衡。

V 平滑的收益（Appendix Table 10）

平滑 VCosSim↑Relative L1↓RMSE↓
×98.25%0.19800.2387
√99.75%0.04060.0773

V 平滑在 V 有通道偏差时显著提升精度。

七、相关工作

方向代表工作SageAttention2 的区别
FlashAttention 系列FA1/2/3 (Dao et al.)不改变精度，纯系统优化
SageAttention (v1)Zhang et al., 2025cINT8 per-block + FP16，仅平滑 K
SmoothQuantXiao et al., 2023激活-权重间权衡，不适用于 QK^T
QuARoTAshkboos et al., 2024Hadamard 旋转 + INT4，精度差
QServeLin et al., 2025W4A8KV4，系统级共设计
FlashAttention3-fp8Shah et al., 2024仅 Hopper，视频生成质量降级

八、总结

核心贡献

  1. INT4 线程级量化：首次将注意力 Q,K 量化到 INT4，基于 PTX MMA 内存布局实现 token 级精度无额外开销
  2. Q+K 联合平滑：发现 Q 也存在显著异常值，联合平滑后 CosSim 达 99.46%
  3. FP8 用于 PV：替代 FP16，利用 Tensor Core 实现速度翻倍，E4M3 精度接近 FP16
  4. 发现 FP22 累加器：首次揭示 mma.f32.f8.f8.f32 的 FP32 累加器仅有 22 位有效精度
  5. 两级累加策略：FP22 MMA + FP32 缓冲，将误差限制在 block 范围内
  6. 跨模型零损失加速：覆盖 LLM、视频生成、图像生成、音频、图像分类五大领域

性能总结

对比对象RTX4090L20H100/H20
vs FlashAttention23×——
vs xformers4.5×——
vs FlashAttention3(fp8)——速度匹配，精度远优
vs SageAttention11.5×1.7×1.7×

局限性

  1. 主要在 NVIDIA Ada/Hopper 架构验证，AMD GPU 未测试
  2. 训练场景未涉及，仅验证推理加速
  3. INT4 变体在 Hopper 上不可用（缺乏 INT4 Tensor Core 支持），需使用 INT8 变体
  4. SageAttn2-4b 在视频生成上有轻微质量降级（VA/VQA 略低于 8b 版本和全精度）

九、参考资源