SageAttention2++: A More Efficient Implementation of SageAttention2

code 2505.21136
attention-kernelFP8-quantizationtensor-core-mmaCUDA-PTXinference-acceleration
← Synthesis ← code Survey Source ↗

SageAttention2++: A More Efficient Implementation of SageAttention2 #

Jintao Zhang, Xiaoming Xu, Jia Wei, Haofeng Huang, Pengle Zhang, Chendong Xiang, Jun Zhu, Jianfei Chen (Tsinghua / Shengshu) | 2025-05 | Category: code | Tags: attention-kernel, FP8-quantization, tensor-core-mma, CUDA-PTX, inference-acceleration

§1 TL;DR #

SageAttention2++ 将 PV matmul 的累加器从 FP32 切换到 FP16（mma.f16.f8.f8.f16），通过缩窄 P/V 量化范围（$P_r \times V_r \leq 1023.5$）避免溢出，实现 3.9× over FlashAttention2 的 attention kernel 加速，精度无损。

§2 Q1 / Q2 / Q3 #

Q1 痛点: SageAttention2 使用 mma.f32.f8.f8.f32 指令做 PV matmul，该指令仅比 FP16 快 2×。而同一硬件（Ada/Blackwell）上存在 mma.f16.f8.f8.f16 指令，理论上 4× 于 FP16。但直接使用会导致 FP16 累加器溢出（32 个 FP8×FP8 乘积之和可超过 ±65504）。

Figure 1: Speed comparison RTX4090 headdim=128

Paper Figure 1: Speed comparison between SageAttention2++ and baselines on RTX4090, headdim=128. SageAttn2++(4+8) achieves ~3.9× over FlashAttention2 across sequence lengths from 1K to 32K.

Fig 1 直接展示了核心性能主张：在 RTX4090 上 headdim=128 的不同序列长度下，SageAttn2++ 系列全面超越 FlashAttention2 和 SageAttention2 基线。

