SageAttention: Accurate 8-Bit Attention for Plug-and-Play Inference Acceleration

algorithm 2410.02367
attention-quantizationint8smooth-quantizationmixed-precisionpost-training-quantization
← Synthesis ← algorithm Survey Source ↗

SageAttention: Accurate 8-Bit Attention for Plug-and-Play Inference Acceleration #

Jintao Zhang, Jia Wei, Haofeng Huang, Pengle Zhang, Jun Zhu, Jianfei Chen | Tsinghua University | 2024-10 | https://arxiv.org/abs/2410.02367 Category: algorithm | Tags: attention-quantization, int8, smooth-quantization, mixed-precision, post-training-quantization

§1 TL;DR #

Smooth K（减均值消除 K 的 channel outlier，softmax 平移不变性保证精确）+ INT8 $QK^\top$ / FP16-with-FP16-accumulator $PV$ 混合精度 attention，RTX4090 达 341 TOPS（2.1× FlashAttention2），端到端精度损失 <0.2%，即插即用。

§2 Q1 / Q2 / Q3 #

Q1 痛点 #

Attention 的 $O(N^2)$ 复杂度在长序列场景（视频生成 8K–128K tokens、LLM prefill）下成为推理瓶颈。现有量化加速工作几乎全部集中在 Linear 层，attention 仍以 FP16 运行。直接将 attention 中的 $Q, K, P, V$ 量化为 8-bit 会导致灾难性精度退化：

证据：Unidiffuser 使用 INT8 attention 生成完全模糊的图像；Llama2 在 MMLU 上降至 25.5%（随机猜测水平）。FlashAttention3 的 FP8 版本在扩散模型上同样灾难性失败（Unidiffuser FID: 394 vs 163）。

Q2 方法 #

