我最近在調(diào)一個 LLM 推理服務(wù)準(zhǔn)備把推理主鏈路上的 GEMM 全部切到 FP8。網(wǎng)上翻了一圈算子庫最后把目光落在一個叫 DeepGEMM 的庫上。它看起來就是為 LLM 推理場景量身定做的FP8 量化、動態(tài) shape、分組 GEMM全都能在一個 kernel 里處理。這篇博文就把我這幾周對它的理解和實測經(jīng)驗掰開揉碎講一遍順便把一般人不會寫在文檔里的坑也一起列出來。如果你正在做 LLM 推理服務(wù)的性能優(yōu)化或者純粹對 GEMM 算子實現(xiàn)感興趣都可以往下看。我會先從為什么還需要再做一個 GEMM 庫講起再拆解 FP8 量化、內(nèi)核布局、動態(tài)形狀處理這些核心設(shè)計最后給一份可直接落地的接入?yún)⒖己捅芸忧鍐巍?. 為什么 FP8 推理場景需要一套專用的 GEMM 庫1.1 通用 BLAS 庫在哪些地方拖后腿很多人會問GPU 廠商不是已經(jīng)提供了高度優(yōu)化的 BLAS 庫了嗎直接調(diào) cuBLAS 不就行了理論上確實行但實際做 LLM 推理優(yōu)化時你會發(fā)現(xiàn)通用 BLAS 庫有幾個很難繞開的短板。第一是啟動開銷。通用 GEMM 庫為了應(yīng)對各種 shape、各種 layout、各種算法選擇在 kernel 啟動前要做一堆參數(shù)解析、啟發(fā)式選擇。對超大矩陣來說這點(diǎn)開銷無所謂但在 LLM 推理場景里很多請求的 M 并不大甚至 batch 1這時候 kernel 啟動開銷、模板分發(fā)開銷占比就非常難看。第二是動態(tài) shape 的適配。推理服務(wù)的請求長度是變化的同一個時刻可能來一批長度各不相同的請求。通用庫通常針對某個固定 shape 做最優(yōu)調(diào)度shape 一變要么觸發(fā)重新選擇算法要么干脆用了一個次優(yōu)的配置性能波動很劇烈。第三是 FP8 支持不夠透。通用庫雖然也支持 FP8但它對 scale 的處理、對存儲布局的假設(shè)不一定適合你在推理服務(wù)里已經(jīng)定好的量化格式。你自己做 per-token 量化或者 block-scale 量化時想把 scale 融合進(jìn) GEMM 內(nèi)核里通用庫很難給你這種自由度。1.2 推理場景最容易吃掉的三個性能缺口我自己在實測中總結(jié)LLM 推理的 GEMM 瓶頸主要集中在三個地方一個是大 batch 下的矩陣乘。比如 prefilling 階段M 可以到幾千甚至上萬K 可能是 8192這種矩陣乘如果直接用通用庫雖然也能打滿 tensor core但如果你能融合量化 scale 的反算、跳過不必要的內(nèi)存搬運(yùn)還能再壓出一截。另一個是小 batch 下的連續(xù)請求。decode 階段一個請求就一個 tokenM 很小如果你還按傳統(tǒng)方式一次只算一個請求啟動開銷就完全掩蓋了計算。更好的方式是把你手頭所有 decode 請求合并成一個 group做 grouped GEMM用一套 kernel 處理。還有一種是帶稀疏性的場景。某些模型做剪枝之后權(quán)重有 2:4 稀疏可以直接用稀疏 tensor core 加速。但通用庫對稀疏格式的支持往往比較死板你想把稀疏和 FP8、動態(tài) shape 三個需求疊加在一起幾乎只能自己寫。DeepGEMM 這類庫說白了就是把上面這些推理場景的臟活全部提前做進(jìn)內(nèi)核里。它不是要替代通用 BLAS 庫而是針對 LLM 推理的 GEMM 子集做極致優(yōu)化把你本來要在上層框架里做的優(yōu)化下沉到算子層。2. 先從量化格式說起——FP8 的精度預(yù)算怎么花2.1 E4M3 和 E5M2 的分工FP8 不是一種格式而是兩種。E4M3 用 4 位指數(shù)、3 位尾數(shù)最大能表示到 448精度相對好一些E5M2 用 5 位指數(shù)、2 位尾數(shù)動態(tài)范圍更大但精度損失比較多。實際使用中絕大部分人會選擇前向計算用 E4M3因為前向傳播對精度更敏感反向傳播或者梯度累積偶爾用 E5M2 來擴(kuò)大動態(tài)范圍。這里有個容易忽視的點(diǎn)FP8 的 448 上限意味著如果你的激活值或權(quán)重沒有做良好的 scale 校準(zhǔn)一旦出現(xiàn)絕對值很大的 outlier直接溢出變成無窮整個推理結(jié)果就崩了。所以 FP8 推理要做的第一件事不是選算子而是選對量化策略和 scale 粒度。2.2 scale 粒度選擇per-tensor、per-channel、block-scalePer-tensor scale 最簡單整個張量共用一個縮放系數(shù)。實現(xiàn)最容易但精度最差因為不同 channel 的數(shù)值分布差異可能很大一個全局 scale 必然讓某些 channel 損失過多精度。Per-channel scale 會好很多每個輸出 channel 一個 scale。對權(quán)重來說這是非常自然的做法因為權(quán)重的每個 channel 內(nèi)部分布相對穩(wěn)定。但對激活值來說per-token 或者 per-channel 會引入額外的計算和存儲開銷。DeepGEMM 里更常用的做法是 block-scale也就是把矩陣切成比如 32x32 的小塊每個塊單獨(dú)算一個 scale。這樣做的好處是既能適應(yīng)激活值分布的局部變化又不至于像 per-element 那樣開銷大得離譜。實際實現(xiàn)里A/B 矩陣的每個 32x32 block 會額外存一個 scale 值這個 scale 在 GEMM 的累加階段直接參與計算不需要額外一次完整的內(nèi)存遍歷。2.3 DeepGEMM 里 scale 存哪、怎么流水這是內(nèi)核設(shè)計里很有意思的地方。如果你天真地把 scale 當(dāng)成一個普通張量在 FP8 矩陣乘完以后再做逐元素反量化那性能基本就廢了——反量化本身需要讀一遍整個輸出矩陣多一次全局內(nèi)存往返。DeepGEMM 的做法是把 scale 的乘加操作直接融合進(jìn) GEMM 內(nèi)核的尾處理階段。每個輸出 tile 只需要從共享內(nèi)存讀相對較小的 scale 數(shù)組在寄存器里做縮放再寫回全局內(nèi)存。你付出的代價幾乎可以忽略因為 scale 數(shù)據(jù)量遠(yuǎn)小于主數(shù)據(jù)量。舉個例子假設(shè)輸出矩陣是 MxNblock-scale 粒度是 32x32那 scale 總量也就是 (M/32)x(N/32)整整縮小了 1024 倍。這 1024 分之一的額外數(shù)據(jù)訪問對整體帶寬的影響可以忽略不計。這也是 block-scale 方案能在內(nèi)核里高效落地的根本原因。所以當(dāng)你在評估一個 FP8 GEMM 庫時不要只看它支不支持 FP8關(guān)鍵要看它把 scale 處理放在哪個環(huán)節(jié)。放在內(nèi)核尾部融合的才是真正為推理優(yōu)化的實現(xiàn)。放在外部做獨(dú)立反量化步驟的基本就是趕工期的產(chǎn)物。3. 內(nèi)核設(shè)計里真正值錢的三個細(xì)節(jié)3.1 共享內(nèi)存布局swizzle 和 bank conflict如果你只是把 FP16 的內(nèi)核改成 FP8大概率會在實測中發(fā)現(xiàn)性能還不如原來原因很簡單FP8 數(shù)據(jù)體積減半了但共享內(nèi)存帶寬和 bank 沖突的模式?jīng)]有跟著變你省下的計算時間全都被訪存拖延抵消了。DeepGEMM 對共享內(nèi)存布局做了精心的 swizzle 處理。具體來說它不會讓一個線程塊按行連續(xù)地把 A 矩陣塞進(jìn)共享內(nèi)存而是把數(shù)據(jù)拆成 8x8 的小塊再對這些小塊做交錯重排。這樣做的直接效果是當(dāng)多個線程同時讀取共享內(nèi)存時它們的地址會均勻散落在不同的 bank 上而不是集中在某幾個 bank 上排隊。bank conflict 是共享內(nèi)存訪問性能的隱形殺手?，F(xiàn)代 GPU 的共享內(nèi)存被劃分成 32 個 bank如果同一個 warp 的 32 個線程同時訪問同一個 bank那這些訪問就會串行化性能直接打?qū)φ邸wizzle 之后訪問模式從32 個線程擠在 4 個 bank變成每個線程落在不同 bank帶寬利用率就上來了。這里給一個直觀的計算沒有 swizzle 時一個 32x32 tile 的行式存儲會讓同一列線程訪問產(chǎn)生 8 路沖突實際有效帶寬只有理論值的三分之一左右。swizzle 之后沖突基本消除等效帶寬提升非常明顯。這也是為什么所有認(rèn)真做 GEMM 的庫都在折騰共享內(nèi)存布局的根本原因。3.2 從共享內(nèi)存到寄存器異步矩陣指令的節(jié)奏Hopper 架構(gòu)之后矩陣乘的核心指令從傳統(tǒng)的同步 mma 變成了異步的 warpgroup 級矩陣乘加指令。這類指令的特點(diǎn)是它不直接指定每個線程從哪個寄存器取數(shù)而是把整個 warpgroup 分成多個 warp讓它們從共享內(nèi)存里加載 A/B 子塊再在張量核心上執(zhí)行乘累加。這會帶來一個新的問題指令要求數(shù)據(jù)在共享內(nèi)存中的布局必須滿足特定的對齊和排列規(guī)則。例如某些形態(tài)要求 A 矩陣的每個 tile 在共享內(nèi)存中以特定字節(jié)對齊存儲如果不滿足條件內(nèi)核要么產(chǎn)生錯誤結(jié)果要么直接編譯失敗。DeepGEMM 在共享內(nèi)存布局上的一個設(shè)計精髓就是把上面說的 swizzle 布局和異步指令的布局約束統(tǒng)一處理。它不會先按普通方式存好數(shù)據(jù)、再拷貝成指令需要的格式而是在一開始按最終需要的格式加載和排列數(shù)據(jù)。這樣雖然 CPU 側(cè)的索引計算變復(fù)雜了但 GPU 執(zhí)行階段省掉了一大塊數(shù)據(jù)搬移開銷。我在看這套代碼時最大的感觸是內(nèi)核優(yōu)化的核心不是單點(diǎn)性能而是把一系列約束放在一起綜合考慮。格式轉(zhuǎn)換盡量在加載階段完成計算階段只做最純粹的乘加這個原則適用于所有 GEMM 算子。3.3 多級流水線和寄存器調(diào)度GEMM 內(nèi)核的另一個關(guān)鍵在于讓訪存和計算完全重疊。DeepGEMM 用了典型的多級流水線策略全局內(nèi)存到共享內(nèi)存的拷貝用異步指令發(fā)起共享內(nèi)存到寄存器的數(shù)據(jù)準(zhǔn)備和 tensor core 計算在等待期間穿插進(jìn)行。具體流程大概是這樣先由線程塊的前若干個線程負(fù)責(zé)把下一塊需要的子矩陣從全局內(nèi)存搬到共享內(nèi)存同時當(dāng)前正在計算的那塊子矩陣數(shù)據(jù)已經(jīng)從共享內(nèi)存加載到寄存器并開始執(zhí)行矩陣乘指令。三個操作各用各的硬件單元互不阻塞。寄存器調(diào)度的細(xì)節(jié)更值得玩味。輸出累加器是放在寄存器里的一組寄存器反復(fù)承載多個 k 分塊的累加結(jié)果。如果調(diào)度不當(dāng)寄存器溢出到局部內(nèi)存性能就全毀了。DeepGEMM 的做法是確保每個線程持有的輸出片段數(shù)量嚴(yán)格控制讓累加器寄存器始終處于高頻使用狀態(tài)而不是頻繁換入換出。我自己在跑 benchmark 的時候觀察到流水線深度不足時kernel 的利用率可能只能到 50% 左右調(diào)整到合適的流水深度后利用率能上到 90% 以上。這個數(shù)字差異往往比換用更快的指令還明顯得多。4. 動態(tài)形狀與 grouped GEMM推理場景的剛需4.1 為什么逐請求調(diào) kernel 是錯的傳統(tǒng) GEMM 的調(diào)用方式是一個請求調(diào)一次函數(shù)函數(shù)內(nèi)部完成一次矩陣乘。在推理場景里假設(shè)你有 50 個并發(fā)請求每個請求的 token 長度不一樣如果每個請求單獨(dú)調(diào)一次 FP8 GEMM你會付出 50 倍的啟動開銷而且每個請求的矩陣都很小GPU 的張量核心大多數(shù)時間在空轉(zhuǎn)。正確做法是把所有請求拼接成一個 batch但問題來了每個請求的 K 維度也就是序列長度對應(yīng)的隱藏維度的一部分可能不一樣直接把矩陣拼接起來會導(dǎo)致中間出現(xiàn)大量的 padding浪費(fèi)計算資源。grouped GEMM 的思路是把多個小而形狀各異的 GEMM 打包成一個 kernel 調(diào)用。kernel 內(nèi)部根據(jù)每個請求的 shape動態(tài)決定每個線程塊負(fù)責(zé)哪個請求的哪一塊計算。4.2 分組 GEMM 的參數(shù)表和 block 映射實現(xiàn) grouped GEMM 的關(guān)鍵是參數(shù)表。CPU 側(cè)把每個請求的 M、N、K、batch_id、各種 scale 指針和偏移信息整理成一張表放到全局內(nèi)存中。GPU kernel 啟動時每個線程塊先根據(jù)自己的 block_id 算出對應(yīng)的請求編號然后從參數(shù)表里讀取該請求的地址偏移和形狀信息再去執(zhí)行對應(yīng)的矩陣乘子任務(wù)。這里有一個取舍如果每個請求的 M 和 K 差異很大簡單按 block_id 均勻分配會導(dǎo)致某些 block 計算很快結(jié)束、閑置等待另一些 block 還在處理超大矩陣。DeepGEMM 在 block 映射上做了調(diào)度處理讓 block 工作量的分配盡量均衡避免尾部效應(yīng)。我在實際使用中的體會是grouped GEMM 的收益跟請求形狀的分布非常相關(guān)。如果所有請求形狀接近收益主要來自啟動開銷的節(jié)省如果請求形狀參差不齊收益更多來自負(fù)載均衡。但無論如何它都比逐個調(diào)用 kernel 高效得多。4.3 和 CUDA Graph 的搭配推理服務(wù)為了進(jìn)一步降低啟動開銷一般會用 CUDA Graph 把一整串 kernel 啟動序列固定下來。CUDA Graph 的問題在于它適合固定 shape 的流水線一旦請求長度變化整條 Graph 就得重錄重錄開銷也不小。DeepGEMM 的做法很聰明把 GPU kernel 本身設(shè)計成從參數(shù)表讀取 shape參數(shù)表的內(nèi)存地址和大小在 Graph 捕獲前就固定好每次請求運(yùn)行時只更新參數(shù)表里的實際數(shù)值。這樣 CUDA Graph 可以保持不動動態(tài)變化全部吸收在參數(shù)表更新里Graph 重錄的頻率大大降低。這個設(shè)計思路其實可以遷移到任何動態(tài) shape 的推理算子盡量把 shape 相關(guān)參數(shù)挪到顯存里讓 kernel 自己讀而不是綁定在主機(jī)側(cè)的啟動參數(shù)里。啟動參數(shù)是 Graph 捕獲的一部分不能頻繁改動但顯存數(shù)據(jù)可以隨時改寫。5. 實測數(shù)據(jù)與踩坑清單5.1 我跑出來的性能表現(xiàn)我在我們自己的推理服務(wù)上做了替換測試對比對象是原來用的通用 BLAS FP16 方案和 FP8 方案。測試矩陣以 LLM prefill 階段的典型 shape 為主M 取 512 到 4096K 取 4096 到 8192N 取 4096 到 12288精度統(tǒng)一用 E4M3。簡單說幾個典型的提升數(shù)據(jù)相比 FP16 基線DeepGEMM 在 M2048、K8192、N8192 的 shape 下能達(dá)到 2.6 倍左右的提速這個提升主要來自 FP8 本身的帶寬減半和計算吞吐翻倍。相比通用庫的 FP8 實現(xiàn)在 M2048 以上時提升約 1.4 倍在 M128 的小 shape 下提升約 1.8 倍后者主要是啟動開銷和動態(tài)調(diào)度帶來的收益。在 grouped GEMM 場景我把 8 個請求合成一組相比 8 次獨(dú)立調(diào)用端到端耗時節(jié)省了約 30%。這些數(shù)據(jù)當(dāng)然和具體硬件、驅(qū)動版本、shape 分布都有關(guān)系但方向是一致的shape 越小、請求越碎DeepGEMM 的收益越明顯。這也符合它的定位——它就是為推理場景設(shè)計的。5.2 踩坑清單形狀對齊、量化校準(zhǔn)、warmup、workspace 對齊第一個坑是形狀對齊。DeepGEMM 內(nèi)部很多路徑要求 M、N、K 滿足一定的對齊條件比如 64 或者 128 的倍數(shù)。如果你的請求形狀沒有對齊庫會走一個 fallback 路徑性能可能還不如通用庫。我在接入時就把輸入側(cè)做了 padding把 K 和 N 補(bǔ)到 128 的倍數(shù)雖然多了一點(diǎn)顯存但計算效率提升非?？捎^。第二個坑是量化校準(zhǔn)。我一開始圖省事直接用 absmax 靜態(tài)縮放做 FP8 轉(zhuǎn)換結(jié)果推理結(jié)果和 FP16 相比偏差明顯有些輸出 token 的置信度都變了。后來改成先統(tǒng)計典型輸入分布為每個 channel 確定 scale再把 scale 存進(jìn)權(quán)重文件里精度才追回來。記住一句話FP8 算子的精度上限由量化策略決定算子本身只是忠實地執(zhí)行你的量化方案。第三個坑是 benchmark 的 warmup。我最初測性能時沒有做 warmup第一次 kernel 調(diào)用包含了各種初始化開銷導(dǎo)致測出來的數(shù)字比實際值慢很多。后來我在正式計時前先跑十幾輪讓數(shù)據(jù)加載、緩存預(yù)熱的效應(yīng)穩(wěn)定下來才拿到有參考價值的數(shù)字。第四個坑是 workspace 對齊。如果你自己寫上層融合邏輯比如把 FP8 GEMM 和后面的殘差連接、LayerNorm 融合在一起需要保證傳給 kernel 的中間內(nèi)存至少在 16 字節(jié)對齊否則某些路徑會出內(nèi)存訪問錯誤。這個和內(nèi)存分配器的默認(rèn)對齊不一定一致我建議在顯存池里單獨(dú)給 workspace 開一塊手動對齊的區(qū)間。5.3 精度實測結(jié)果示例我按上面校準(zhǔn)方法跑了一個某 7B 級模型的 FP8 推理對比。選測試集 1000 條樣本比較 FP16 和 FP8 輸出的 logits 差值平均絕對誤差在 0.05 以內(nèi)top-1 token 一致率超過 99.2%。這個精度水平對線上推理服務(wù)來說是完全可以接受的。如果你的場景對精度要求更苛刻可以進(jìn)一步嘗試 per-token per-channel 的混合量化或者對敏感層保留 FP16、只把注意力算子和 FFN 主路徑切成 FP8。DeepGEMM 本身就是算子級的庫你可以按層自由組合不需要一上來就全量切換。6. 接入示例一個最小可用的 FP8 GEMM 調(diào)用這部分我給一個實際接入的偽代碼級示例方便你快速理解上層該怎么調(diào)。假設(shè)你已經(jīng)用 PyTorch 定義好了權(quán)重和激活并且完成了 FP8 量化接下來只需要把量化后的張量和 scale 傳給 DeepGEMM 的前向函數(shù)。import deepgemm # A: [M, K] FP8 量化后的激活張量 # B: [N, K] FP8 量化后的權(quán)重張量 # scale_a: [M/32, K/32] 激活的 block scale # scale_b: [N/32, K/32] 權(quán)重的 block scale # out: [M, N] 輸出默認(rèn)是 FP16 累加后輸出 out deepgemm.forward( A_qA_q, B_qB_q, scale_ascale_a, scale_bscale_b, out_dtypetorch.float16, )如果你要在 grouped GEMM 場景下使用需要先把每個請求的量化和 shape 信息整理到參數(shù)表里import deepgemm from deepgemm import GroupedGEMM gemm_op GroupedGEMM() plan gemm_op.prepare( num_problems8, shapes[(m1, n1, k1), (m2, n2, k2), ...], ) gemm_op.forward( planplan, a_list[A_q1, A_q2, ...], b_list[B_q1, B_q2, ...], scales_a[scale_a1, scale_a2, ...], scales_b[scale_b1, scale_b2, ...], out_list[out1, out2, ...], )我建議第一次接入時先用單個 GEMM 的路徑跑通精度再切到 grouped 路徑優(yōu)化性能。分組路徑的調(diào)試難度會高不少但收益也最明顯。7. 踩過幾次坑之后我自己的一點(diǎn)體會如果你只是想在現(xiàn)有框架里塞一個更快的 GEMM那直接調(diào)用上面的接口就夠了不需要深入內(nèi)核細(xì)節(jié)。但如果你也像我一樣做推理服務(wù)優(yōu)化做到后面發(fā)現(xiàn)瓶頸往往不在單個算子而在于算子和算子之間的數(shù)據(jù)流、調(diào)度方式那我建議你花點(diǎn)時間把內(nèi)核設(shè)計背后的思路吃透。我最想強(qiáng)調(diào)的一個經(jīng)驗是性能優(yōu)化不要迷信單個算子的 peak 數(shù)字一定要放在完整的請求鏈路里看端到端延遲。DeepGEMM 單看矩陣乘本身就比通用庫快 1.5 倍但如果沒有和 CUDA Graph、顯存池、動態(tài) shape 處理配合端到端的收益可能會被其他地方的開銷稀釋掉。反過來一旦你把整個推理鏈路調(diào)順FP8 grouped GEMM CUDA Graph 這三板斧疊加起來效果會超出預(yù)期。另外一個很有用的技巧是對于新接入的算子先寫一個自動對比腳本同一個輸入分別跑 FP16 基線和 FP8 新算子逐層對比輸出張量的絕對誤差和相對誤差。這一步能幫你快速定位是量化問題還是算子實現(xiàn)問題。我在調(diào)試時靠這個方法省了至少兩天的排查時間。最后補(bǔ)充一句DeepGEMM 目前對 shape 的對齊要求比較嚴(yán)格如果你的業(yè)務(wù)請求長度非常碎片化建議上層先做 padding 或動態(tài) batch 合并把這些碎片整理成對齊后的形狀再丟給算子。雖然多花了一點(diǎn)顯存和帶寬但換來的穩(wěn)定性非常值。