聞樂 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

所有用英偉達Blackwell B200的人，都在花冤枉錢？？

普林斯頓大學(xué)等聯(lián)合團隊指出，這款GPU居然因為軟硬件適配問題白白浪費了60%的計算資源

算力浪費了，咋辦呢——FlashAttention-4給出了答案。

這款專為Blackwell架構(gòu)GPU量身打造的注意力算法，一舉將利用率從行業(yè)普遍的20%-30%推至71%

FlashAttention-4由Tri Dao領(lǐng)銜、攜手Meta、Together AI等團隊共同研發(fā)。

嗯，英偉達自己也參與其中了

Blackwell B200有力使不出

英偉達Blackwell B200作為新一代數(shù)據(jù)中心GPU，其tensor core張量核心算力達到2.25 PFLOPS，是上一代Hopper H100的2倍。

理論上能讓注意力計算的速度實現(xiàn)跨越式提升。

但理想很豐滿……

這款GPU發(fā)生了嚴(yán)重的偏科。

核心算力猛增的同時，關(guān)鍵的配套計算單元卻原地踏步。

其中，負責(zé)指數(shù)運算的MUFU單元吞吐量與Hopper架構(gòu)完全一致，沒有任何提升；

共享內(nèi)存的帶寬也保持原樣，并未跟隨張量核心同步升級。

這一硬件設(shè)計的不對稱性，直接導(dǎo)致了性能瓶頸的反轉(zhuǎn)。

在大模型核心的注意力計算負載中，原本的性能瓶頸矩陣乘法，如今耗時遠低于輔助環(huán)節(jié)，共享內(nèi)存的讀寫操作和指數(shù)運算的耗時，反而比矩陣乘法多出25%-60%

算力翻倍的Tensor Core長期處于等待狀態(tài)，大量計算資源就這么被閑置了。

于是，大量開發(fā)者花費重金部署的B200 GPU，因核心算力與配套單元的脫節(jié)，超六成資源被白白浪費

算力翻倍？

No！明明是有力使不出……

FlashAttention-4三招破解瓶頸

針對Blackwell GPU的偏科問題，F(xiàn)lashAttention-4量身打造了三大優(yōu)化策略。

第一招，多管齊下化解指數(shù)運算與內(nèi)存讀寫難題。

團隊一方面通過軟件模擬指數(shù)函數(shù)，借助多項式近似的方法，讓高速的FMA計算單元參與到原本由MUFU單元負責(zé)的指數(shù)運算中，大幅提升指數(shù)計算的吞吐量；

同時通過混合硬件計算與軟件模擬的方式，在提速的同時保證計算精度。

另一方面推出條件性softmax rescaling策略，僅在必要時執(zhí)行softmax的縮放操作，直接跳過大量無用的計算步驟，減少非矩陣乘法的運算量。

此外，團隊充分利用Blackwell架構(gòu)的2-CTA MMA模式，讓兩個計算單元搭檔完成矩陣運算，各自僅加載一半的運算數(shù)據(jù)。

這就將共享內(nèi)存的讀寫量直接砍半，同時還減少了后續(xù)的原子操作，從根源上緩解共享內(nèi)存的帶寬壓力。

第二招，重構(gòu)計算流水線，實現(xiàn)算力的并行最大化。

FlashAttention-4深度適配Blackwell架構(gòu)的全異步MMA操作和新增的張量內(nèi)存TMEM，重新設(shè)計了注意力計算的前向和反向流水線。

讓softmax計算與矩陣乘法這兩個核心環(huán)節(jié)實現(xiàn)完全的計算重疊。

當(dāng)硬件的張量核心在處理一個矩陣塊時，另一部分硬件資源可同時對另一個數(shù)據(jù)塊執(zhí)行softmax計算，避免硬件算力的空閑。

第三招，兼顧硬件迭代，為下一代GPU預(yù)留優(yōu)化空間。

研發(fā)團隊同時考慮到Blackwell架構(gòu)的硬件升級趨勢，目前B300/GB300 GPU的指數(shù)運算單元吞吐量已翻倍至32 ops/clock/SM

針對這一變化，團隊明確表示，F(xiàn)lashAttention-4當(dāng)前的軟件模擬指數(shù)運算方案，在下一代硬件上會根據(jù)實際性能表現(xiàn)重新權(quán)衡，確保算法能持續(xù)適配硬件的迭代升級。

告別 C++，編譯速度狂飆30倍

除了算法層的深度優(yōu)化，F(xiàn)lashAttention-4在開發(fā)層面也帶來了變化。

與此前基于C++模板開發(fā)的FlashAttention-3不同，FlashAttention-4的全部代碼基于Python的領(lǐng)域?qū)Ｓ冒姹綜uTe-DSL框架編寫，實現(xiàn)了零C++代碼開發(fā)

這一設(shè)計帶來的是編譯的效率躍升。

前向傳播內(nèi)核的編譯時間從FlashAttention-3的55秒縮短至2.5秒，提速22倍；

反向傳播的編譯時間從45秒降至1.4秒，提速32倍，整體編譯速度最高狂飆30倍。

在B200 GPU上的實測數(shù)據(jù)顯示，其前向傳播算力最高達到1613 TFLOPS/s，一舉實現(xiàn)71%的理論峰值利用率。

對比主流的計算框架，F(xiàn)lashAttention-4的優(yōu)勢也比較明顯。

比英偉達官方的cuDNN 9.13快1.1-1.3倍，比常用的Triton框架快2.1-2.7 倍。

且在長序列、因果掩碼等大模型訓(xùn)練推理的核心場景中，性能優(yōu)勢更為突出。

One More Thing

論文還指出，cuDNN從9.13版本開始就已經(jīng)開始反向吸收了FA4的核心技術(shù)。

看來，英偉達自己也忍不住抄作業(yè)了（doge）。

論文地址：https://arxiv.org/abs/2603.05451
參考鏈接：https://x.com/alex_prompter/status/2033885345935462853?s=20

— 完 —