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

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

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

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

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

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

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

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

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

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

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

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

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

FP4推理加速

微縮比例FP4注意力機制

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

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

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

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

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

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

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

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

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

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

FP4數(shù)據(jù)類型中，研究人員選擇NVFP4，因為其在注意力量化中的精度遠(yuǎn)高于MXFP4；通過實驗驗證，NVFP4在CogVideoX模型的所有層中表現(xiàn)出了更高的精度。

兩級縮放

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

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

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

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

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

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

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

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

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

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

前向

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

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

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

后向

對關(guān)于值（V）的梯度部分進(jìn)行量化，會對查詢（Q）和鍵（K）的梯度精度產(chǎn)生很大影響，主要是因為操作的精度直接影響到后續(xù)的計算，而誤差會在反向傳播的過程中不斷累積，尤其是當(dāng)處理很長的序列時，誤差會越來越大。

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

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

實驗結(jié)果

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

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

端到端性能損失

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

在GSM8K、DROP、MMLU和HELLASWAG數(shù)據(jù)集上對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)損失收斂，但其收斂速度相對較慢，限制了其在預(yù)訓(xùn)練任務(wù)中的適用性。

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

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

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

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