清華髮布全新注意力SageAttention3：提速5倍，精度不降，訓推都能用！

2025-07-23 00:11 CVer

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達
點選進入—>【頂會/頂刊】投稿交流群

新增微訊號：CVer2233，小助手會拉你進群！
掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

轉載自：新智元 ｜ 編輯：LRS

【導讀】清華大學朱軍教授團隊提出SageAttention3，利用FP4量化實現推理加速，比FlashAttention快5倍，同時探索了8位元注意力用於訓練任務的可行性，在微調中實現了無損效能。
注意力機制是大模型的核心，能夠很好地捕捉上下文資訊，但其複雜度會隨輸入長度呈二次方增長，導致了現有的生成式模型受到上下文視窗的限制，無法高效處理長文字。
量化（Quantization）可以對模型推理過程進行加速，用更低的權重精度換取計算效率，比如在英偉達新一代Blackwell架構GPU中，FP4張量核心比FP16的計算效能要快得多。
不過FP4量化只有15個可表示的值，無論是按「張量量化」還是按「token量化」，都無法有效保留模型的準確性。
注意力圖中的小值集中在[0, 1]範圍內，直接量化為FP4會導致縮放因子（scaling factors）的動態範圍極其狹窄，硬體上要求量化因子是FP8資料型別，會導致準確率損失下降明顯。
現有的研究方向侷限於「推理加速」，在訓練過程中使用8位元注意力時，注意力圖的梯度特別容易受到量化誤差的影響，從而導致輸入梯度中累積誤差。
針對量化加速特性，清華大學朱軍教授團隊釋出了首個可用於推理的FP4注意力機制SageAttention3，也是首次探索「低位元注意力」在大型模型階段中加速的可行性。
論文連結：https://arxiv.org/pdf/2505.11594
程式碼連結：https://github.com/thu-ml/SageAttention
SageAttention3將量化組大小限制為1×16，而不是按張量或通道進行量化，可以有效避免塊內異常值的影響，同時提高FP4量化的精度。
然後採用兩級量化方法，先透過按token量化將每個token的值範圍歸一化到[0, 448×6]，然後使用FP4微縮（microscaling）量化以提高精度，充分利用FP8縮放因子的表示範圍。
在反向傳播的五個矩陣乘法中，識別出最敏感的矩陣乘法，並將其精度保持在FP16，從而避免因量化誤差導致的精度損失。
SageAttention3，在RTX5090上實現了1038萬億次每秒運算（TOPS），比FlashAttention快了5倍。
研究人員使用8位元可訓練注意力機制（SageBwd）微調基礎模型時，在效能上沒有任何損失，但在預訓練任務中並不適用。

FP4推理加速

微縮比例FP4注意力機制
對一個矩陣進行量化處理，將其轉換為FP4格式，並使用FP8格式的縮放因子矩陣。
具體操作是：將矩陣X劃分為多個小塊，每個小塊對應一個縮放因子。
量化（Quantization）過程是將矩陣的每個值除以縮放因子後進行舍入，得到量化後的值；
反量化（Dequantization）則是將量化後的值乘以縮放因子，恢復為近似的原始值。
在矩陣乘法中，利用FP4微縮比例量化來加速計算：與傳統的FP16精度矩陣乘法（200 TOPS）相比，FP4微縮比例矩陣乘法的速度（1600TOPS）可以提升8倍。
實現方式為FP4MM乘法指令，輸入為兩個量化後的矩陣和兩個縮放因子，輸出為矩陣乘法的結果。
在注意力計算時，研究人員透過將FP4微縮比例量化應用於「查詢矩陣和鍵矩陣的乘法」和「中間矩陣P和值矩陣的乘法」來加速注意力計算。
先將查詢矩陣和鍵矩陣分別量化為FP4格式，並計算縮放因子；使用FP4微縮比例矩陣乘法指令計算查詢矩陣和鍵矩陣的乘積，得到中間結果；對中間結果應用線上softmax操作，得到中間矩陣P；
將P矩陣和值矩陣分別量化為FP4格式，並計算縮放因子；再次使用FP4微縮比例矩陣乘法指令計算P矩陣和值矩陣的乘積，得到最終的輸出。
硬體實現基於FlashAttention技術，並採用平滑技術來提高注意力的準確性。
FP4資料型別中，研究人員選擇NVFP4，因為其在注意力量化中的精度遠高於MXFP4；透過實驗驗證，NVFP4在CogVideoX模型的所有層中表現出了更高的精度。