Q2 方法: 两步优化——

  1. 量化范围缩窄：将 $\hat{P}$ 的量化范围从 448 缩到 $P_r = 224$，$\hat{V}$ 从 448 缩到 $V_r = 4.5$，使 mma.m16n8k32 的 32 次乘加结果不超 FP16 范围。约束：$P_r \times V_r \leq 2047$。
  2. Delayed FP32 Buffering：连续 2 次 MMA 结果在 FP16 中累加后再转 FP32，减半类型转换 PTX 指令开销。约束收紧为 $P_r \times V_r \leq 1023.5$。选定 $224 \times 4.5 = 1008 \leq 1023.5$。

    3. 核心技术壁垒: 精确认识到 mma.m16n8k32 指令的 k=32 累加维度与 FP16 65504 上限共同决定了 $P_r \times V_r$ 的硬约束，并找到不损精度的非对称分配方案 ($P_r \gg V_r$)。这要求对 PTX ISA 的 MMA 语义和 post-softmax $P$ 值的实际分布有深度理解。

    Q3 结果:

    • Kernel: SageAttn2++(4+8) 3.9× / (8+8) 3.0× over FlashAttention2
    • 精度: cosine similarity 99.97%（与 SageAttn2 相同）
    • End-to-end: 在 Llama3.1, CogvideoX, HunyuanVideo, Wan, Flux, SD3.5 上指标几乎无损

    §3 架构 / 方法图 #

    Table 1: Matmul instruction speedup

    Paper Table 1: Speedup of different matmul instructions on RTX4090/RTX5090. FP8 with FP16 accumulator achieves 4× speedup over FP16 baseline — the hardware foundation of SageAttention2++.

    Table 1 是整篇工作的硬件洞察基础：同一 GPU 上三种 MMA 指令的吞吐差异（1×/2×/4×），4× 的指令此前未被 SageAttention2 利用。

    flowchart TD subgraph "SageAttention2++ Kernel" A["Load Q, K tiles → SMEM"] --> B["QK^T via INT4/INT8 MMA
    (unchanged from SageAttn2)"] B --> C["Online Softmax → P̃ (FP32)"] C --> D["Quantize P̃ → FP8
    δP = max|P̃| / 224"] E["Load V tile → SMEM"] --> F["Quantize V → FP8
    δV = colmax|V| / 4.5"] D --> G["PV Matmul: mma.f16.f8.f8.f16
    (4× speed vs FP16)"] F --> G G --> H{"Every 2 MMAs?"} H -->|Yes| I["acc_fp32 += float32(acc_fp16)
    Delayed FP32 Buffering"] H -->|No| G I --> J["Dequantize: O = acc × δP × δV"] J --> K["Online softmax rescale → Write O to HBM"] end

    关键变化仅在 PV matmul 路径：accumulator datatype FP32→FP16 + 量化范围 448→224/4.5。QK^T 路径和 online softmax 完全继承 SageAttention2。

    §4 作者证明 #

    符号含义
    $P_r$$\hat{P}$ 量化后的最大绝对值上限（代替 E4M3 的 448）
    $V_r$$\hat{V}$ 量化后的最大绝对值上限
    $\delta_P, \delta_V$per-block / per-channel scale factors
    $k=32$mma.m16n8k32 的累加维度

    核心方程及物理意义:

    $$|32 \times p \times v| \leq 65504$$

    含义：MMA 指令沿 k=32 维累加乘积，worst-case（同号最大值）不得溢出 FP16 range。化简得 $|p \times v| \leq 2047$，即 $P_r \times V_r \leq 2047$。

    $$\delta_P = |\max(\widetilde{P})| / P_r, \quad \delta_V = \text{colmax}(|V|) / V_r$$

    含义：对称量化 scale factors，将实际值映射到缩窄后的整数范围 $[-P_r, P_r]$ 和 $[-V_r, V_r]$。

    $$P_r \times V_r \leq 2047 / 2 = 1023.5$$

    含义：Delayed FP32 Buffering 连续累加 2 次 MMA 结果再转 FP32，约束收紧 2×。

    6 minimum checks:

    1. ✅ 维度一致性: $P_r \times V_r = 224 \times 4.5 = 1008 \leq 1023.5$
    2. ✅ 边界条件: post-softmax $P$ 值域 $[0,1]$，量化到 $[0, 224]$ 而非 $[-224, 224]$，单调非负进一步降低溢出风险
    3. ✅ Scale factor 定义: $\delta_P = \max(\widetilde{P})/224$ 是 per-block scalar; $\delta_V = \text{colmax}(|V|)/4.5$ 是 per-channel vector
    4. ✅ Dequant 恢复: $PV = \hat{P}\hat{V} \times \delta_P \times \delta_V$，scale factors 在 matmul 外乘回
    5. ✅ $P$ 非负性未被利用: worst-case bound 保守地假设 $|p|$ 可达 $P_r$，实际 $p \in [0, P_r]$ 因 softmax
    6. ✅ 表 2 验证: 所有满足约束的 $(P_r, V_r)$ 组合 cosine similarity 均为 99.97%

      7. §5 实验与数据 #

      Table 2: Attention accuracy with different quantization ranges

      Paper Table 2: Average attention accuracy (CogvideoX) for various $(P_r, V_r)$ pairs. All configurations maintain 99.97% cosine similarity, demonstrating that narrowing quantization range has negligible effect on attention output quality.

      Table 2 是精度安全性的核心证据。$V_r$ 从 448 缩窄到 4.5（100× 压缩）对精度零影响，因 post-softmax $P$ 的实际数值远小于 E4M3 满量程，per-channel $V$ scaling 已捕获通道间差异。

      Figure 3: Speed comparison RTX5090 headdim=128

      Paper Figure 3: Kernel speed on RTX5090, headdim=128. The speedup pattern is consistent with RTX4090, confirming the technique generalizes across Ada and Blackwell architectures.

      RTX5090 上的一致表现说明优化并非针对特定 SM 版本的 micro-architectural quirk，而是利用了跨代稳定的指令级吞吐差异。

      Figure 5: Visible quality comparison

      Paper Figure 5: Qualitative comparison of generated outputs between Full-Precision, SageAttention2, and SageAttention2++. Visual quality is indistinguishable.

      定性对比进一步确认端到端无肉眼可见差异。

      Table 3: End-to-end metrics

      Paper Table 3: Comprehensive end-to-end evaluation across 6 models (Llama3.1, CogvideoX, HunyuanVideo, Wan, Flux, SD3.5). SageAttn2++(8+8) achieves near-identical metrics to full precision across all tasks.

      Table 3 是端到端验证的全量数据。值得注意的例外：Wan (4+8) 的 VQA-a 从 53.3 降至 29.7（44% drop），但这同样出现在 SageAttn2(4+8) 中，说明是 INT4 QK 量化的问题而非 SageAttn2++ 引入。

      §6 论证链 #

      StepClaimEvidenceStrength
      1mma.f16.f8.f8.f16 比 mma.f32.f8.f8.f32 快 2×Table 1 实测 RTX4090/5090Strong — hardware spec + microbenchmark
      2直接使用 FP16 acc 会溢出分析：$448 \times 448 \times 32 = 6,422,528 \gg 65504$Strong — arithmetic proof
      3缩窄 $P_r, V_r$ 使 $P_r \times V_r \leq 1023.5$ 可避免溢出Eq.1-3 推导Strong — closed-form bound
      4缩窄量化范围不损精度Table 2: cosine sim 99.97% across all configsStrong — empirical on CogvideoX
      5Kernel 整体加速 3.9×/3.0×Fig.1-4: 多序列长度、两款 GPUStrong — reproducible benchmark
      6端到端指标无损Table 3: 6 models × 多指标Moderate — (8+8) strong; (4+8) on Wan has 44% VQA drop

      §7 实现 cross-reference #

      Repository:

      语言: CUDA + Python (PyTorch extension)

      关键实现细节:

      1. PTX 指令选择: 核心在于 mma.m16n8k32 指令调用时指定 .f16 accumulator 而非 .f32。这是单行 PTX inline asm 的差异，但需要整个 quantization pipeline 配合调整 scale factors。
        1. Per-channel vs per-block granularity: $V$ 使用 per-channel（per-column）scale factor 是精度保持的关键——当 $V_r$ 缩窄到 4.5 时，per-channel scaling 确保每列独立利用有限量化范围。若使用 per-tensor scaling，4.5 的范围会严重损失低幅值通道的信息。

          2. Code未公开细节: 论文发表时代码承诺 "will be available"。基于 SageAttention 仓库历史，实现预期在 csrc/ 目录下的 CUDA kernel 文件中，通过 sage_attn_cuda.cu 或类似入口注册为 PyTorch custom op。

          §8 Project Identity (Code-specific) #

          FieldValue
          Repo URL
          Primary languageCUDA, Python
          LicenseApache-2.0
          Sponsor orgTsinghua ML Group + Shengshu Tech
          Build systemPyTorch CUDA extension (setup.py / pip)
          Target HWNVIDIA Ada (SM89) + Blackwell consumer (RTX 5090)

          §9 Architecture & Module Map (Code-specific) #

          flowchart LR subgraph "SageAttention Library" API["Python API
          sage_attn()"] --> Dispatch["Dispatch Layer
          variant selection"] Dispatch --> QK_INT["QK Kernel
          INT4/INT8 matmul"] Dispatch --> PV_FP8["PV Kernel
          FP8 + FP16 acc"] Dispatch --> Softmax["Online Softmax
          + rescale"] QK_INT --> SMEM["Shared Memory
          Tiling (FlashAttn style)"] PV_FP8 --> SMEM PV_FP8 --> MMA["mma.f16.f8.f8.f16
          Tensor Core"] end User["User Code
          model.forward()"] --> API API --> Output["Attention Output
          (drop-in replacement)"]

          用户通过 sage_attn() Python 接口调用，内部根据配置 (4+8 / 8+8) 选择具体 kernel variant，kernel 内融合 QK matmul → softmax → PV matmul 全流程。

          §10 Critical Path Analysis (Code-specific) #

          Hot path: sage_attn() → CUDA kernel launch → per-tile loop:

          1. Load Q/K tiles → shared memory (HBM bandwidth bound)
          2. INT4/INT8 QK^T matmul → Tensor Core (compute bound, 2×-4× vs FP16)
          3. Online softmax → registers (compute, but lightweight vs matmul)
          4. Quantize P̃ → FP8 → registers (scale factor computation + round)
          5. Load V tile → shared memory
          6. Quantize V → FP8 → per-channel scale + round
          7. PV matmul via mma.f16.f8.f8.f16 → Tensor Core (HOT: 4× instruction throughput)
          8. Delayed FP32 Buffering → register promotion every 2 MMAs
          9. Dequantize + rescale → write O to HBM

            10. Step 7 是 SageAttn2++ 的核心优化点。Latency budget: PV matmul 占 attention 总时间的 ~40-50%（与 QK^T 大致对称），将其指令吞吐翻倍带来整体 ~1.3-1.5× 提升（从 SageAttn2 的 3× 到 3.9×）。

            §11 Performance Characteristics (Code-specific) #

            ConfigGPUheaddimSpeedup vs FA2Absolute (TOPS, est.)
            SageAttn2++(4+8)RTX40901283.9×~640
            SageAttn2++(8+8)RTX40901283.0×~490
            SageAttn2++(4+8)RTX5090128~3.5-4×N/A
            SageAttn2++(8+8)RTX5090128~3×N/A

            Scaling: 性能随序列长度增长而提升（1K → 32K），在 8K+ 趋于稳定——长序列充分利用 Tensor Core 占用率。短序列（1K）下加速比降至 ~2×，因 kernel launch 和 quantization overhead 占比增大。

            §12 Comparison with Alternatives (Code-specific) #

            FeatureSageAttn2++FlashAttention2FlashAttention3SageAttention2
            Speed (vs FA2)3.9×1×~1.5-2× (Hopper only)3×
            Accuracy99.97% cossimexactexact99.97% cossim
            GPU supportAda + Blackwell consumerAmpere+Hopper onlyAda+
            AccumulatorFP16FP32FP32FP32
            QK precisionINT4/INT8FP16FP8INT4/INT8
            PV precisionFP8 (narrowed range)FP16FP8FP8
            Drop-in?YesYesYesYes

            SageAttn2++ 的优势在消费级 GPU（RTX4090/5090）上最为突出，这些 GPU 没有 FlashAttention3 支持。在 Hopper GPU 上（H100/H200），FA3 是更强基线但 SageAttn2++ 的指令不一定可用。

            §13 Verdict & Recommendations (Code-specific) #

            When to adopt (Yes regime):

            • RTX4090/5090 上的 long-context inference（8K+ tokens）
            • 视频/图像生成模型推理（attention 占比高）
            • 需要 plug-and-play 加速且不可改模型架构的场景
            • 对 ~0.03% 精度损失可接受（几乎所有推理场景）

            When NOT to adopt (No regime):

            • Hopper/Blackwell datacenter GPU（H100/H200/B200）—— 有 FlashAttention3
            • 训练场景（backward pass 未实现）
            • 对数值精度极度敏感的应用（金融、科学计算）
            • $V$ 值分布极端（跨通道方差超大）的特殊架构

            Suggested contributions:

            1. 添加 backward pass 支持以扩展到 training
            2. 对 Wan (4+8) 的 VQA 降级做 root cause 分析（大概率是 INT4 QK 量化问题而非 PV 路径）
            3. 自动化 $(P_r, V_r)$ 选择——runtime 检测 $V$ 实际分布后动态调整

              4. §14 Concurrency & Memory (Code-specific) #

              Concurrency model: CUDA kernel, grid-level parallelism over (batch × heads × Q-tiles). Warp-level cooperative via Tensor Core MMA instructions.

              Memory management:

              • Shared memory: tiled blocks of Q, K, P, V (FlashAttention tiling scheme)
              • Registers: FP16 accumulator (16-bit vs FP32's 32-bit → lower register pressure → potentially higher occupancy)
              • HBM: source tensors Q, K, V and output O; no intermediate P matrix materialized (fused kernel)

              Register pressure advantage: FP16 accumulator 使用的寄存器位宽仅为 FP32 的一半。对于 register-bound 的 attention kernel，这可能允许更大的 tile size 或更高的 warp occupancy。论文未量化此二阶效应。

              核心技术壁垒: 对 PTX mma.m16n8k32 指令语义的精确建模（k=32 累加维度 × FP16 range = 硬约束），结合 post-softmax attention weight 的值域特性（$P \in [0,1]$ → 量化到 $[0, P_r]$），找到不损精度的非对称范围分配。这不是算法创新而是工程洞察——需要同时理解 ISA spec、数值表示、和 transformer attention 的统计特性。

              关键实现细节:

              1. 选择 $P_r = 224 \gg V_r = 4.5$ 的非对称分配（而非 $P_r = V_r = 32$）暗示 $V$ 通道间 per-channel scaling 已经将有效范围压缩到很小，$V_r = 4.5$ 的 headroom 对精度无影响。
              2. Delayed FP32 Buffering 的 "每 2 次 MMA 转一次" 是 overhead vs. tightened-constraint 的最优 tradeoff——每 3 次转更少 overhead 但约束过紧（$P_r \times V_r \leq 682$），每 1 次转则无收益。