新智元報(bào)道

編輯：LRS

【新智元導(dǎo)讀】清華大學(xué)朱軍教授團(tuán)隊(duì)提出SageAttention3，利用FP4量化實(shí)現(xiàn)推理加速，比FlashAttention快5倍，同時(shí)探索了8比特注意力用于訓(xùn)練任務(wù)的可行性，在微調(diào)中實(shí)現(xiàn)了無損性能。

注意力機(jī)制是大模型的核心，能夠很好地捕捉上下文信息，但其復(fù)雜度會(huì)隨輸入長(zhǎng)度呈二次方增長(zhǎng)，導(dǎo)致了現(xiàn)有的生成式模型受到上下文窗口的限制，無法高效處理長(zhǎng)文本。

量化（Quantization）可以對(duì)模型推理過程進(jìn)行加速，用更低的權(quán)重精度換取計(jì)算效率，比如在英偉達(dá)新一代Blackwell架構(gòu)GPU中，F(xiàn)P4張量核心比FP16的計(jì)算性能要快得多。

不過FP4量化只有15個(gè)可表示的值，無論是按「張量量化」還是按「token量化」，都無法有效保留模型的準(zhǔn)確性。

注意力圖中的小值集中在[0, 1]范圍內(nèi)，直接量化為FP4會(huì)導(dǎo)致縮放因子（scaling factors）的動(dòng)態(tài)范圍極其狹窄，硬件上要求量化因子是FP8數(shù)據(jù)類型，會(huì)導(dǎo)致準(zhǔn)確率損失下降明顯。

現(xiàn)有的研究方向局限于「推理加速」，在訓(xùn)練過程中使用8比特注意力時(shí)，注意力圖的梯度特別容易受到量化誤差的影響，從而導(dǎo)致輸入梯度中累積誤差。

針對(duì)量化加速特性，清華大學(xué)朱軍教授團(tuán)隊(duì)發(fā)布了首個(gè)可用于推理的FP4注意力機(jī)制SageAttention3，也是首次探索「低比特注意力」在大型模型階段中加速的可行性。

論文鏈接：https://arxiv.org/pdf/2505.11594

代碼鏈接：https://github.com/thu-ml/SageAttention

SageAttention3將量化組大小限制為1x16，而不是按張量或通道進(jìn)行量化，可以有效避免塊內(nèi)異常值的影響，同時(shí)提高FP4量化的精度。

然后采用兩級(jí)量化方法，先通過按token量化將每個(gè)token的值范圍歸一化到[0, 448×6]，然后使用FP4微縮（microscaling）量化以提高精度，充分利用FP8縮放因子的表示范圍。

在反向傳播的五個(gè)矩陣乘法中，識(shí)別出最敏感的矩陣乘法，并將其精度保持在FP16，從而避免因量化誤差導(dǎo)致的精度損失。

SageAttention3，在RTX5090上實(shí)現(xiàn)了1038萬億次每秒運(yùn)算（TOPS），比FlashAttention快了5倍。

研究人員使用8比特可訓(xùn)練注意力機(jī)制（SageBwd）微調(diào)基礎(chǔ)模型時(shí)，在性能上沒有任何損失，但在預(yù)訓(xùn)練任務(wù)中并不適用。

FP4推理加速

微縮比例FP4注意力機(jī)制

對(duì)一個(gè)矩陣進(jìn)行量化處理，將其轉(zhuǎn)換為FP4格式，并使用FP8格式的縮放因子矩陣。

具體操作是：將矩陣X劃分為多個(gè)小塊，每個(gè)小塊對(duì)應(yīng)一個(gè)縮放因子。

量化（Quantization）過程是將矩陣的每個(gè)值除以縮放因子后進(jìn)行舍入，得到量化后的值；

反量化（Dequantization）則是將量化后的值乘以縮放因子，恢復(fù)為近似的原始值。