兩級縮放
直接對中間矩陣P使用FP4微縮比例量化會導致精度損失，因為縮放因子需要從FP32轉換為E4M3格式，會降低精度。
先將矩陣P的每一行的值範圍擴充套件到一個更大的區間，以充分利用E4M3的表示範圍，然後對擴充套件後的矩陣使用標準的FP4量化。
透過這種兩級量化方法，能夠減少縮放因子的數值表示誤差和矩陣的量化誤差，從而提高整體精度。

硬體實現和最佳化
在FP4矩陣乘法（MatMul）中，FP32累加器的記憶體佈局和運算元A的暫存器佈局不一樣，如果強行對齊，會降低核心的效能。
研究人員的解決辦法是調整P塊的列順序，同時為了保證矩陣乘法的正確性，也相應地調整了K的列順序，整個過程可以和量化操作一起完成。
在對P進行微縮比例量化時，需要找到16個連續行元素的最大值，但這16個元素分佈在四個執行緒中，需要先在每個執行緒內進行最大值計算，然後線上程之間進行shuffle操作，會顯著降低核心的速度。
透過將量化與線上softmax操作融合可以最佳化該過程，先計算S矩陣中16個元素的最大值，並在後續的softmax最大值計算中重用最大值，可以減少50%的冗餘shuffle和最大值計算操作，可以實現大約10%的整體核心加速。

用於訓練的INT8注意力
低位元量化注意力機制（例如FlashAttention3和SageAttention）目前主要用於推理階段，研究人員提出了一種適用於訓練階段的INT8注意力機制（SageBwd），可以將注意力機制中的七個矩陣乘法中的六個量化為INT8格式，同時在微調任務中沒有效能下降。
前向
在注意力機制的前向傳播過程中，有兩個關鍵的矩陣乘法操作。
對P採用逐token量化（每個token單獨量化），對V採用逐塊量化，可以提高注意力機制的精度；利用線上softmax計算過程中已經得到的最大值，可以避免重複計算最大值。
簡單來說，對於每個FlashAttention塊，先找到塊中所有數值的最大絕對值，然後除以127，得到縮放因子；再用該縮放因子去量化塊中的值。
後向
對關於值（V）的梯度部分進行量化，會對查詢（Q）和鍵（K）的梯度精度產生很大影響，主要是因為操作的精度直接影響到後續的計算，而誤差會在反向傳播的過程中不斷累積，尤其是當處理很長的序列時，誤差會越來越大。
為了避免這個問題，研究人員決定不對該操作進行量化，而是保持其使用更高精度的格式（FP16），同時對其他四個操作使用低精度（INT8）量化來加速計算。
這種方式既提高了計算效率，又保證了梯度計算的準確性。

實驗結果
SageAttention3在RTX5090上的執行速度比FlashAttention快4到5倍，比xformers快8到11倍，同時在各種模型中保持了端到端的效能指標。
在RTX4090上，使用SageBwd和基線方法測試前向加反向傳播速度，結果表明，SageBwd比FlashAttention2最多快1.67倍，比用Triton實現的FlashAttention2和xformers速度提升更高。

端到端效能損失
為了評估SageBwd在訓練任務中的有效性，研究人員進行了兩組實驗。
在GSM8K、DROP、MMLU和HELLASWAG資料集上對Qwen2.5（3B）和Llama3.2（1B的基礎模型進行了微調，損失結果表明，SageBwd與BF16完全一致。
圖(a)中，研究人員使用Llama（400M）模型在FineWebEdu資料集上進行了預訓練任務，損失曲線表明，雖然SageBwd能夠實現損失收斂，但其收斂速度相對較慢，限制了其在預訓練任務中的適用性。
研究人員還在多個測試資料集上評估了微調模型的答案質量，結果表明SageBwd的效能與BF16相同。
SageAttention3在HunyuanVideo上進行影片生成和在Stable-Diffusion3.5上進行影像生成的對比表明，SageAttention3完全保持了生成質量。
端到端加速實驗結果顯示，SageAttention3在RTX5090上實現了大約3倍（HunyuanVideo）和2.4倍（CogVideoX）的端到端推理生成加速。
SageBwd在RTX4090上使用8K/16K token微批次時，將Llama（1B）的訓練加速了大約1.15倍。
參考資料：
https://arxiv.org/abs/2505.11594

何愷明在MIT授課的課件PPT下載
在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

ICCV 2025 論文和程式碼下載

在CVer公眾號後臺回覆：ICCV2025，即可下載ICCV 2025論文和程式碼開源的論文合

CVPR 2025 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2025，即可下載CVPR 2025論文和程式碼開源的論文合集
CV垂直方向和論文投稿交流群成立
掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）人數破萬！如果你想要了解最新最快最好的CV/DL/AI論文、實戰專案、行業前沿、從入門到精通學習教程等資料，一定要掃描下方二維碼，加入CVer知識星球！最強助力你的科研和工作！

▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請點贊和在看


相關文章