清華SageAttention3，FP4量化5倍加速！且首次支持8比特訓練

2025-06-19 09:07:00

近日，清華大學陳鍵飛團隊進一步提出了針對 BlackWell 架構的首個全 FP4 量化的即插即用注意力算子（SageAttention3）。

隨著大型模型需要處理的序列長度不斷增加，注意力運算（Attention）的時間開銷逐漸成為主要開銷。此前，清華大學陳鍵飛團隊提出的即插即用的 SageAttention 和 SageAttention2 已經被業界及社區廣泛的使用于各種開源及商業的大模型中，比如 Vidu，CogvideoX，Mochi，Wan，HunyuanVideo，Flux，Llama3，Qwen 等。

近日，清華大學陳鍵飛團隊進一步提出了針對 BlackWell 架構的首個全 FP4 量化的即插即用注意力算子（SageAttention3）。實現了 5 倍相比于 FlashAttention 的即插即用的推理加速（此前的 SageAttention V1/V2/V2++ 分別達到了 2.1，3，3.9 倍的加速效果），比如在 RTX 5090 上，SageAttention3 達到了 1040 TOPS 的速度，甚至是比 RTX 5090 昂貴十幾倍的 H100 上使用 Hopper 獨有的 FlashAttention3 還要快 1.65 倍！SageAttention3 在多種視頻和圖像生成等大模型上（包括 HunyuanVideo，CogVideoX，Mochi 和各類圖像生成模型）均保持了端到端的精度表現。同時還首次提出可訓練的 8 比特注意力（SageBwd）用于大模型的訓練加速（注：FlashAttention3 的 FP8 版本也只支持前向傳播），在各項微調任務中均保持了與全精度注意力相同的結果。

論文標題：SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-bit Training
論文鏈接：https://arxiv.org/abs/2505.11594
開源代碼：https://github.com/thu-ml/SageAttention

效果預覽

SageAttention3 實現了高效的 Attention 算子，可以實現即插即用的 5 倍于 FlashAttention 的推理加速。即輸入任意 Q, K, V 矩陣，可以快速返回 Attention Output (O)，真正做到了兩行代碼加速任意模型推理。（注：按照官方倉庫中的開源計劃，SageAttention2++ 的代碼將于6月20日左右開源，SageAttention3 的代碼將于7月15日左右開源。）

效果上，以 HunyuanVideo 為例，在 RTX5090 上 SageAttention3 可以 3 倍加速端到端的生成，且視頻質量幾乎無損：

視頻 1（使用 FlashAttention2，490s）

視頻 2（使用 SageAttention3,164s）

（注：FlashAttention2 已經是在 RTX5090 上最優的 FlashAttention 實現。）

接下來，將從前言，挑戰，方法，以及實驗效果四個方面介紹 SageAttention3。

SageAttention3 總體流程圖

前言

隨著大模型需要處理的序列長度越來越長，Attention 的速度優化變得越來越重要。下圖展示了一個標準的 Transformer 模型中各運算的時間占比隨序列長度的變化：

為了方便指代注意力運算中的矩陣，我們先回顧一下注意力的計算公式：

研究動機：（1）Blackwell 架構有著速度極快的 FP4 Tensor Core，以 RTX5090 為例，其速度是 FP16 Tensor Core 的 8 倍。（2）訓練階段的注意力運算開銷也同樣重要，在此之前并沒有工作嘗試過低比特注意力加速模型訓練，包括 FlashAttention3 的 FP8 版本也只有 Forward 過程。我們還希望同時量化注意力的前向 + 反向過程來加速訓練。

FP4 注意力量化有什么問題？

（1）FP4 數值類型僅有 15 個有效數值，這使得以 Tensor（Per-tensor）或以 Token（Per-token）粒度的量化都難以有效保證量化的準確度。

（2）P 矩陣的值分布在 [0, 1] 之間，直接的 FP4 量化會使量化縮放因子被限制在一個狹窄的范圍內。然而，硬件要求這些量化因子必須采用 FP8 數據類型表示。此時，將縮放因子轉為 FP8 時會導致顯著的精度損失。

8-Bit 注意力用于訓練有什么問題？

（1）P 矩陣的梯度對量化誤差過于敏感，并且在反向過程中還會沿著序列長度對 Q 和 K 的梯度造成誤差累積。

技術方案

為了解決上述的挑戰，研究團隊提出了對應的解決辦法。

（1）為了提高 FP4 的量化精度。研究團隊采用了 Microscaling FP4 量化，這是 BlackWell 硬件層面支持的一種量化方式。即可以采用或的量化粒度進行矩陣量化，NIVIDA 在硬件層面自動支持了反量化過程。此外，Microscaling FP4 有兩種數據表示的形式，一種是 MXFP4, 另外一種是 NVFP4。兩種格式都采用了 E2M1 的 FP4 數據類型。不同的是，NVFP4 的量化的塊大小為，縮放因子的數據類型為 E4M3。MXFP4 的量化的塊大小為，縮放因子的數據格式為 E8M0。研究團隊采用了 NVFP4 數據格式，因為其量化準確率遠高于 MXFP4：

（2）針對 P 的縮放因子范圍狹窄的問題，研究團隊提出了兩階段量化（Two-level Quantization）的辦法。FlashAttention 中的 P 矩陣的值在 [0, 1] 的范圍內，導致 P 的縮放因子的范圍也只在 0~0.167 之間。把縮放因子直接轉換為 FP8 格式會帶來極大的精度損失。