在矩陣乘法中，利用FP4微縮比例量化來加速計(jì)算：與傳統(tǒng)的FP16精度矩陣乘法（200 TOPS）相比，F(xiàn)P4微縮比例矩陣乘法的速度（1600TOPS）可以提升8倍。

實(shí)現(xiàn)方式為FP4MM乘法指令，輸入為兩個(gè)量化后的矩陣和兩個(gè)縮放因子，輸出為矩陣乘法的結(jié)果。

在注意力計(jì)算時(shí)，研究人員通過將FP4微縮比例量化應(yīng)用于「查詢矩陣和鍵矩陣的乘法」和「中間矩陣P和值矩陣的乘法」來加速注意力計(jì)算。

先將查詢矩陣和鍵矩陣分別量化為FP4格式，并計(jì)算縮放因子；使用FP4微縮比例矩陣乘法指令計(jì)算查詢矩陣和鍵矩陣的乘積，得到中間結(jié)果；對(duì)中間結(jié)果應(yīng)用在線softmax操作，得到中間矩陣P；

將P矩陣和值矩陣分別量化為FP4格式，并計(jì)算縮放因子；再次使用FP4微縮比例矩陣乘法指令計(jì)算P矩陣和值矩陣的乘積，得到最終的輸出。

硬件實(shí)現(xiàn)基于FlashAttention技術(shù)，并采用平滑技術(shù)來提高注意力的準(zhǔn)確性。

FP4數(shù)據(jù)類型中，研究人員選擇NVFP4，因?yàn)槠湓谧⒁饬α炕械木冗h(yuǎn)高于MXFP4；通過實(shí)驗(yàn)驗(yàn)證，NVFP4在CogVideoX模型的所有層中表現(xiàn)出了更高的精度。

兩級(jí)縮放

直接對(duì)中間矩陣P使用FP4微縮比例量化會(huì)導(dǎo)致精度損失，因?yàn)榭s放因子需要從FP32轉(zhuǎn)換為E4M3格式，會(huì)降低精度。

先將矩陣P的每一行的值范圍擴(kuò)展到一個(gè)更大的區(qū)間，以充分利用E4M3的表示范圍，然后對(duì)擴(kuò)展后的矩陣使用標(biāo)準(zhǔn)的FP4量化。

通過這種兩級(jí)量化方法，能夠減少縮放因子的數(shù)值表示誤差和矩陣的量化誤差，從而提高整體精度。

硬件實(shí)現(xiàn)和優(yōu)化

在FP4矩陣乘法（MatMul）中，F(xiàn)P32累加器的內(nèi)存布局和操作數(shù)A的寄存器布局不一樣，如果強(qiáng)行對(duì)齊，會(huì)降低內(nèi)核的性能。

研究人員的解決辦法是調(diào)整P塊的列順序，同時(shí)為了保證矩陣乘法的正確性，也相應(yīng)地調(diào)整了K的列順序，整個(gè)過程可以和量化操作一起完成。

在對(duì)P進(jìn)行微縮比例量化時(shí)，需要找到16個(gè)連續(xù)行元素的最大值，但這16個(gè)元素分布在四個(gè)線程中，需要先在每個(gè)線程內(nèi)進(jìn)行最大值計(jì)算，然后在線程之間進(jìn)行shuffle操作，會(huì)顯著降低內(nèi)核的速度。

通過將量化與在線softmax操作融合可以優(yōu)化該過程，先計(jì)算S矩陣中16個(gè)元素的最大值，并在后續(xù)的softmax最大值計(jì)算中重用最大值，可以減少50%的冗余shuffle和最大值計(jì)算操作，可以實(shí)現(xiàn)大約10%的整體內(nèi)核加速。

用于訓(xùn)練的INT8注意力

低比特量化注意力機(jī)制（例如FlashAttention3和SageAttention）目前主要用于推理階段，研究人員提出了一種適用于訓(xùn)練階段的INT8注意力機(jī)制（SageBwd），可以將注意力機(jī)制中的七個(gè)矩陣乘法中的六個(gè)量化為INT8格式，同時(shí)在微調(diào)任務(wù)中沒有性能下降。

