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