最近為了把一個(gè)本地模擬的 decoder-only 大模型推理速度提上來(lái)我把某國(guó)產(chǎn) NPU 異構(gòu)計(jì)算架構(gòu)的算子庫(kù)從上層接口一直翻到了底層 kernel。翻完最大的感受是大模型上層那些 Attention 變體、KV Cache 優(yōu)化都很重要但真正兜底性能的還是那個(gè)最不起眼的 MatMul 算子——大語(yǔ)言模型矩陣運(yùn)算的核心組件幾乎每一層都在反復(fù)調(diào)用它。這篇文章不算什么教科書講解而是結(jié)合我實(shí)際排障過(guò)程把這個(gè)算子在算子庫(kù)里的實(shí)現(xiàn)骨架、性能邊界、優(yōu)化取舍拆開(kāi)講清楚。想快速定位大模型推理性能瓶頸的人、做算子開(kāi)發(fā)或者框架性能優(yōu)化的人應(yīng)該都能從里面找到可以立刻拿去用的判斷方法。先說(shuō)明一下文中的算子庫(kù)并不特指某一個(gè)商業(yè)產(chǎn)品而是我在本地環(huán)境中研究過(guò)的、常見(jiàn)的 NPU 異構(gòu)計(jì)算架構(gòu)里的算子集合接口邏輯和主流算子庫(kù)大同小異。為了避免歧義下文統(tǒng)一用“NN 算子庫(kù)”代稱。MatMul 在所有深度學(xué)習(xí)框架里都是高頻算子但由于 NPU 的硬件調(diào)度方式和通用 GPU 不同它的內(nèi)部設(shè)計(jì)側(cè)重點(diǎn)很不一樣。這篇文章會(huì)盡量把“為什么這樣設(shè)計(jì)”講透。1. 先理解為什么 LLM 每一層都在碰 MatMul1.1 從 Attention 和 FFN 看矩陣乘的落點(diǎn)很多人覺(jué)得矩陣乘不過(guò)是線性代數(shù)課本里的一個(gè)運(yùn)算但放到生成式大模型里它承擔(dān)了幾乎全部的可學(xué)習(xí)權(quán)重。Decoder-only 模型的每個(gè) Transformer Block常規(guī)結(jié)構(gòu)可以拆成兩部分Attention 和 FFN而這兩部分的主要計(jì)算全是 MatMul。比如 Attention 的第一步是 QKV 投影。輸入矩陣 X 的形狀是 [seq_len, hidden_size]權(quán)重矩陣 W_Q、W_K、W_V 的形狀都是 [hidden_size, hidden_size]于是要做三個(gè)矩陣乘得到 Q、K、V。接下來(lái)計(jì)算注意力分?jǐn)?shù)要把 Q 和轉(zhuǎn)置后的 K 相乘拿到分?jǐn)?shù)之后又要用分?jǐn)?shù)和 V 相乘得到注意力輸出最后還有一個(gè)輸出投影矩陣把 Attention 結(jié)果再映射回 hidden_size。光一個(gè) Attention已經(jīng)出現(xiàn)了六次 MatMul。FFN 也不遑多讓。常見(jiàn)的做法是先把 hidden 狀態(tài)投影到一個(gè)更大的中間維度過(guò)激活函數(shù)再投影回 hidden_size這就是兩次大矩陣乘。再加上 Embedding 后的映射、最后一層輸出 logits 的投影一個(gè) 12 層的模型單次生成一個(gè) token 要執(zhí)行的 MatMul 次數(shù)輕松超過(guò)幾十次。如果 QKV 投影沒(méi)有被融合這個(gè)數(shù)字還會(huì)更高。所以我在看 NN 算子庫(kù)時(shí)第一個(gè)反應(yīng)就是只要把 MatMul 這個(gè)算子的性能吃透整個(gè)大模型推理項(xiàng)目的性能畫像就完成了一大半。這不是夸張因?yàn)榇竽Ｐ徒^大部分算力消耗都落在這個(gè)算子上。1.2 GEMM、BMM、量化 MatMul同一算子的幾種面孔算子庫(kù)里的 MatMul 并不是只有一種實(shí)現(xiàn)。從接口上看它通常接收兩個(gè)輸入 tensora 和 b再通過(guò) transpose_a、transpose_b 這類參數(shù)決定是否轉(zhuǎn)置。根據(jù)輸入維度和數(shù)據(jù)格式至少能分出三種常見(jiàn)形態(tài)標(biāo)準(zhǔn) GEMM即通用矩陣乘對(duì)應(yīng)線性層、投影層。輸入是二維矩陣公式為 C alpha * A * B beta * Calpha 和 beta 默認(rèn)一般是 1 和 0。批量矩陣乘 BMM輸入帶有 batch 維比如 Attention 里對(duì)多個(gè) head 的 Q 和 K 做批量乘。形狀可以是 [B, M, K] 乘 [B, K, N]輸出 [B, M, N]。這種形態(tài)對(duì)并行調(diào)度的要求更高因?yàn)?batch 維天然適合分發(fā)給不同計(jì)算單元。量化 MatMul輸入可能是 int8 或 int4 的權(quán)重激活可能仍是 fp16。計(jì)算時(shí)需要在硬件矩陣乘單元里做反量化或者走混合精度通道累加器保持 fp32最終再寫回目標(biāo) dtype。為什么算子庫(kù)要區(qū)分這些形態(tài)因?yàn)榈讓佑布?duì)不同數(shù)據(jù)布局、不同維度的支持差別很大。比如 BMM 的 batch 維可以當(dāng)作并行度量化 MatMul 則受限于硬件是否原生支持低比特乘加指令。如果算子庫(kù)只用一套樸素實(shí)現(xiàn)性能很可能在某種形態(tài)下崩掉。1.3 推理場(chǎng)景對(duì) MatMul 的特殊要求大模型推理和訓(xùn)練對(duì) MatMul 的壓力完全不同。訓(xùn)練時(shí)通常是大的 batch sizeM 很大形狀相對(duì)規(guī)整追求的是吞吐推理時(shí)有兩個(gè)截然不同的階段prefill 階段要處理整段輸入M 等于輸入序列長(zhǎng)度可能很大decode 階段每次只生成一個(gè) tokenM 等于 1。這種“M 一會(huì)兒大一會(huì)兒小”的動(dòng)態(tài) shape會(huì)讓同一個(gè) MatMul 算子呈現(xiàn)出完全不同的性能特征。M1 時(shí)計(jì)算量極小權(quán)重搬運(yùn)占主導(dǎo)M1024 時(shí)計(jì)算量驟增矩陣乘法單元是否飽和成為關(guān)鍵。NN 算子庫(kù)在設(shè)計(jì) MatMul 時(shí)必須能根據(jù) M、N、K 和硬件資源動(dòng)態(tài)選擇不同 kernel而不是靜態(tài)綁定一個(gè)最優(yōu)配置。這也是我后面調(diào)試時(shí)踩坑最多的地方。2. 矩陣乘的性能邊界先算清楚 Roofline 再談優(yōu)化2.1 一次 MatMul 的計(jì)算量與訪存量在動(dòng)手調(diào)優(yōu)之前我習(xí)慣先做一次“紙上計(jì)算”。MatMul 的計(jì)算量非常固定C 的每個(gè)元素要做 K 次乘加C 本身有 M 乘 N 個(gè)元素所以總乘加次數(shù)是 MNK如果按浮點(diǎn)運(yùn)算數(shù) FLOPs 算還要乘以 2因?yàn)橐淮纬思铀銉蓚€(gè)運(yùn)算FLOPs 2 * M * N * K訪存量則取決于輸入數(shù)據(jù)要從主存搬運(yùn)多少字節(jié)。最粗略的估算是不考慮緩存假設(shè) A、B、C 都從主存走一遍訪存字節(jié)數(shù) ≈ (M*K K*N) * sizeof(type) M*N * sizeof(acc)真正關(guān)鍵的不是這兩個(gè)絕對(duì)值而是它們的比值也就是算術(shù)強(qiáng)度arithmetic_intensity FLOPs / 訪存字節(jié)數(shù)用 Roofline 模型的話來(lái)說(shuō)算術(shù)強(qiáng)度低就是訪存受限算術(shù)強(qiáng)度高才是計(jì)算受限。我舉個(gè)具體例子。假設(shè) hidden_size4096decode 階段 M1權(quán)重是 fp16FLOPs 2 * 1 * 4096 * 4096 ≈ 33.6 MFLOPs訪存字節(jié)數(shù)主要是讀取權(quán)重 B4096*4096 個(gè) fp16 約等于 33.6MB算術(shù)強(qiáng)度 ≈ 1 FLOP/Byte這個(gè)強(qiáng)度低得嚇人。換句話說(shuō)每從主存讀一個(gè)字節(jié)只做了一次浮點(diǎn)運(yùn)算。再對(duì)比 prefill 階段 M1024FLOPs 2 * 1024 * 4096 * 4096 ≈ 34.4 GFLOPsA 的大小變成 1024*4096 個(gè) fp16大約 8.4MBB 仍然是 33.6MBC 大約 8.4MB算術(shù)強(qiáng)度 ≈ 34.4e9 / 50e6 ≈ 700 FLOP/Byte同樣一個(gè)權(quán)重矩陣M 變大之后算術(shù)強(qiáng)度提升了幾百倍瓶頸屬性完全反過(guò)來(lái)。這就是為什么大模型推理里“decode 慢”那么普遍它一直在訪存受限區(qū)里跑。2.2 同一個(gè)算子在 Roofline 兩側(cè)反復(fù)橫跳MatMul 比較特殊的地方在于它既會(huì)被歸為訪存密集又會(huì)被歸為計(jì)算密集關(guān)鍵看 shape。NN 算子庫(kù)里的 MatMul 為了兼顧兩種場(chǎng)景往往會(huì)準(zhǔn)備兩套優(yōu)化路線。訪存受限時(shí)優(yōu)化目標(biāo)不是把矩陣乘法單元打滿而是減少主存訪問(wèn)。手段包括把權(quán)重盡量壓在緩存里、多請(qǐng)求合并 batch、使用低比特權(quán)重、避免中間結(jié)果來(lái)回寫。計(jì)算受限時(shí)優(yōu)化目標(biāo)變成提高矩陣乘單元利用率比如選擇合適的分塊、把數(shù)據(jù)搬運(yùn)和計(jì)算重疊起來(lái)、用低精度乘加指令。這兩種思路如果搞混會(huì)出現(xiàn)很典型的性能問(wèn)題一個(gè)人拼命調(diào)分塊大小想把計(jì)算流水線加滿結(jié)果 M1 時(shí)根本沒(méi)有足夠多的計(jì)算任務(wù)所有的努力都在給一個(gè)“等待權(quán)重搬運(yùn)”的過(guò)程做裝飾另一個(gè)人拼命做 kernel 融合減少中間結(jié)果訪存但大 M 場(chǎng)景下其實(shí)計(jì)算早就飽和減少的那點(diǎn)訪存根本改變不了瓶頸。所以每次優(yōu)化前我都要求自己先算出算術(shù)強(qiáng)度再?zèng)Q定動(dòng)作方向。2.3 分塊到底改變了什么分塊是矩陣乘優(yōu)化里最基礎(chǔ)也最核心的手段。樸素實(shí)現(xiàn)里計(jì)算 C 逐個(gè)元素時(shí)每個(gè) C 元素都要遍歷 A 的一行和 B 的一列導(dǎo)致 A 和 B 的數(shù)據(jù)在內(nèi)存里被反復(fù)讀取。分塊之后C 被切成 M_tile × N_tile 的子塊A 和 B 也按對(duì)應(yīng)的 K 方向切成小條讓每個(gè)子塊都能裝進(jìn)片上存儲(chǔ)被多次復(fù)用。一個(gè)典型的分塊偽代碼大概長(zhǎng)這樣for (mo 0; mo M; mo MT) { for (no 0; no N; no NT) { for (ko 0; ko K; ko KT) { load A_tile A[mo:moMT, ko:koKT]; load B_tile B[ko:koKT, no:noNT]; C_tile A_tile * B_tile; } } }MT、NT、KT 怎么選直接影響性能。分塊太大裝不進(jìn)片上存儲(chǔ)分塊太小復(fù)用率不夠。NPU 上還得對(duì)齊矩陣乘單元的形狀比如某些硬件矩陣乘單元一次處理 16×16、32×16 這樣的塊。算子庫(kù)里的 MatMul 之所以看起來(lái)復(fù)雜很大一部分就是在處理這些“對(duì)齊”和“切分”的細(xì)節(jié)。3. 算子庫(kù)中 MatMul 的實(shí)現(xiàn)骨架從接口到 Kernel3.1 算子描述、參數(shù)校驗(yàn)和 Shape 推導(dǎo)NN 算子庫(kù)里任何一個(gè)算子對(duì)外都有一個(gè)標(biāo)準(zhǔn)描述算子名、輸入 tensor、屬性、輸出 tensor。MatMul 的屬性通常包括 transpose_a、transpose_b以及可能存在的 alpha、beta。在真正進(jìn)入計(jì)算前算子框架會(huì)先做 shape 推導(dǎo)和參數(shù)校驗(yàn)。這一步看起來(lái)簡(jiǎn)單但直接影響圖編譯階段能不能靜態(tài)分配內(nèi)存。比如動(dòng)態(tài) shape 下如果 shape 推導(dǎo)失敗整個(gè)圖的執(zhí)行計(jì)劃可能退化成運(yùn)行時(shí)推斷帶來(lái)額外開(kāi)銷。我在調(diào)試模擬項(xiàng)目時(shí)就遇過(guò)一個(gè)問(wèn)題輸入的 seq_len 在 prefill 和 decode 之間切換某些綁定靜態(tài) shape 的 kernel 被強(qiáng)行兜底到通用的慢版本profile 里一眼就能看出來(lái)。后來(lái)把 shape 推導(dǎo)修好讓調(diào)度器正確識(shí)別 M1 的小形狀MatMul 耗時(shí)立刻降了一截。參數(shù)校驗(yàn)還包括 dtype 檢查。fp16、bf16、int8、fp32 混用是常事算子的實(shí)現(xiàn)必須明確累加器精度和輸出 dtype否則精度問(wèn)題會(huì)一層層傳上去。3.2 Tiling 策略是性能好壞的分水嶺分塊聽(tīng)上去簡(jiǎn)單做起來(lái)全是細(xì)節(jié)。同一個(gè) MatMulM1 和 M1024 的最優(yōu)分塊可能完全不同。NN 算子庫(kù)一般會(huì)先按 M、N、K 和 dtype 查一張經(jīng)驗(yàn)配置表如果表里沒(méi)有再跑一個(gè)代價(jià)模型去估算不同 tile 方案的耗時(shí)。舉例來(lái)說(shuō)如果目標(biāo)硬件的片上 SRAM 約 256KB權(quán)重和激活都是 fp16那么一個(gè) 192×64×64 的 tile 可能需要 64KB再加上雙緩沖的副本可能已經(jīng)占滿。而如果權(quán)重是 int8同樣容量能放更大的 tile。tile 不是越大越好還需要考慮多核并行的粒度。常見(jiàn)做法是外層按 M 和 N 切給多個(gè)核內(nèi)層在 K 方向做流水這樣核間通信最少。在實(shí)際調(diào)試中我自己不會(huì)一上來(lái)就調(diào) tile 參數(shù)而是先用 profile 看硬件的矩陣乘單元利用率。如果利用率高但整體耗時(shí)高問(wèn)題往往在數(shù)據(jù)搬運(yùn)如果利用率低再看是不是 tile 大小導(dǎo)致計(jì)算任務(wù)喂不飽。3.3 多核并行與流水線重疊NPU 上跑 MatMul通常不是單核硬算而是把輸出 C 切成多個(gè)子塊分給多個(gè)計(jì)算核心。切分原則是沿 M 和 N 方向切盡量讓每個(gè)核獨(dú)立算完自己的 C 塊避免跨核共享 A 或 B 帶來(lái)的同步開(kāi)銷。除了多核還有流水線。MatMul 的完整數(shù)據(jù)流是“主存→片上存儲(chǔ)→矩陣乘單元→寫回主存”。如果不做流水計(jì)算單元會(huì)一直等數(shù)據(jù)搬運(yùn)做了雙緩沖之后當(dāng)前 tile 在計(jì)算時(shí)下一個(gè) tile 已經(jīng)在搬運(yùn)單位時(shí)間能處理的數(shù)據(jù)量大幅提升。這個(gè)過(guò)程很像餐廳后廚炒菜和備菜同時(shí)進(jìn)行而不是炒完一道再去備下一道。算子庫(kù)里的調(diào)度代碼通常會(huì)顯式管理這些 buffer主存地址、片上 buffer index、計(jì)算任務(wù)隊(duì)列。不同 NPU 的 buffer 大小和搬移指令差別很大所以這部分代碼很少能跨平臺(tái)復(fù)用也是 MatMul 算子實(shí)現(xiàn)里最難維護(hù)的部分。3.4 不通用的底層 Kernel 和統(tǒng)一的分發(fā)邏輯同一個(gè) MatMul算子庫(kù)往往需要維護(hù)多套底層 kernel。CPU 上可以用 SIMD 向量化NPU 上要用專用的矩陣乘單元指令如果還支持通用 GPU還得有另一套基于通用 GPU 指令的實(shí)現(xiàn)。所以圍繞 MatMul 會(huì)有一個(gè)策略分發(fā)層根據(jù)硬件類型、dtype、shape、是否轉(zhuǎn)置、是否 batch選擇具體 kernel。分發(fā)層的存在很有價(jià)值。它可以做到“上層算子只寫一次底層按平臺(tái)擴(kuò)展”。我在本地模擬環(huán)境里主要看的是 NPU 路徑但很多優(yōu)化思路對(duì) CPU 同樣適用只是分塊大小和指令選擇不同。4. 排查 MatMul 性能瓶頸的一段真實(shí)經(jīng)歷4.1 模擬場(chǎng)景為什么 decode 階段一個(gè) token 要 90 毫秒我在本地搭了一個(gè)模擬項(xiàng)目 X結(jié)構(gòu)是 12 層 decoder-only transformerhidden_size409632 個(gè) headhead_dim128權(quán)重全部 fp16。prefill 階段一次性輸入 1024 個(gè) token表現(xiàn)還過(guò)得去真正讓人頭疼的是 decode 階段每生成一個(gè) token端到端耗時(shí)要 90 毫秒左右。第一件事是打開(kāi)算子級(jí) profiling。結(jié)果里 MatMul 相關(guān)算子QKV 投影、注意力分?jǐn)?shù)、注意力輸出、輸出投影、FFN 上下投影、logits 投影加起來(lái)占了總耗時(shí)的 70% 以上。這個(gè)占比完全符合預(yù)期12 層模型每層一堆矩陣乘每次生成 token 時(shí)所有權(quán)重基本都要從頭讀一遍。當(dāng)時(shí)我大概估算了一筆賬每層 QKV 三個(gè) 4096×4096 的 fp16 權(quán)重加上輸出投影和 FFN 的幾組大權(quán)重全部加起來(lái)有數(shù) GB。假設(shè)我的模擬環(huán)境有效帶寬約 200GB/s光是讀取一遍權(quán)重就要二十多毫秒90 毫秒端到端說(shuō)明除了 MatMul 本身還有很多中間搬運(yùn)和調(diào)度浪費(fèi)可以壓縮。4.2 用算術(shù)強(qiáng)度判斷該往哪個(gè)方向調(diào)優(yōu)Profile 數(shù)據(jù)里有每個(gè) MatMul 的 M、N、K我用前面那個(gè)公式手算了一遍。decode 階段幾乎每個(gè) MatMul 的 M1算術(shù)強(qiáng)度在 1 到幾個(gè) FLOP/Byte 之間清一色訪存受限。結(jié)論很明確再去糾結(jié)矩陣乘單元的利用率沒(méi)有意義M1 的任務(wù)根本喂不飽計(jì)算單元。真正的目標(biāo)是讓權(quán)重訪問(wèn)變得“便宜”要么一次少讀一點(diǎn)要么讓讀過(guò)來(lái)的數(shù)據(jù)多算一點(diǎn)要么干脆別反復(fù)讀同樣的權(quán)重。順著這個(gè)思路我做了三件事。第一檢查調(diào)度器是否給 M1 選用了小形狀專用 kernel避免大 tile 帶來(lái)的無(wú)謂循環(huán)開(kāi)銷第二嘗試把部分權(quán)重從 fp16 壓到 int8同時(shí)用 per-channel scale 做量化補(bǔ)償?shù)谌?QKV 投影、FFN 這類多個(gè) MatMul 和中間激活合并成融合算子減少中間變量寫回主存的次數(shù)。4.3 動(dòng)刀之后90 毫秒降到了 42 毫秒這三件事不是一次全上的而是一步步驗(yàn)證。先開(kāi) M1 專用 kernel端到端從 90 毫秒降到 71 毫秒效果明顯因?yàn)檎{(diào)度開(kāi)銷和無(wú)效的并行等待被砍掉不少。然后把注意力相關(guān)的幾個(gè)權(quán)重?fù)Q成 int8權(quán)重讀取字節(jié)數(shù)下降decode 耗時(shí)又降了十幾毫秒但這里要特別小心精度我是先用一小段測(cè)試集做對(duì)比確認(rèn)生成質(zhì)量沒(méi)有明顯回退才繼續(xù)的。最后做融合算子把 QKV 三個(gè) MatMul 加 bias 和殘差合并把 FFN 兩個(gè) MatMul 加激活合并端到端來(lái)到 42 毫秒左右。整個(gè)排查過(guò)程復(fù)盤一下真正讓性能改善的不是某個(gè)“黑科技 kernel”而是先意識(shí)到瓶頸是訪存然后順著訪存路徑把所有浪費(fèi)撿干凈。這里我建議任何做推理優(yōu)化的人都養(yǎng)成一個(gè)習(xí)慣動(dòng)手調(diào)代碼前先花五分鐘算一下當(dāng)前算子的算術(shù)強(qiáng)度再?zèng)Q定優(yōu)化方向否則很容易南轅北轍。5. 從 MatMul 往外看融合、量化與大模型落地的關(guān)鍵5.1 算子融合把多個(gè) MatMul 變成一次循環(huán)在做 MatMul 優(yōu)化時(shí)繞不開(kāi)“融合”這兩個(gè)字。Attention 里的 QK^T 和后續(xù)的 softmax、乘 V如果按原始算子拆開(kāi)做QK^T 這個(gè)中間矩陣要先寫回主存做完 softmax 再讀出來(lái)再和 V 乘。這個(gè)中間矩陣哪怕只寫一次訪存成本也相當(dāng)可觀。目前主流的分塊注意力算法就是把 Q、K、V 按小塊加載到片上Q 和 K 的子塊相乘得到分?jǐn)?shù)子塊softmax 也在片上做然后直接拿結(jié)果子塊去乘 V 的子塊最后累加。整個(gè)過(guò)程中間矩陣不落主存訪存壓力大幅下降。NN 算子庫(kù)里的融合版 MatMul本質(zhì)上就是讓這種“嵌套計(jì)算”在同一個(gè) kernel 里完成而不是靠一次次算子 launch 拼起來(lái)。這種思路不僅限 Attention。FFN 里的 MatMul、激活函數(shù)和下一個(gè) MatMul 也可以融合。融合后少了好幾次全局內(nèi)存 round-tripdecode 場(chǎng)景下收益尤其大。5.2 低比特 MatMul 為什么在推理里那么香量化是另一個(gè)大殺器原因也藏在 Roofline 里。訪存受限時(shí)推理延遲大約正比于需要讀取的權(quán)重字節(jié)數(shù)。fp16 換成 int8字節(jié)數(shù)直接減半int4 如果硬件支持得好再減半。對(duì)著一個(gè)幾乎完全訪存受限的 decode 階段這等于直接把“卡脖子”的通道拓寬。但低比特不是白來(lái)的。int8 和 int4 的矩陣乘通常需要 scale 和 zero point計(jì)算時(shí)累加器還得用 fp32 來(lái)保證精度。如果一個(gè)模型沒(méi)有做過(guò)量化感知訓(xùn)練直接截?cái)鄼?quán)重生成質(zhì)量很可能崩。我的經(jīng)驗(yàn)是先在離線數(shù)據(jù)上小批量驗(yàn)證觀察困惑度或輸出分布再?zèng)Q定要不要上量化。算子庫(kù)里的量化 MatMul 實(shí)現(xiàn)通常會(huì)封裝反量化邏輯但模型側(cè)校準(zhǔn)仍要自己做。5.3 可遷移到日常開(kāi)發(fā)的三條判斷準(zhǔn)則把 MatMul 的優(yōu)化經(jīng)驗(yàn)收攏一下其實(shí)可以抽象出三條通用準(zhǔn)則放到任何性能敏感的開(kāi)發(fā)里都成立動(dòng)手前先算算術(shù)強(qiáng)度。判斷瓶頸是訪存還是計(jì)算不同結(jié)論對(duì)應(yīng)的優(yōu)化動(dòng)作完全不同。用 profile 數(shù)據(jù)驗(yàn)證假設(shè)。不要憑猜每一步改動(dòng)都要能解釋為什么變量在變。尊重 shape 的多樣性。一套靜態(tài)最優(yōu)解很難通吃所有場(chǎng)景調(diào)度器最好能根據(jù)輸入形狀動(dòng)態(tài)切換策略。這三條聽(tīng)上去很樸素但我在實(shí)際項(xiàng)目里反復(fù)吃過(guò)虧才把它們變成肌肉記憶。尤其是第一條很多優(yōu)化方案看起來(lái)很高級(jí)一放進(jìn) Roofline 框架里就原形畢露要么是在給不存在的計(jì)算瓶頸加速要么是繞過(guò)了一個(gè)根本不是瓶頸的訪存。最后分享一個(gè)小技巧。如果你也想研究這類算子庫(kù)別一上來(lái)就扎進(jìn)匯編級(jí)的 kernel 代碼。先把算子的接口、shape 推導(dǎo)、tiling 配置、多核調(diào)度這些“外圍邏輯”看完再結(jié)合自己的模型 profile 數(shù)據(jù)去理解 kernel。這樣讀代碼的效率會(huì)高很多而且更容易定位到實(shí)際問(wèn)題。我自己就是在把 MatMul 算子拆完之后才發(fā)現(xiàn)推理瓶頸背后的真正邏輯。