SageAttention 提出三个核心技术组成的 post-training quantization 算法：

  1. Smooth K：$\gamma(K) = K - \mathrm{mean}(K)$，减去所有 token 共享的 channel 均值消除 outlier。由于 softmax 的平移不变性 $\sigma(qK^\top - q \cdot \mathrm{mean}(K)) = \sigma(qK^\top)$，该变换精确保留注意力分布，开销 < 0.2%
  2. 混合精度 Matmul：$QK^\top$ 使用 INT8 per-block 量化（RTX4090 上 INT8 比 FP16 快 4×），$\widetilde{P}V$ 保持 FP16 并使用 FP16 accumulator（比 FP32 accumulator 快 2×，精度完全一致）
  3. 自适应量化：对每层 profiling 余弦相似度，> 99.8% 的层使用全 INT8 kernel（额外提速 ~4%），其余层回退到混合精度

    4. 核心技术壁垒：$K$ 的 channel-wise outlier 不是 token 间的信号变异，而是所有 token 共享的大偏置。这一洞察使得均值减法成为精确变换（非近似），是整个方法准确性的基石。如果不认识到这一点，自然会尝试 per-channel 量化（但 $QK^\top$ 的内轴无法用 per-channel scale 反量化）或 SmoothQuant 式迁移（但 $Q$ 同样存在 outlier，无法将 $K$ 的 scale 迁移到 $Q$）。

    Q3 结果 #

    • RTX4090 上达到 341 TOPS（INT8 理论峰值 660 TOPS 的 52%），FlashAttention2 仅 165 TOPS
    • 比 FlashAttention2 快 2.1×，比 xformers 快 2.7×，实际模型平均加速 2.83×
    • 在 Llama2、CogVideoX、Unidiffuser、UltraPixel、TIMM、Llava1.6 六个模型上，端到端指标损失平均仅 ~0.2%
    • 即插即用，无需再训练或校准数据（仅 adaptive 选择需少量 profiling）

    Problem Formulation #

    目标函数：最小化量化 attention 输出与全精度输出的误差，同时最大化 kernel 吞吐量：

    $$\min_{\psi_Q, \phi_K, \psi_P, \psi_V} \mathbb{E}[\|O_{\text{quant}} - O_{\text{fp16}}\|] \quad \text{s.t. throughput}(\text{kernel}) \gg \text{throughput}(\text{FlashAttention2})$$

    其中 $O_{\text{quant}}$ 由量化后的 matmul 计算（Eq. 4-5），$\phi_K = \psi_K \circ \gamma$ 包含 smooth 变换。

    假设：

    • Softmax 平移不变性：$\sigma(x + c) = \sigma(x)$，对任意标量 $c$
    • INT8 Tensor Core 可用且比 FP16 Tensor Core 快（Ampere/Ada 架构上 4× 理论峰值比）
    • $\widetilde{P} = \exp(S - \mathrm{rowmax}(S))$ 的最大值恒为 1（softmax 性质），因此 per-block 量化的 static scale $s = 1/127$ 等价于 per-token 精度

    算法一步的输入/输出：

    • 输入：$Q, K, V \in \mathbb{R}^{N \times d}$（FP16），block sizes $b_q, b_{kv}$
    • 输出：$O \in \mathbb{R}^{N \times d}$（FP16），与全精度 attention 输出误差极小（CosSim > 99.8%）

    §3 架构 / 方法图 #

    Method Core: Smooth K — 唯一关键创新 #

    每篇算法论文都有一个核心创新机制。SageAttention 的核心是 Smooth K：将 $K$ 矩阵的 channel-wise outlier 识别为共享偏置并精确消除。

    Before（直接 INT8 量化）After（SageAttention）
    K 处理直接 per-token/per-block INT8 量化$\gamma(K) = K - \mathrm{mean}(K)$，消除共享偏置后再量化
    $QK^\top$ matmulFP16 或 FP8INT8 per-block（4× 硬件加速）
    $\widetilde{P}V$ matmulFP16 with FP32 accumulator，或 INT8FP16 with FP16 accumulator（2× 加速，零精度损失）
    Unidiffuser FID221.18（per-token INT8）/ 394.13（FA3 FP8）166.93（vs full-precision 163.33）
    CogVideo Fscore1.924（per-token INT8，严重退化）3.718（vs full-precision 3.768）

    Figure 5: SageAttention 的完整数据流

    Paper's Figure 5, verbatim (caption: "Workflow of SageAttention").

    Figure 5 展示了 SageAttention 的两种 kernel variant：(a) SAGEAttn-vT 将 $QK^\top$ 和 $\widetilde{P}V$ 都用 INT8 量化；(b) SAGEAttn-B 对 $QK^\top$ 用 INT8，对 $\widetilde{P}V$ 保持 FP16 并使用 FP16 accumulator。SAGEAttn-B 是主推 kernel，精度更高且速度已达 2× FA2。

    Figure 4: Q, K, V 矩阵的 channel-wise 数据分布

    Paper's Figure 4, verbatim (caption: "Typical examples of data distribution of (Q, K, V)").

    Figure 4 直观展示了 Smooth K 的动机：$K$ 矩阵在 Unidiffuser 和 CogVideoX 中呈现显著的 channel-wise outlier（某些 channel 所有 token 的值都偏大），而这些 outlier 本质上是共享偏置——减去均值后 $K$ 的数值范围显著缩小，量化误差随之降低。$Q$ 虽也有 outlier 但程度较轻，$V$ 分布相对均匀。

    完整算法流程（Algorithm 1: SAGEAttn-B） #

    1. 预处理：$K \leftarrow K - \mathrm{mean}(K)$
    2. 量化：$(\delta_Q, \hat{Q}) = \psi_Q(Q/\sqrt{d})$，$(\delta_K, \hat{K}) = \psi_K(K)$，INT8 per-block
    3. 分块：$\hat{Q}$ 分 $T_m = N/b_q$ 块，$\hat{K}, V$ 分 $T_n = N/b_{kv}$ 块
    4. 并行外循环（每个 SM 处理一个 $\hat{Q}_i$）：
    5. 内循环遍历所有 $\hat{K}_j, V_j$：
    6. $S_i^j = \mathrm{INT8\_Matmul}(\hat{Q}_i, \hat{K}_j^\top) \times \delta_Q[i] \times \delta_K[j]$
    7. Online softmax 更新 $m_i^j, \widetilde{P}_i^j, l_i^j$
    8. $O_i^j = \mathrm{rescale}(O_i^{j-1}) + \mathrm{FP16\_Matmul}(\widetilde{P}_i^j, V_j, \text{accum=FP16})$
    9. 归一化：$O_i = \mathrm{diag}(l_i^{T_n})^{-1} O_i^{T_n}$

      10. §4 作者证明 #

      符号表 #

      符号含义维度
      $Q, K, V$Query, Key, Value 矩阵$N \times d$
      $S = QK^\top / \sqrt{d}$注意力分数矩阵$N \times N$
      $P = \sigma(S)$Softmax 后的注意力权重$N \times N$
      $O = PV$注意力输出$N \times d$
      $\gamma(K) = K - \mathrm{mean}(K)$Smooth 变换$N \times d \to N \times d$
      $\psi(\cdot)$量化器：$(A) \mapsto (\delta_A, \hat{A})$输出 scale + INT8 tensor
      $\psi^{-1}_{\delta}(\hat{A})$反量化器：$\delta \cdot \hat{A}$INT8 → FP16
      $\phi_K = \psi_K \circ \gamma$K 的完整变换：smooth 后量化$N \times d \to$ INT8
      $\widetilde{P}_i^j$Online softmax 的未归一化权重$b_q \times b_{kv}$
      $m_i^j, l_i^j$Online softmax 的 running max 和 running sum$b_q \times 1$

      核心方程与物理意义 #

      Eq. 6 — Smooth K 变换：

      $$\gamma(K) = K - \mathrm{mean}(K), \quad \mathrm{mean}(K) = \frac{1}{N}\sum_{t=1}^{N} K[t,:]$$

      物理意义：减去所有 token 共享的 channel 偏置，使剩余信号（token 间差异）的动态范围缩小，量化误差随之降低。

      正确性证明：对任意 query $q$：

      $$\sigma(q \cdot \gamma(K)^\top) = \sigma(q K^\top - q \cdot \mathrm{mean}(K)^\top) = \sigma(qK^\top)$$

      第二个等号成立是因为 $q \cdot \mathrm{mean}(K)^\top$ 是标量（对所有 key 位置相同），softmax 对输入加常数不变。

      Eq. 4-5 — 量化 attention 计算：

      $$S = \psi^{-1}_{\delta_Q \delta_K}(\hat{Q}\hat{K}^\top), \quad (m', \widetilde{P}) = \tilde{\sigma}(m, S), \quad O = \mathrm{diag}(\exp(m' - m))O_{\text{prev}} + \psi^{-1}_{\delta_P \delta_V}(\hat{P}\hat{V})$$

      物理意义：将 FlashAttention 的 tiled computation 中的两个 matmul 替换为量化版本，online softmax 保持全精度。

      6 项验证检查 #

      #检查项结论
      1Smooth K 的 softmax 平移不变性是否严格成立？成立。$\sigma(x + c\mathbf{1}) = \sigma(x)$ 是 softmax 的基本性质，$q \cdot \mathrm{mean}(K)^\top$ 对每个 key 位置是相同标量
      2INT8 比 FP8 更精确的声明是否有充分证据？Table 2 在 Llama2 和 Unidiffuser 所有层上对比了 INT8 vs E4M3 vs E5M2 的三项指标（CosSim, L1, RMSE），INT8 在 $Q,K$ 量化上全面领先
      3FP16 accumulator 与 FP32 精度一致的声明是否可靠？Tables 4-5 给出平均和最差精度，FP16 与 FP32 的三项指标数值完全相同（到小数点后 4 位），在 $\widetilde{P} \in [0,1]$ 的有界条件下合理
      4Smooth K 的 <0.2% 开销声明是否可信？Table 10 在 Appendix 中验证。均值计算是 $O(Nd)$ 的 reduction + broadcast，相比 $O(N^2 d)$ 的 attention 计算确实可忽略
      5Per-block 量化是否优于 per-tensor/per-token？Table 1 显示 per-block + Smooth K 和 per-token + Smooth K 精度接近且都远优于无 smooth 版本，per-block 在所有模型上表现稳定
      6Adaptive 选择的 99.8% 阈值是否充分验证？阈值选择基于 SAGEAttn-B 的 worst-case CosSim（99.8%），确保 SAGEAttn-vB 仅在精度不劣于主 kernel 时启用。但阈值的 sensitivity analysis 未详述

      本文不包含收敛定理或形式化 bound——属于 post-training 量化方法，核心证明是 Smooth K 的精确性（softmax 不变性）和经验验证。

      §5 实验与数据 #

      关键结果一览 #

      Table 1: Smooth K 的消融实验

      量化方式 (Q,K)Smooth KLlama WikiText ↓CogVideo Fscore ↑Unidiffuser FID ↓UltraPixel FID ↓TIMM ImageNet ↑
      Full-Precision-5.8233.768163.33179.7884.79%
      Per-token✗5.8241.924221.18193.3684.21%
      Per-token✓5.8243.734166.52179.7984.74%
      Per-block✗5.8252.014229.08195.6784.18%
      Per-block✓5.8243.718166.93179.9884.76%
      FlashAttn3 FP8-5.8503.394394.13383.6184.70%

      Smooth K 将 CogVideo 从 1.924 恢复到 3.734（接近 full-precision 的 3.768），Unidiffuser FID 从 221→167（full-precision 163）。FlashAttention3 的 FP8 在扩散模型上灾难性失败（FID 394 vs 163）。

      Table 7: 真实模型加速

      ModelQ,K,V shapeOriginal Attention (TOPS)SageAttention (TOPS)Speedup
      CogVideoX(2, 30, 17776, 64)163.37 (FA2)340.762.09×
      Unidiffuser(2, 24, 1105, 64)105.68 (xformers)246.932.34×
      UltraPixel(8, 30, 4096, 64)79.05 (xformers)252.473.19×
      Llama2(1, 32, 4096, 128)152.18 (FA2)341.002.24×
      TIMM(12, 64, 197, 64)18.91 (Torch)111.415.89×

      平均加速 2.83×。长序列模型（CogVideoX 17K tokens）获得稳定 2× 加速，短序列模型（TIMM 197 tokens）因 baseline 较弱获得最大加速。

      Speed Benchmarks #

      Figure 6: RTX4090 headdim=64 速度对比

      Paper's Figure 6, verbatim (caption: "Speed of SageAttention on RTX4090 (headdim=64)").

      Figure 6 展示在 headdim=64 配置下，SageAttention 在所有序列长度（1K–32K）上一致达到 ~340 TOPS，而 FlashAttention2 峰值仅 ~165 TOPS，xformers ~130 TOPS。causal 和 non-causal mask 下加速比一致。

      Figure 7: RTX4090 headdim=128 速度对比

      Paper's Figure 7, verbatim (caption: "Speed of SageAttention on RTX4090 (headdim=128)").

      Figure 7 确认 headdim=128 下同样保持 ~2× 优势。两种常用 head dimension 都验证了加速的一致性。SageAttention 在 RTX4090（消费级 GPU）上的吞吐量接近 FlashAttention3 在 H100（数据中心 GPU）上 490 TOPS 的水平。

      端到端质量验证 #

      Figure 3: 图像生成质量对比

      Paper's Figure 3, verbatim (caption: "A comparison example").

      Figure 3 直观展示了量化方法对图像生成质量的影响：Full-Precision 和 SageAttention 生成清晰一致的图像，而直接 INT8 量化和 FlashAttention3 的 FP8 版本都生成完全模糊的图像。这是 Smooth K 效果最有力的视觉证据。

      Table 8: 端到端指标损失（代表性结果）

      ModelTaskMetricFull-PrecisionSageAttention损失
      Llama2 (7B)WikiTextPerplexity ↓5.8235.824+0.001
      Llama2 (7B)MMLUAccuracy ↑45.97%45.89%-0.08%
      CogVideoXVideo GenFscore ↑3.7683.718-1.3%
      UnidiffuserImage GenFID ↓163.33166.93+2.2%
      UltraPixelImage GenFID ↓179.78179.98+0.1%
      TIMMImageNetAccuracy ↑84.79%84.76%-0.03%

      所有模型的端到端指标损失均在 2.2% 以内，LLM 的 perplexity 变化仅 0.001。

      INT8 vs FP8 数据类型对比（Table 2） #

      $Q,K$ dtype$\widetilde{P},V$ dtypeCosSim ↑Relative L1 ↓RMSE ↓
      INT8FP1699.99%0.01160.0091
      INT8E4M399.94%0.03453.53e-3
      E4M3E4M399.81%0.06075.93e-3
      E5M2E5M299.22%0.12131.20e-2
      INT8INT899.70%0.10356.82e-3

      INT8 用于 $Q,K$ 在精度和速度上都优于 FP8（E4M3/E5M2）。$\widetilde{P},V$ 保持 FP16 可达 99.99% CosSim，而 INT8 的 worst-case 仅 56.4%（Table 3），这决定了混合精度策略的必要性。

      Dataset Analysis #

      评估数据集：

      ModelDatasetTaskMetric
      Llama2 (7B)WikiText / LAMBADA / MMLULanguage modeling / comprehension / knowledgePPL / Acc / Acc
      CogVideoXOpen-Sora promptsText-to-videoFscore
      Unidiffuser, UltraPixelCOCO annotationsText-to-imageFID
      TIMMImageNet / CIFAR-10/100Image classificationAccuracy
      Llava1.6MMMUVisual QAAccuracy

      数据集覆盖了 LLM、视频生成、图像生成、图像分类和视觉问答五类任务。论文未提及评估集与训练集的重叠检查，但所使用的基准（WikiText, ImageNet, MMLU 等）均为标准公开测试集。

      §6 论证链 #

      Step前提论证结论
      1$K$ 的 channel-wise outlier 本质是共享偏置，而非 token 间信号差异（Figure 4）Softmax 平移不变性保证 $\sigma(qK^\top - c) = \sigma(qK^\top)$，减均值是精确变换Smooth K ($\gamma(K) = K - \mathrm{mean}(K)$) 消除 outlier 且不改变 attention 输出
      2INT8 for $\widetilde{P},V$ worst-case CosSim 仅 56.4%（Table 3）；FP16 accum 与 FP32 accum 精度完全一致（Tables 4-5）$\widetilde{P} \in [0,1]$ 有界，FP16 精度足够且 accumulator 不溢出FP16 $\widetilde{P}V$ with FP16 accumulator 是最优折衷：2× 快于 FP32 accum，远准于 INT8
      3INT8 per-block + Smooth K 在所有模型上 CosSim > 99.8%（Table 1, Table 2）结合 Step 1 和 Step 2，INT8 $QK^\top$ + FP16 $PV$ 的混合精度流水线精度有保证SAGEAttn-B kernel 实现 2.1× FA2 加速且端到端损失 <0.2%（Table 7, 8）
      4部分层的全 INT8 variant（SAGEAttn-vB）CosSim > 99.8%（等于 SAGEAttn-B 的 worst case）Profiling 每层精度后，安全地切换到更快 kernelAdaptive 选择额外提速 ~4%，无精度风险

      §7 实现 cross-reference #

      代码仓库 #

      • GitHub: https://github.com/thu-ml/SageAttention
      • 安装: pip install sageattention
      • 语言: Python + Triton（GPU kernel DSL → PTX）
      • 依赖: PyTorch 2.4.0+cu121, Triton nightly (20240816), Python 3.11

      核心技术壁垒（详述） #

      $K$ 的 channel-wise outlier 是共享偏置这一洞察是整个方法的基石。实际开发中，面对 $K$ 的 outlier，直觉做法是 (a) per-channel 量化——但 $QK^\top$ 的内轴（channel dim）的 scale factor 无法在 dequantization 中使用；(b) SmoothQuant 式 channel-to-token 迁移——但 $Q$ 也有 outlier，迁移后 $Q$ 的量化更差。唯一正确的路径是认识到 outlier 的偏置本质 + softmax 的平移不变性，使得直接减均值成为精确变换。这种洞察来自对真实模型 $K$ 分布的仔细分析（Figure 4），不是理论推导能自动得到的。

      关键实现细节 #

      1. ROPE 与量化融合：$Q, K$ 经过 ROPE 计算后在 shared memory 写回 global memory 之前直接量化为 INT8，避免了一次额外的 HBM round-trip。同时将 $1/\sqrt{d}$ 系数融合进 $Q$ 的量化步骤。这一融合使得量化的 IO 开销为零。
        1. $\widetilde{P}$ 的 static scale：由于 $\widetilde{P} = \exp(S - \mathrm{rowmax}(S))$ 的最大值恒为 1，per-block 量化可使用固定 scale $s = 1/127$，省去了运行时计算 scale 的开销，同时精度等价于 per-token 量化。

          2. Reproducibility & Ecosystem #

          • 代码完全开源，pip install sageattention 即可使用
          • SageAttention 已被广泛集成到视频/图像生成推理流程中
          • 后续工作 SageAttention2（2411.10958）扩展到 4-bit $\widetilde{P}V$，SageAttention3（2505.11594）进一步利用 Blackwell FP4 Tensor Core
          • 社区已在 ComfyUI、Diffusers 等框架中集成 SageAttention 作为加速后端
          • 论文的 Smooth K 技术已被后续多个 attention 量化工作采用

          Built: 2026-05-26T13:50:00Z | L1 hash: b00160c73a8d2a69ce95c7398453c3b1f46e8be7d39eaf5e2793b44611fc6c5c