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