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