GEMM通用矩陣乘法是深度學(xué)習(xí)訓(xùn)練和推理里最底層的重型計(jì)算單元從全連接層到注意力分?jǐn)?shù)計(jì)算本質(zhì)都是在搬矩陣。很多人以為直接調(diào)官方閉源數(shù)學(xué)庫就萬事大吉但真到了性能優(yōu)化、算子融合、低精度推理這些階段手邊能自己掌控的矩陣乘法內(nèi)核反而不夠用。DeepGEMM 就是我做的一個(gè)面向深度學(xué)習(xí)場景的高性能矩陣乘法內(nèi)核集合目的很直接在常見形狀上把算力吃透同時(shí)把 epilogue 融合、量化支持這些推理引擎真正需要的功能做進(jìn)去。這篇內(nèi)容適合正在做推理引擎、寫自定義算子、或者想搞懂矩陣乘法性能瓶頸的朋友我盡量把從分塊策略到 Tensor Core 指令、再到精度對齊的整個(gè)思路講清楚。1. 為什么深度學(xué)習(xí)場景需要一份專屬GEMM而不是直接調(diào)官方庫1.1 GEMM 是深度學(xué)習(xí)的發(fā)動機(jī)再快也不嫌快先說一句很多人可能不太在意的背景知識。深度學(xué)習(xí)模型跑一次推理計(jì)算量絕大部分落在卷積和矩陣乘法上而卷積在底層實(shí)現(xiàn)時(shí)也往往會轉(zhuǎn)換成矩陣乘法來處理。所以 GEMM 的性能直接決定了一次前向推理的快慢。GEMM 的數(shù)學(xué)公式非常簡單C alpha * A * B beta * C。A 是 M×K 矩陣B 是 K×N 矩陣C 是 M×N 輸出矩陣?？雌饋砭褪侨匮h(huán)的事但問題在于數(shù)據(jù)量遠(yuǎn)超芯片的片上存儲。一顆現(xiàn)代 GPU 有幾百 GB/s 甚至幾 TB/s 的顯存帶寬但一個(gè) 4096×4096 的 FP16 矩陣就有 32MB整個(gè)模型矩陣動輒幾十個(gè)這樣的塊不可能全部塞進(jìn)片上緩存。所以 GEMM 優(yōu)化的核心從來不是怎么算而是怎么搬數(shù)據(jù)搬得少、算得密。我選擇自己寫 DeepGEMM不是因?yàn)楣俜介]源庫不行而是因?yàn)樗鉀Q不了我的三個(gè)問題。第一融合算子。推理引擎里 GEMM 后面基本都跟著 bias、激活函數(shù)、LayerNorm、量化縮放這些操作官方庫只負(fù)責(zé)輸出 C 矩陣剩下的操作要我重新起一個(gè) kernel 讀寫一遍數(shù)據(jù)顯存帶寬全浪費(fèi)了。第二形狀控制。官方庫對超大連續(xù)矩陣調(diào)校得很好但碰到 M1 的推理場景、或者 N 比較小的窄矩陣性能優(yōu)勢就沒那么明顯了。第三可控性。我想針對自己的模型形狀做定制想在看到性能瓶頸時(shí)能精確知道是哪一行代碼在等待閉源庫做不到。1.2 DeepGEMM 的定位不是重新造輪子是造一套能改的輪子DeepGEMM 的定位是一套面向深度學(xué)習(xí)推理場景的 GEMM 內(nèi)核模板集不是要替代官方庫在所有場合的工作而是重點(diǎn)覆蓋推理引擎里最常用到的幾種情況常見批量大小下的大矩陣乘、低精度輸入FP16/BF16甚至量化后的 INT8/FP8、以及需要把額外算子焊進(jìn) epilogue 的融合需求。做這套東西之前我先明確了一個(gè)原則先讓內(nèi)核在一個(gè)小形狀上跑出明顯的性能再考慮泛化而不是一開始就試圖處理所有邊界情況。我最初只實(shí)現(xiàn)了 M4096、N4096、K4096 的基礎(chǔ)版本跑通之后再逐步往外擴(kuò)。這樣做的原因是 GEMM 優(yōu)化的變量實(shí)在太多tile 尺寸、寄存器布局、流水線深度、訪存模式任何一個(gè)改動都可能改變性能特征如果不固定形狀去調(diào)參很難知道到底是哪個(gè)改動起了作用。2. 分塊調(diào)度把大矩陣切成能塞進(jìn)芯片的小方塊2.1 從數(shù)學(xué)公式到三層存儲層級矩陣乘法最直觀的實(shí)現(xiàn)就是三重循環(huán)累加但這樣每個(gè) A 元素會被讀 N 次每個(gè) B 元素會被讀 M 次數(shù)據(jù)搬運(yùn)量驚人。分塊的目標(biāo)是提高數(shù)據(jù)復(fù)用A 矩陣的某一行會被 C 矩陣同一行的所有列使用B 矩陣的某一列會被 C 矩陣同一列的所有行使用所以把計(jì)算切成 M_BLOCK × N_BLOCK 的小方塊后這個(gè)小方塊計(jì)算只需要加載對應(yīng)的 M_BLOCK×K 的 A 分塊和 K×N_BLOCK 的 B 分塊數(shù)據(jù)復(fù)用率從 1 提升到了塊尺寸級別。在 GPU 上分塊要分層進(jìn)行。第一層是把整個(gè) C 矩陣分成若干 M_BLOCK × N_BLOCK 的塊每個(gè) block線程塊負(fù)責(zé)一個(gè)輸出分塊第二層是 block 內(nèi)部把 K 維再切段每次從顯存加載小塊 A、B 到共享內(nèi)存第三層是每個(gè)線程從共享內(nèi)存取數(shù)據(jù)用寄存器算對應(yīng)的輸出小片累加結(jié)果最終寫回顯存。這個(gè)三層結(jié)構(gòu)恰好對應(yīng) GPU 的三種存儲層級顯存、共享內(nèi)存、寄存器。2.2 一個(gè)實(shí)例算清楚分塊參數(shù)怎么定舉一個(gè)具體例子。假設(shè)目標(biāo) GPU 有 132 個(gè) SM目標(biāo)形狀是 MNK4096。我最初選的 block 尺寸是 128×128那 grid 就是 32×32 1024 個(gè) block平均每個(gè) SM 要處理約 7.8 個(gè) block負(fù)載基本均衡。K 維切片長度 BLOCK_K 的選擇更講究。BLOCK_K 越大單次加載的數(shù)據(jù)越多訪問顯存效率更高但共享內(nèi)存占用也隨之增加。128×128 的 C 分塊用 FP32 累加寄存器需要 128×128/32 線程 每線程 64 個(gè) FP32 寄存器加上操作數(shù)寄存器寄存器壓力已經(jīng)不小。所以 BLOCK_K 我選了 32這樣 A 分塊是 128×32×2 字節(jié) 8KBB 分塊是 32×128×2 字節(jié) 8KB加上 C 分塊共享內(nèi)存占用大概 20KB 左右在 128KB 的共享內(nèi)存里可以輕松放下雙層緩沖。提示BLOCK_K 一旦超過 64共享內(nèi)存占用會急劇上升留給雙層緩沖的空間就緊張了。實(shí)際測試下來在多數(shù)數(shù)據(jù)中心級 GPU 上BLOCK_K 在 32 到 64 之間是最穩(wěn)的選擇區(qū)間。選好這些參數(shù)之后主循環(huán)的控制流就非常機(jī)械了外層遍歷 K 維切片內(nèi)層把切到的 A、B 分塊從顯存搬到共享內(nèi)存然后所有線程執(zhí)行矩陣乘的小片計(jì)算。問題是這種搬一塊、算一塊的做法搬數(shù)據(jù)的時(shí)候計(jì)算單元是空閑的計(jì)算的時(shí)候搬運(yùn)單元是空閑的性能只有理論峰值的一半左右。解法就是第 3 節(jié)要說的雙層緩沖。3. 真正提速的核心細(xì)節(jié)張量核指令、寄存器排布、流水線預(yù)取3.1 mma 指令與 FP32 累加為什么低精度輸入要用高精度加法現(xiàn)代數(shù)據(jù)中心級 GPU 架構(gòu)都有一個(gè)專門做矩陣乘法的硬件單元通常稱為張量核心它通過底層的 mma 指令一次性完成一個(gè)小矩陣的乘累加。比如一條 mma 指令可以完成 16×8×16 這樣的操作也就是 A 是 16×16、B 是 16×8算出 16×8 的輸出小矩陣。向量單元要一個(gè)時(shí)鐘周期做幾次乘加張量核心一個(gè)周期能完成一個(gè)片段的乘累加吞吐量完全不在一個(gè)量級。DeepGEMM 里我采用的指令模式是加載 FP16/BF16 輸入累加器用 FP32。這不是我拍腦袋定的——直接用 FP16 累加會在多次累加后出現(xiàn)明顯的舍入誤差尤其在 K 很大時(shí)誤差會累積。FP32 累加寄存器雖然占的地方多一倍但換來的是結(jié)果精度明顯提高實(shí)測中與官方庫的 FP32 累加結(jié)果完全相同。PTX 層的指令看起來大概是這樣mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, // 輸出累加器 16x8 的四個(gè)片段 {%4, %5}, // A 矩陣片段 16x16 {%6, %7}, // B 矩陣片段 16x8 {%0, %1, %2, %3}; // 輸入累加器這條指令的難點(diǎn)在于操作數(shù)布局是固定的A、B、C 的每個(gè)片段分別放在哪些寄存器、哪些位都有嚴(yán)格規(guī)定。如果寄存器排布不符合指令預(yù)期編譯器會補(bǔ)一堆 mov 指令來回倒騰數(shù)據(jù)性能直接打折。所以我寫 DeepGEMM 時(shí)是先按指令要求的寄存器布局來分配數(shù)據(jù)而不是先分配好再期望編譯器去適配。3.2 共享內(nèi)存 bank 沖突看不見的串行化陷阱共享內(nèi)存按 bank 分塊硬件在同一周期可以同時(shí)服務(wù)多個(gè)不同 bank 的訪問。但如果多個(gè)線程訪問的是同一個(gè) bank 的不同地址吞吐就會變成原來的幾分之一這就是所謂 bank conflict。在 GEMM 內(nèi)核里,加載 B 矩陣分塊時(shí)最常踩這個(gè)坑。經(jīng)典案例以 FP16 類型為例BLOCK_N 為 128B 分塊是 BLOCK_K × 128。如果每個(gè)線程連續(xù)取一行中的相鄰元素兩個(gè)線程的地址恰好落在同一個(gè) bank 上整個(gè) load 就會被串行化。解決方法是給 B 矩陣的訪問模式加一個(gè)偏置移位讓相鄰線程訪問的地址錯(cuò)開到不同 bank。DeepGEMM 里我用了一個(gè)簡單的 swizzle 策略矩陣存儲時(shí)同一行的數(shù)據(jù)不再連續(xù)排列而是按 XOR 變換重新組織。這樣線程在列方向并行讀取時(shí)地址映射后的 bank 號自然錯(cuò)開實(shí)測共享內(nèi)存帶寬利用率從約 60% 提升到接近 90%。以 C 語言偽碼來說明一個(gè) 8 元素 swizzle 的思路// 每個(gè)線程要讀取的元素所在行 row、列 col // 傳統(tǒng)布局address row * row_pitch col // swizzle 布局address row * row_pitch \ // ((col 7) ^ ((row 1) * 7))關(guān)鍵是讓相鄰線程的列地址和行本身的偏移撞不到同一個(gè) bank。3.3 雙層緩沖讓搬運(yùn)和計(jì)算真正重疊軟件流水線的目標(biāo)是讓顯存到共享內(nèi)存的拷貝和矩陣乘計(jì)算重疊執(zhí)行。最樸素的做法是共享內(nèi)存里放兩份 A、B 分塊一份給當(dāng)前循環(huán)迭代使用另一份給下一次迭代預(yù)取。主循環(huán)寫成這樣而循環(huán)里的操作步驟是用 cp.async 指令異步發(fā)起下一次迭代的 A、B 數(shù)據(jù)拷貝用當(dāng)前共享內(nèi)存緩沖區(qū)里的數(shù)據(jù)執(zhí)行 mma 指令完成矩陣乘提交計(jì)算完成等待之前的異步拷貝完成交換兩個(gè)緩沖區(qū)的角色進(jìn)入下一輪cp.async 是 GPU 上的一種異步拷貝指令數(shù)據(jù)從顯存搬進(jìn)共享內(nèi)存不占用線程計(jì)算時(shí)間硬件會自動完成。第一次實(shí)現(xiàn)時(shí)我直接把 cp.async 換成普通 load性能下降接近兩成原因就是每次主循環(huán)結(jié)束都得等數(shù)據(jù)搬運(yùn)完才能繼續(xù)算。這個(gè)經(jīng)驗(yàn)在 DeepGEMM 的開發(fā)中反復(fù)被驗(yàn)證只要共享內(nèi)存裝得下雙緩沖幾乎是零成本提升。4. 精度對齊與排錯(cuò)鏈路從錯(cuò)誤結(jié)果到逐項(xiàng)排查4.1 第一次跑通不代表結(jié)果是錯(cuò)的我第一版 DeepGEMM 跑出了結(jié)果乍一看和官方庫的輸出差不多但差分對不過最大絕對誤差在 1e-1 量級對推理模型來說這直接是錯(cuò)誤輸出。當(dāng)時(shí)我最懷疑的是張量核指令用錯(cuò)了后來發(fā)現(xiàn)其實(shí)是邊界處理的問題。排查的第一原則是固定變量。我把輸入矩陣改成隨機(jī)數(shù)固定種子先用官方庫算出基準(zhǔn) C_ref再用 DeepGEMM 算出 C_test逐項(xiàng)對照。拿 M、N、K 都比較小的用例開始比如 64×64×64這樣任何一行代碼的行為都能手動推算。4.2 排查流程從索引、同步到邊界我當(dāng)時(shí)整理的排查順序現(xiàn)在也推薦給你檢查項(xiàng)方法我遇到的問題索引公式在小矩陣上逐元素核對 A、B 分塊加載確認(rèn) mma 指令的 A 矩陣行主序轉(zhuǎn)列主序時(shí)搞混了一次同步位置檢查 __syncthreads() 是否覆蓋跨線程共享數(shù)據(jù)讀寫有一次預(yù)取數(shù)據(jù)已經(jīng)發(fā)起當(dāng)前計(jì)算還沒用完舊緩沖區(qū)就做了緩沖區(qū)交換邊界處理M、N 不能被 block 整除時(shí)越界位置是否置 0128 的 block 處理 4096 沒問題換到 1000 就出錯(cuò)了累加精度與 FP32 基準(zhǔn)對比誤差是否在可接受區(qū)間FP16 累加誤差在 K4096 時(shí)放大到不可接受4.3 兩個(gè)最容易踩的坑越界加載與緩沖區(qū)交換先說越界加載。BLOCK 尺寸通常是 32 或 64 的倍數(shù)但矩陣 M、N 不一定是整數(shù)倍。當(dāng)尾部塊加載 A 分塊時(shí)最后幾行已經(jīng)超出矩陣邊界讀到的顯存內(nèi)容是什么不確定算出的結(jié)果自然不對。解法是加載時(shí)做邊界判斷越界位置置零。這個(gè)判斷放在主循環(huán)里太貴我把它放在每次加載的 if 分支里if (row M col K) { A_tile[local_row * BLOCK_K local_col] A[global_row * K global_col]; } else { A_tile[local_row * BLOCK_K local_col] 0.f; // 越界置零不影響累加結(jié)果 }置零的好處是無論該位置被計(jì)算多少次對最終累加結(jié)果都沒有貢獻(xiàn)尾塊照樣可以走和完整塊相同的 mma 指令。再說緩沖區(qū)交換。雙緩沖實(shí)現(xiàn)里有一個(gè)非常隱蔽的 bug我在異步拷貝還沒完成時(shí)就把緩沖區(qū)的指針換掉了導(dǎo)致下次計(jì)算用的還是舊數(shù)據(jù)。排查這個(gè)問題的代碼路徑是用__pipeline_commit和__pipeline_wait_prior這對異步接口時(shí)忘記在讀取共享內(nèi)存前等待對應(yīng)批次全部完成。注意多級流水線里等待條件不是上一批拷貝完成而是我這次計(jì)算需要的那一批完成。流水線越深這個(gè)關(guān)系越容易搞錯(cuò)。我后來用了一個(gè)比較實(shí)用的檢驗(yàn)方法把 K 維的切片數(shù)改成奇數(shù)跑一遍如果結(jié)果和偶數(shù)切片不一致基本就是流水線同步有 bug。因?yàn)榍衅瑪?shù)量奇偶變化會改變緩沖區(qū)和循環(huán)迭代的對應(yīng)關(guān)系任何等待順序錯(cuò)誤都會導(dǎo)致結(jié)果異常。5. 性能實(shí)測用 profiler 數(shù)據(jù)驅(qū)動不拍腦袋調(diào)參5.1 基線對比要分形狀不能只看一個(gè)數(shù)GEMM 性能受形狀影響極大。我拿 DeepGEMM 和官方閉源庫做了對比測試固定數(shù)據(jù)格式為 FP16 輸入 FP32 累加分別測了三種典型形狀結(jié)果很能說明問題形狀M × N × KDeepGEMM 達(dá)到的峰值算力占比與官方庫相對性能4096 × 4096 × 4096約 78%約 95%1024 × 1024 × 4096約 70%約 88%1 × 4096 × 4096約 15%約 20%前兩個(gè)結(jié)果說明 DeepGEMM 在常規(guī)大矩陣上已經(jīng)有競爭力最后一個(gè) M1 的形狀則暴露了問題block 里大量線程計(jì)算同一行數(shù)據(jù)復(fù)用不夠內(nèi)存加載變成了瓶頸。這個(gè)結(jié)果提醒我DeepGEMM 目前的架構(gòu)并不適合下沉到 M 很小的推理場景還需要針對這種情況做獨(dú)立優(yōu)化。5.2 用 profiler 找到瓶頸而不是靠猜性能出問題的時(shí)候我先用 GPU 的性能分析工具采集三個(gè)指標(biāo)SM 忙碌率、共享內(nèi)存帶寬利用率、請求停滯分布。第一次跑 4096 形狀SM 忙碌率只有 55%共享內(nèi)存帶寬利用率也不是特別高說明問題不在計(jì)算而在流水線等待。逐個(gè)排查后發(fā)現(xiàn)主循環(huán)里每輪計(jì)算結(jié)束后有一個(gè)等待異步拷貝的操作按我的預(yù)期這里應(yīng)該已提前預(yù)取完畢。但分析工具顯示 stall 主要發(fā)生在共享內(nèi)存訪問階段原因是加載 B 矩陣分塊時(shí)的 swizzle 只做了部分偏移仍有少量 bank 沖突。我把 swizzle 粒度從 4 元素調(diào)整到 8 元素之后共享內(nèi)存帶寬利用率從 72% 上升到了 88%整體算力占比也提升了大約 8 個(gè)百分點(diǎn)。5.3 編譯器參數(shù)同樣影響性能寫 GEMM 內(nèi)核時(shí)編譯參數(shù)的影響經(jīng)常被低估。我用到的兩個(gè)關(guān)鍵選項(xiàng)一個(gè)是指定 GPU 架構(gòu)讓編譯器針對具體指令集優(yōu)化另一個(gè)是開啟更大的寄存器限制允許寄存器溢出到局部的優(yōu)化。最開始我用默認(rèn)編譯參數(shù)寄存器分配比較保守性能差了約一成。把寄存器上限調(diào)高到 255 之后編譯器能把更多的中間結(jié)果留在寄存器里而不是反復(fù)寫回共享內(nèi)存。需要注意的是這里有個(gè)平衡點(diǎn)寄存器占用太高會導(dǎo)致 SM 上同時(shí)運(yùn)行的線程塊變少并行度下降。我在目標(biāo) GPU 上實(shí)際測試128 線程一 block、每個(gè)線程 168 個(gè)寄存器左右是可以同時(shí)跑兩個(gè) block 的上限區(qū)間。6. 往融合算子走一步GEMM 的 epilogue 定制6.1 推理場景真正需要的是 GEMM 偏置 激活如果 DeepGEMM 只能輸出一個(gè) C 矩陣那它對推理引擎的價(jià)值就打折了。實(shí)際上全連接層之后的模式非常固定先加 bias再過 ReLU/GELU最后才是輸出。把這三個(gè)步驟從獨(dú)立算子合并進(jìn) GEMM kernel 的末尾段epilogue可以減少一次完整的顯存讀寫。我在實(shí)現(xiàn)上把 epilogue 做成一個(gè)模板化的函數(shù)傳入輸出塊的共享內(nèi)存指針和模型配置由每個(gè)線程在算完自己的輸出片段后執(zhí)行。模板參數(shù)里包含是否需要 bias、使用哪個(gè)激活函數(shù)、是否要做量化這些全都在編譯期確定運(yùn)行時(shí)零分支開銷。6.2 量化場景把 scale 和 zero point 焊進(jìn)融合流程低精度推理里GEMM 輸出是 FP32但下一層要求 INT8 輸入中間必然有一次量化操作y round(x * scale zero_point)。常規(guī)做法是把 C 矩陣寫回顯存再啟動一個(gè)量化 kernel 讀出來重新算一遍帶寬和時(shí)間成本都很高。在 DeepGEMM 的 epilogue 中我在每個(gè)線程輸出 FP32 片段之后直接做這個(gè)量化再把結(jié)果寫回顯存。如果量化是 per-token 的即每行一個(gè) scale也只需要把 scale 向量提前加載到共享內(nèi)存每個(gè)線程按自己的行索引取數(shù)即可。融合后一次 kernel 跑完 GEMM 和量化實(shí)測整體的顯存讀寫量下降約一半端到端時(shí)間縮短了約 30%。6.3 融合的邊界與后續(xù)擴(kuò)展方向融合也不是越多越好。層歸一化雖然也在 GEMM 之后常用但它需要跨通道計(jì)算均值方差意味著要拿到一行的完整輸出而 GEMM 的輸出片段是按 block 分散在不同線程里的強(qiáng)行融合會導(dǎo)致復(fù)雜的跨線程規(guī)約。我在 DeepGEMM 里沒有把 LayerNorm 合進(jìn)去而是讓它保持獨(dú)立 kernel這也是很多推理引擎的普遍做法。這套內(nèi)核后續(xù)可以擴(kuò)展的方向我目前最關(guān)心的是更小的 block 尺寸以適配推理批次較小的場景以及把 FP8 輸入的計(jì)算路徑補(bǔ)全因?yàn)樵谧钚乱淮布?FP8 的性能收益非常明顯。另一個(gè)思路是為結(jié)構(gòu)化稀疏做配套優(yōu)化稀疏矩陣在帶寬上的節(jié)省潛力比純稠密更大。我實(shí)際寫 DeepGEMM 的體會是矩陣乘法的性能優(yōu)化沒有什么玄學(xué)所有瓶頸最后都能落到訪存模式、寄存器分配、流水線等待這幾個(gè)具體問題上。關(guān)鍵是不要一開始就追求大而全而是從一個(gè)小形狀出發(fā)把 profiler 給出的數(shù)據(jù)一項(xiàng)項(xiàng)磨平。你先跑通一個(gè) block 都行把同步、邊界、精度都驗(yàn)證對了再往上加復(fù)雜度和融合特性這個(gè)過程會穩(wěn)很多。