第一次在H100上碰FP8 GEMM的時(shí)候我差點(diǎn)被一整套奇怪的數(shù)值結(jié)果勸退矩陣乘出來(lái)不是全0就是NaN偶爾正常一次結(jié)果又對(duì)不上參考實(shí)現(xiàn)。后來(lái)把關(guān)注點(diǎn)從“FP8就是少了一半精度”轉(zhuǎn)到“更底層的數(shù)據(jù)管理和流水線調(diào)度”上才慢慢理解像DeepGEMM這類(lèi)庫(kù)為什么值得拆開(kāi)看。DeepGEMM是一個(gè)面向FP8精度的輕量級(jí)GEMM庫(kù)專(zhuān)門(mén)處理稠密矩陣乘和MoE分段矩陣乘代碼量不大卻把從數(shù)值縮放到Hopper架構(gòu)底層調(diào)度的關(guān)鍵點(diǎn)全部串了起來(lái)。如果你在做推理加速、訓(xùn)練框架底層開(kāi)發(fā)或者只是想搞懂FP8內(nèi)核背后的原理這篇東西應(yīng)該能給你省不少時(shí)間。1. 從FP16到FP8它卡住了不少團(tuán)隊(duì)1.1 硬件算力翻倍但軟件沒(méi)跟上GPU的浮點(diǎn)演進(jìn)其實(shí)和寫(xiě)代碼的懶惰程度是矛盾的。老框架里到處是FP16的假設(shè)算子注冊(cè)表、自動(dòng)混合精度策略、自定義CUDA kernel的模板參數(shù)全都圍著FP16轉(zhuǎn)。換FP8不是把某個(gè)dtype字符串改掉就完事幾乎每個(gè)環(huán)節(jié)都要重新驗(yàn)證。但硬件層面的誘惑實(shí)在太大。Hopper架構(gòu)里FP8的峰值算力直接是FP16的兩倍內(nèi)存帶寬消耗也少一半。對(duì)大模型推理來(lái)說(shuō)這等同于在相同顯存和功耗下塞進(jìn)更多請(qǐng)求對(duì)訓(xùn)練來(lái)說(shuō)則代表更大的batch和更短的迭代時(shí)間。于是問(wèn)題變成既然硬件給得這么爽快為什么很多團(tuán)隊(duì)遲遲沒(méi)切換答案很簡(jiǎn)單軟件沒(méi)有跟上。FP8 GEMM不是“把float換成8位”這么樸素。乘積結(jié)果的范圍、累加過(guò)程中的精度、縮放因子的計(jì)算位置每一樣都會(huì)改變內(nèi)核的數(shù)據(jù)流。很多人上手之后發(fā)現(xiàn)要么精度崩掉要么性能還不如FP16最后只能退回老方案。DeepGEMM這類(lèi)庫(kù)的價(jià)值恰恰在于把這些問(wèn)題在庫(kù)內(nèi)部一次性解決好讓上層使用者不必重造輪子。1.2 大模型GEMM的兩個(gè)“反標(biāo)準(zhǔn)”形狀傳統(tǒng)GEMM庫(kù)的設(shè)計(jì)目標(biāo)是規(guī)整的大矩陣乘法M、N、K都很大用分塊策略把SM喂得飽飽的reduction維度盡可能長(zhǎng)。但大模型場(chǎng)景里的GEMM形狀往往并不規(guī)整。第一個(gè)典型場(chǎng)景是推理生成階段。M經(jīng)常只有1到16K和N可以到幾千甚至上萬(wàn)。這種細(xì)長(zhǎng)形狀會(huì)讓標(biāo)準(zhǔn)tile劃分產(chǎn)生大量空轉(zhuǎn)每個(gè)block算不了多少活就結(jié)束了SM的利用率被嚴(yán)重拉低。哪怕你有很高的FP8峰值算力實(shí)際吞吐也會(huì)非常慘。第二個(gè)典型場(chǎng)景是MoE。多個(gè)專(zhuān)家權(quán)重共享一份激活輸入每個(gè)專(zhuān)家分到的token數(shù)量不固定如果像常規(guī)做法那樣逐個(gè)專(zhuān)家單獨(dú)調(diào)GEMMkernel launch次數(shù)會(huì)爆炸單個(gè)小GEMM又湊不滿一個(gè)完整waveSM不斷閑置。DeepGEMM把這一類(lèi)問(wèn)題抽象成“分組/分段GEMM”用單次內(nèi)核調(diào)用處理多個(gè)專(zhuān)家的計(jì)算目的是讓本來(lái)細(xì)碎的工作變得足夠稠密。這個(gè)方向恰恰是很多通用庫(kù)沒(méi)有為FP8做精細(xì)優(yōu)化的死角。1.3 “小而精”的設(shè)計(jì)哲學(xué)很關(guān)鍵我最初看到DeepGEMM源碼時(shí)第一反應(yīng)是核心邏輯怎么這么短它并沒(méi)有把cuBLAS幾千行內(nèi)部邏輯重寫(xiě)一遍而是圍繞FP8 GEMM的關(guān)鍵路徑做了精準(zhǔn)裁剪。對(duì)于庫(kù)開(kāi)發(fā)者來(lái)說(shuō)代碼短不代表功能弱反而說(shuō)明每一行都有明確目的便于閱讀、審計(jì)和集成。這也是我認(rèn)為它值得學(xué)習(xí)的原因之一。很多高性能庫(kù)動(dòng)輒上萬(wàn)行內(nèi)部?jī)?yōu)化堆疊得密密麻麻你根本分不清哪些是架構(gòu)約束、哪些是歷史包袱。DeepGEMM把優(yōu)化集中在“FP8數(shù)據(jù)布局”“縮放因子處理”“Hopper異步流水線”“MoE分段調(diào)度”這幾個(gè)真正影響性能的點(diǎn)上剩下的保持樸素。這種極致的取舍對(duì)想自己動(dòng)手改內(nèi)核的人來(lái)說(shuō)是很好的教材。2. FP8數(shù)值細(xì)節(jié)決定整個(gè)庫(kù)的設(shè)計(jì)方向2.1 e4m3、e5m2怎么選FP8并不是一種格式而是兩種e4m3和e5m2。它們的差異集中在指數(shù)位和尾數(shù)位的分配上。我整理了一張常用的對(duì)照表格式指數(shù)位尾數(shù)位最大有限值適用場(chǎng)景FP16510約65504通用訓(xùn)練/推理BF1687約3.39e38訓(xùn)練中間計(jì)算FP8 e4m343約448激活、權(quán)重FP8 e5m252約57344梯度、誤差較大數(shù)據(jù)e4m3因?yàn)槲矓?shù)只有3位精度很粗糙但動(dòng)態(tài)范圍在4位指數(shù)下也能覆蓋常見(jiàn)的激活分布e5m2的指數(shù)多一位能表達(dá)更大的數(shù)值但尾數(shù)更少精度進(jìn)一步下降。工程里最常見(jiàn)的組合是激活和權(quán)重用e4m3梯度、誤差或者需要較大動(dòng)態(tài)范圍的數(shù)據(jù)用e5m2。這里有個(gè)很容易踩的坑不要想當(dāng)然地覺(jué)得所有張量都適合e4m3。如果數(shù)值分布本身很寬比如某些層的中間激活偶爾會(huì)冒出一個(gè)明顯偏大的channele4m3直接就會(huì)溢出到NaN。真正可靠的方案都是配合縮放因子做范圍壓縮而不是裸著硬上。2.2 scaling不是乘個(gè)系數(shù)那么簡(jiǎn)單FP8最容易被低估的部分是scaling。剛接觸時(shí)我也以為就是給每個(gè)張量乘一個(gè)全局系數(shù)把最大絕對(duì)值限制在448以內(nèi)。實(shí)際做進(jìn)去才發(fā)現(xiàn)per-tensor整張量一個(gè)scale在真實(shí)模型里經(jīng)常不夠用因?yàn)榧せ罘植疾⒉痪鶆蛞粋€(gè)離群特征會(huì)拉低全圖的有效精度。工程上更常見(jiàn)的是per-token結(jié)合per-channel的組合策略對(duì)激活A(yù)按token維度分組對(duì)權(quán)重B按channel維度分組甚至把K維度切塊后做更細(xì)粒度的分組量化。這樣雖然會(huì)多出一些scale計(jì)算的額外開(kāi)銷(xiāo)但能把量化誤差控制在可接受范圍。用個(gè)生活化的類(lèi)比per-tensor像用一把長(zhǎng)尺量全班身高班里有個(gè)兩米多的體育生你為了不爆尺只能把刻度拉得很寬結(jié)果普通同學(xué)的身高誤差全變大per-token / per-channel相當(dāng)于分組測(cè)量名正言順地按人群特征用不同的尺整體誤差自然更低。DeepGEMM在設(shè)計(jì)上支持這類(lèi)復(fù)雜的scale組合把縮放因子的傳遞和矩陣乘本身的流水線融合在一起。也就是說(shuō)scale不是前處理階段單獨(dú)算好再傳進(jìn)來(lái)而是內(nèi)核內(nèi)部就處理掉了一部分縮放邏輯避免額外讀寫(xiě)和kernel launch帶來(lái)的開(kāi)銷(xiāo)。2.3 累加精度與延遲縮放FP32寄存器里的故事FP8能承受的精度有限但GEMM有一個(gè)天然優(yōu)勢(shì)乘法可以用8位累加卻不需要停在8位。標(biāo)準(zhǔn)的做法是內(nèi)部累加器用FP32整個(gè)K維的乘積一路累加完最后寫(xiě)回時(shí)才落到目標(biāo)輸出精度。這樣即使輸入是8位累加誤差也不會(huì)在K很大時(shí)失控。另一個(gè)關(guān)鍵策略是延遲縮放。你可以把它理解成“用上一次統(tǒng)計(jì)好的scale先跑今天的計(jì)算”。訓(xùn)練場(chǎng)景里激活分布短時(shí)間不會(huì)突變沒(méi)必要每一步都重新算一遍scale適當(dāng)沿用上一輪的統(tǒng)計(jì)值就能穩(wěn)住精度同時(shí)省掉一次額外的數(shù)值統(tǒng)計(jì)kernel。在線縮放則適合動(dòng)態(tài)范圍變化很劇烈的場(chǎng)景兩者各有取舍。把這些細(xì)節(jié)串在一起你就會(huì)發(fā)現(xiàn)FP8 GEMM的性能高低不只是看乘法和加法本身還得看scale如何參與流水線、累加器放在哪一級(jí)、寫(xiě)回時(shí)有沒(méi)有多余的格式轉(zhuǎn)換。DeepGEMM在數(shù)值層面的每一個(gè)設(shè)計(jì)最終都是為了讓這些步驟盡量不打斷內(nèi)存與計(jì)算的重疊。3. 雙路徑實(shí)現(xiàn)Triton把功能跑通CUDA把性能吃透3.1 Triton路徑先把功能跑通Triton是一種面向GPU的Python式DSL寫(xiě)起來(lái)簡(jiǎn)潔社區(qū)普及率高很多項(xiàng)目用它做快速原型。DeepGEMM保留了Triton實(shí)現(xiàn)意味著普通工程師也能在不太懂CUDA細(xì)節(jié)的情況下把核心流程讀懂、改改參數(shù)、跑通自己的用例。Triton自動(dòng)化處理了很多瑣事tile劃分、共享內(nèi)存分配、基本的數(shù)據(jù)搬運(yùn)。對(duì)多數(shù)常規(guī)形狀來(lái)說(shuō)Triton路徑已經(jīng)能拿到可觀的性能因?yàn)榫幾g器會(huì)幫你做不少調(diào)度決策。它的主要限制是不能像手寫(xiě)CUDA那樣精細(xì)控制異步流水線的每個(gè)環(huán)節(jié)也很難針對(duì)某個(gè)具體GPU架構(gòu)做非常極致的寄存器和共享內(nèi)存布局優(yōu)化。所以可以把Triton路徑看作“大眾版”覆蓋廣、上手快、維護(hù)成本低。如果你只是想把一個(gè)模型從FP16遷移到FP8先跑Triton路徑驗(yàn)證數(shù)值是完全可行的一條路線。3.2 CUDA路徑把Hopper的潛力吃干凈真正要榨干Hopper架構(gòu)的算力還是得回到CUDA。Triton路徑解決的是“能跑”CUDA路徑解決的是“跑滿”。這部分的優(yōu)化點(diǎn)我列一下TMA批量拷貝Hopper上的Tensor Memory Accelerator可以異步把全局內(nèi)存里的矩陣塊搬到共享內(nèi)存搬運(yùn)過(guò)程不占用計(jì)算單元和計(jì)算流水線重疊。Warp Speciaization把warp分成producer和consumer兩類(lèi)一部分只負(fù)責(zé)搬運(yùn)數(shù)據(jù)另一部分只負(fù)責(zé)矩陣乘讓smem始終有數(shù)據(jù)可用消除流水線空泡。wgmma指令warp group級(jí)別的矩陣乘指令一個(gè)指令能完成更大塊的FP8乘法比早期mma指令更省指令數(shù)和寄存器壓力。動(dòng)態(tài)wave調(diào)度根據(jù)矩陣形狀和SM數(shù)量計(jì)算需要多少個(gè)block盡量讓每個(gè)SM都分配到合適的任務(wù)避免最后一個(gè)wave出現(xiàn)大量空閑。這些點(diǎn)單獨(dú)任何一個(gè)都不算新鮮但組合起來(lái)以后效果是明顯的GeMM的算術(shù)強(qiáng)度被拉高內(nèi)存延遲被隱藏SM的利用率穩(wěn)定在高位。DeepGEMM的CUDA路徑就是把這些東西按FP8的特點(diǎn)重新組織了一遍并且把代碼控制在緊湊范圍內(nèi)這是它最難得的地方。3.3 為什么保留兩條路徑我最初覺(jué)得保留兩套實(shí)現(xiàn)是重復(fù)勞動(dòng)后來(lái)在實(shí)際工程里才理解這樣做的好處。Triton路徑是維護(hù)和教學(xué)的最優(yōu)解CUDA路徑是極限性能的兜底選項(xiàng)。運(yùn)行時(shí)可以根據(jù)矩陣形狀、是否有Hopper特性可用、輸入是否規(guī)整等條件自動(dòng)選擇走哪條路。這種“雙軌”思路對(duì)很多團(tuán)隊(duì)都有參考價(jià)值不要一上來(lái)就追求手寫(xiě)CUDA先用高抽象路徑驗(yàn)證功能再對(duì)熱點(diǎn)場(chǎng)景用低層優(yōu)化做加固。兩條路徑可以長(zhǎng)期共存而不是二選一。DeepGEMM的調(diào)度機(jī)制本質(zhì)上也是這種漸進(jìn)式優(yōu)化思路的一次實(shí)踐。4. 復(fù)現(xiàn)與實(shí)測(cè)我看懂了為什么它敢叫高效4.1 環(huán)境準(zhǔn)備和三個(gè)最容易失敗的地方想復(fù)現(xiàn)DeepGEMM硬件門(mén)檻比較清楚Hopper架構(gòu)的GPU也就是通常說(shuō)的H100、H200這類(lèi)SM90設(shè)備。軟件方面需要較新的CUDA環(huán)境Triton也要跟得上版本PyTorch作為上層依賴基本跑不掉。我實(shí)際遇到過(guò)的三個(gè)失敗場(chǎng)景分享在這里編譯時(shí)CUDA arch信息沒(méi)配對(duì)導(dǎo)致報(bào)錯(cuò)或者生成出來(lái)的kernel根本起不了作用。解決辦法是先確認(rèn)torch.cuda.get_device_capability()輸出再設(shè)置對(duì)應(yīng)架構(gòu)不要偷懶用默認(rèn)值。Triton版本太舊JIT編譯內(nèi)核時(shí)直接炸。這個(gè)最隱蔽因?yàn)榘惭b依賴時(shí)不一定報(bào)錯(cuò)但運(yùn)行到特定路徑就會(huì)掛。建議按倉(cāng)庫(kù)README的版本要求再往上對(duì)齊一點(diǎn)別用過(guò)于保守的版本。環(huán)境混裝導(dǎo)致第三方庫(kù)沖突。最好單獨(dú)建一個(gè)虛擬環(huán)境避免系統(tǒng)里多個(gè)項(xiàng)目的CUDA依賴互相污染。跑通環(huán)境后的第一步不要直接上benchmark先跑內(nèi)置的smoke test。它能一次性驗(yàn)證數(shù)值正確性、基本性能和當(dāng)前架構(gòu)兼容性是最省時(shí)間的路徑。4.2 實(shí)測(cè)大shape足夠穩(wěn)小shape才是亮點(diǎn)我復(fù)現(xiàn)時(shí)的觀察分為兩類(lèi)形狀。大shape比如M4096、N4096、K4096這種規(guī)整稠密GEMM整體吞吐已經(jīng)能逼近硬件理論峰值穩(wěn)定性也高反復(fù)跑幾輪波動(dòng)很小。這類(lèi)場(chǎng)景本來(lái)就是通用庫(kù)的強(qiáng)項(xiàng)DeepGEMM做得好的地方在于沒(méi)有明顯拖后腿。真正讓我眼前一亮的是小shape。M只有16、甚至更小時(shí)傳統(tǒng)庫(kù)的利用率會(huì)掉得很厲害但DeepGEMM還能保持相對(duì)高的有效吞吐原因就在于它對(duì)wave調(diào)度和tile粒度做了針對(duì)性處理不會(huì)讓SM閑著。對(duì)推理場(chǎng)景來(lái)說(shuō)這就是實(shí)打?qū)嵉氖找嬉驗(yàn)閷?shí)際線上請(qǐng)求很少能湊出那種廠商標(biāo)桿測(cè)試?yán)锏拇骲atch。場(chǎng)景傳統(tǒng)通用庫(kù)的典型問(wèn)題DeepGEMM的表現(xiàn)大而規(guī)整利用率接近峰值但難以超越穩(wěn)定、接近理論峰值小而細(xì)長(zhǎng)SM大量空轉(zhuǎn)吞吐慘淡通過(guò)調(diào)度保持較高利用率MoE分組kernel launch多單個(gè)小GEMM效率低分段GEMM一次處理明顯改善4.3 性能對(duì)比的“公平陷阱”跑benchmark時(shí)最容易犯的錯(cuò)誤是拿FP8和FP16直接比FLOPS。FP8的峰值算力本來(lái)就翻倍了比贏了也不代表你的模型效果就一樣好。更合理的對(duì)比是看“最終推理吞吐”或者“訓(xùn)練迭代時(shí)間”同時(shí)記錄數(shù)值誤差是否在可接受范圍內(nèi)。具體的建議是固定GPU頻率多輪warmup后取穩(wěn)定值對(duì)比時(shí)用相同的內(nèi)存布局和同樣的輸入數(shù)據(jù)不僅看FLOPS還要看L2命中、kernel launch次數(shù)和峰值內(nèi)存帶寬。只有把這幾項(xiàng)放在一起才能判斷一個(gè)實(shí)現(xiàn)是“算法層面的優(yōu)勢(shì)”還是“配置層面的運(yùn)氣”。關(guān)于誤差我用BF16參考實(shí)現(xiàn)與FP8輸出做比對(duì)相對(duì)誤差在可接受范圍內(nèi)符合這一精度檔次的預(yù)期。如果你發(fā)現(xiàn)誤差異常大先檢查scale初始化和e4m3、e5m2是否混用這兩個(gè)問(wèn)題占比最高。5. 遷移到自己項(xiàng)目前需要繞開(kāi)的五個(gè)坑5.1 形狀對(duì)齊與padding策略FP8 GEMM對(duì)形狀對(duì)齊很敏感。M、N、K是16的倍數(shù)時(shí)性能最穩(wěn)達(dá)不到時(shí)要么由庫(kù)內(nèi)部補(bǔ)齊對(duì)齊邏輯要么在外部對(duì)張量做padding。能外部pad就外部pad因?yàn)閹?kù)內(nèi)部對(duì)齊會(huì)消耗共享內(nèi)存和寄存器的有效利用率。padding的時(shí)候注意輸出形狀也要跟著變最后要把多余部分切掉不要讓它影響后續(xù)算子。如果你的業(yè)務(wù)場(chǎng)景里M經(jīng)常是1、2、4這樣的值建議在上層就預(yù)先做微小batch合并處理拉高單次GEMM的有效寬度。5.2 FP8縮放因子的初始化細(xì)節(jié)我最開(kāi)始測(cè)試時(shí)偷懶把scale全部設(shè)為1結(jié)果輸出要么溢出成NaN要么精度差到?jīng)]法用。正確的初始化思路是根據(jù)輸入張量的實(shí)際分布估算統(tǒng)計(jì)最大絕對(duì)值然后用目標(biāo)格式上限除以這個(gè)值得到初始scale。e4m3的有限上限約是448所以初始scale可以是max_val / 448。如果清楚數(shù)據(jù)分布比較集中還可以留一點(diǎn)余量避免訓(xùn)練或推理過(guò)程中偶爾出現(xiàn)的離群值觸發(fā)溢出。不同張量的scale要獨(dú)立計(jì)算不能把所有權(quán)重共用一個(gè)scale。這里有個(gè)額外體會(huì)scale并不是越大越好。scale過(guò)大會(huì)把有效數(shù)值壓得太小尾數(shù)位的潛力浪費(fèi)精度反而下降。比例尺要貼合數(shù)據(jù)分布而不是一味追求“不溢出”。5.3 集成MoE時(shí)的分組約定MoE場(chǎng)景下每個(gè)專(zhuān)家分到的token數(shù)量不固定這是分段GEMM的主要復(fù)雜度來(lái)源。接DeepGEMM時(shí)要格外注意group id和token偏移信息的傳遞。傳錯(cuò)的話輕則性能回退重則直接算錯(cuò)結(jié)果而且這類(lèi)錯(cuò)誤在數(shù)值上不一定立刻看得出來(lái)很可能隱藏到后續(xù)層才會(huì)放大量級(jí)。我建議先構(gòu)造一個(gè)形狀最簡(jiǎn)單的固定分組用例比如兩個(gè)專(zhuān)家、每組各8個(gè)token逐步驗(yàn)證索引映射關(guān)系。跑通之后再上動(dòng)態(tài)分組動(dòng)態(tài)分組的關(guān)鍵是確保每個(gè)專(zhuān)家內(nèi)部的數(shù)據(jù)是連續(xù)的、可以按塊搬運(yùn)的布局而不是松散排布。5.4 混精度使用建議FP8不是萬(wàn)靈藥。對(duì)精度極其敏感的層比如輸出概率分布附近的計(jì)算直接用BF16可能更穩(wěn)妥。我的經(jīng)驗(yàn)是模型較深層的激活經(jīng)過(guò)多次歸一化后分布更穩(wěn)定用FP8不容易出問(wèn)題而靠近輸入和輸出的位置數(shù)值變化幅度大混合使用BF16收益更好。實(shí)際接入時(shí)可以按層配置精度策略前幾層用FP8加速中間敏感層保留BF16末尾再回FP8。這種方法能兼顧性能和穩(wěn)定也方便逐步擴(kuò)大FP8覆蓋范圍。再補(bǔ)充一個(gè)工程細(xì)節(jié)切換不同精度時(shí)不要只改dtype還要重新確認(rèn)對(duì)應(yīng)的scale機(jī)制是否配套。FP8獨(dú)有的scaling在BF16里沒(méi)有反之亦然很多團(tuán)隊(duì)就是在這類(lèi)“只換一半”的地方跌了交椅。6. 一些最后想說(shuō)的花了兩周時(shí)間把DeepGEMM拆完并遷移到自己項(xiàng)目里之后我最深的感觸是優(yōu)化不是靠一句“換上8位浮點(diǎn)”就能完成的系統(tǒng)工程。它需要你對(duì)底層硬件的數(shù)據(jù)通路有足夠理解也需要對(duì)模型本身的數(shù)值邊界有耐心。把DeepGEMM的源碼從頭讀一遍比翻十篇綜述對(duì)你的幫助都大因?yàn)檫@里面每一個(gè)優(yōu)化點(diǎn)幾乎都能對(duì)應(yīng)到你未來(lái)可能在自家框架里遇到的真實(shí)瓶頸。如果條件允許我建議你從Triton路徑跑通開(kāi)始再對(duì)照CUDA路徑逐行去看那些“為什么多寫(xiě)一行”“為什么這樣調(diào)度”的決策最后再遷移回自己的業(yè)務(wù)場(chǎng)景驗(yàn)證收益。這樣一個(gè)流程走完你得到的不僅是性能還有一套可遷移的優(yōu)化思路。