前向

在注意力機(jī)制的前向傳播過程中，有兩個(gè)關(guān)鍵的矩陣乘法操作。

對(duì)P采用逐token量化（每個(gè)token單獨(dú)量化），對(duì)V采用逐塊量化，可以提高注意力機(jī)制的精度；利用在線softmax計(jì)算過程中已經(jīng)得到的最大值，可以避免重復(fù)計(jì)算最大值。

簡(jiǎn)單來說，對(duì)于每個(gè)FlashAttention塊，先找到塊中所有數(shù)值的最大絕對(duì)值，然后除以127，得到縮放因子；再用該縮放因子去量化塊中的值。

后向

對(duì)關(guān)于值（V）的梯度部分進(jìn)行量化，會(huì)對(duì)查詢（Q）和鍵（K）的梯度精度產(chǎn)生很大影響，主要是因?yàn)椴僮鞯木戎苯佑绊懙胶罄m(xù)的計(jì)算，而誤差會(huì)在反向傳播的過程中不斷累積，尤其是當(dāng)處理很長(zhǎng)的序列時(shí)，誤差會(huì)越來越大。

為了避免這個(gè)問題，研究人員決定不對(duì)該操作進(jìn)行量化，而是保持其使用更高精度的格式（FP16），同時(shí)對(duì)其他四個(gè)操作使用低精度（INT8）量化來加速計(jì)算。

這種方式既提高了計(jì)算效率，又保證了梯度計(jì)算的準(zhǔn)確性。

實(shí)驗(yàn)結(jié)果

SageAttention3在RTX5090上的運(yùn)行速度比FlashAttention快4到5倍，比xformers快8到11倍，同時(shí)在各種模型中保持了端到端的性能指標(biāo)。

在RTX4090上，使用SageBwd和基線方法測(cè)試前向加反向傳播速度，結(jié)果表明，SageBwd比FlashAttention2最多快1.67倍，比用Triton實(shí)現(xiàn)的FlashAttention2和xformers速度提升更高。

端到端性能損失

為了評(píng)估SageBwd在訓(xùn)練任務(wù)中的有效性，研究人員進(jìn)行了兩組實(shí)驗(yàn)。

在GSM8K、DROP、MMLU和HELLASWAG數(shù)據(jù)集上對(duì)Qwen2.5（3B）和Llama3.2（1B的基礎(chǔ)模型進(jìn)行了微調(diào)，損失結(jié)果表明，SageBwd與BF16完全一致。

圖(a)中，研究人員使用Llama（400M）模型在FineWebEdu數(shù)據(jù)集上進(jìn)行了預(yù)訓(xùn)練任務(wù)，損失曲線表明，雖然SageBwd能夠?qū)崿F(xiàn)損失收斂，但其收斂速度相對(duì)較慢，限制了其在預(yù)訓(xùn)練任務(wù)中的適用性。

研究人員還在多個(gè)測(cè)試數(shù)據(jù)集上評(píng)估了微調(diào)模型的答案質(zhì)量，結(jié)果表明SageBwd的性能與BF16相同。

SageAttention3在HunyuanVideo上進(jìn)行視頻生成和在Stable-Diffusion3.5上進(jìn)行圖像生成的對(duì)比表明，SageAttention3完全保持了生成質(zhì)量。

端到端加速實(shí)驗(yàn)結(jié)果顯示，SageAttention3在RTX5090上實(shí)現(xiàn)了大約3倍（HunyuanVideo）和2.4倍（CogVideoX）的端到端推理生成加速。

SageBwd在RTX4090上使用8K/16K token微批次時(shí)，將Llama（1B）的訓(xùn)練加速了大約1.15倍。

參考資料：

https://arxiv.org/abs/2505.11594