英偉達(dá)最強(qiáng)B200算力浪費(fèi)60%！普林斯頓團(tuán)隊(duì)出手，利用率升至71%

聞樂(lè) 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

所有用英偉達(dá)Blackwell B200的人，都在花冤枉錢(qián)？？

普林斯頓大學(xué)等聯(lián)合團(tuán)隊(duì)指出，這款GPU居然因?yàn)?strong>軟硬件適配問(wèn)題白白浪費(fèi)了60%的計(jì)算資源

算力浪費(fèi)了，咋辦呢——FlashAttention-4給出了答案。

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

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

嗯，英偉達(dá)自己也參與其中了

Blackwell B200有力使不出

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

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

但理想很豐滿(mǎn)……

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

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

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

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

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

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

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

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

算力翻倍？

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

FlashAttention-4三招破解瓶頸

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

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

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

同時(shí)通過(guò)混合硬件計(jì)算與軟件模擬的方式，在提速的同時(shí)保證計(jì)算精度。

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

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

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

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

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

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

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

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

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

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

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

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

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

這一設(shè)計(jì)帶來(lái)的是編譯的效率躍升。

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

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

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

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

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

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

One More Thing

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

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

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

— 完 —