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