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

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

• 論文標(biāo)題：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

效果預(yù)覽

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

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

視頻 1（使用 FlashAttention2，490s）        

視頻 2（使用 SageAttention3,164s）

（注：FlashAttention2 已經(jīng)是在 RTX5090 上最優(yōu)的 FlashAttention 實現(xiàn)。）

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

SageAttention3 總體流程圖

前言

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

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

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

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

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

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

8-Bit 注意力用于訓(xùn)練有什么問題？

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

技術(shù)方案

為了解決上述的挑戰(zhàn)，研究團(tuán)隊提出了對應(yīng)的解決辦法。

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

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

于是研究團(tuán)隊決定先把 P 通過 Per-token 量化到 [0, ] 的范圍內(nèi)，再進(jìn)行 FP4 的量化：

下表展示了 Two-Level Scaling 對精度的提升：

下圖展示了 SageAttention3 的算法流程：

（3）在 8-Bit 訓(xùn)練 Attention 當(dāng)中，研究團(tuán)隊對 Q，K，V 采用了 Per-block INT8 量化，對 P 巧妙地采用了無量化 Overhead 的 Per-token 量化。前向過程的算法如下：

在反向傳播的過程中總共涉及到 5 個矩陣乘法：

研究團(tuán)隊發(fā)現(xiàn)是否量化 dOVT 對精度有著較大的影響：

于是研究團(tuán)隊將 dOVT 保留為 FP16 精度，而對其它四個矩陣乘法進(jìn)行了量化。以下是反向傳播的算法：

實驗效果

SageAttention3 實現(xiàn)了 GPU 底層的 CUDA Kernel，在算子速度以及各個模型端到端準(zhǔn)確度上都有十分不錯的表現(xiàn)。

具體來說，算子速度相比于 FlashAttention2（5090 上最快的 FlashAttention） 和 xformers 有大約 5 倍以及 10 倍的加速：

各模型在真實場景的端到端精度表現(xiàn)中，在視頻、圖像生成等大模型上均保持了端到端的精度表現(xiàn)：

下圖是在 HunyuanVideo 當(dāng)中的可視化實例：

下圖是在 Flux 上的可視化實例：

下圖是在 Cogvideo 中的可視化實例：

下表展示了各個視頻、圖像生成模型中 SageAttention3 的端到端精度表現(xiàn)：

端到端的速度表現(xiàn)上，SageAttention3 的實現(xiàn)均可以有效地對長序列的模型進(jìn)行加速，比如可以端到端 3 倍加速 HunyuanVideo：

8-Bit 訓(xùn)練 Attention 在 Base Model 微調(diào)到 Instruct Model 的任務(wù)上展現(xiàn)出與 BF16 的注意力完全一致的精度表現(xiàn)，下表是在多個不同的任務(wù)以及模型上微調(diào)的結(jié)果：

并且在訓(xùn)練速度上也能起到較好的加速效果：

研究團(tuán)隊還發(fā)現(xiàn)，目前的 8 比特用于訓(xùn)練的 Attention 雖然在微調(diào)任務(wù)上完全無損，但是在預(yù)訓(xùn)練任務(wù)上與全精度的 Attention 在 Loss 上還有一定差距，需要未來進(jìn)一步的研究：