GEMM

3 秒看懂

GEMM (General Matrix Multiply) 是深度學習最底層的運算原語,幾乎所有的全連線層、卷積重排、自注意力矩陣計算最終都會分解為通用矩陣乘法。它是算力需求的核心尺度:訓練一個大型模型,90% 以上的浮點操作都在跑各種各樣的 GEMM。提升 GEMM 效率,≈ 提升 AI 晶片算力利用率。

3 分鐘產業解釋

在深度學習架構底層,無論你是跑一個線性層 y=Wx,還是計算注意力分數 QK^T,最終都會被編譯器/執行時對映成高度最佳化的 GEMM 核函式。硬體廠商(如 NVIDIA、Intel、AMD)與架構社群圍繞 GEMM 建置了龐大的軟體棧,如 cuBLAS、oneMKL、rocBLAS,專門為不同矩陣形狀、資料版面配置、精度(FP32/FP16/BF16/INT8)提供接近理論峰值的實現。

GEMM 之所以成為核心瓶頸和工程焦點,是因為它計算密度高(單位資料搬運可攤到大量 FMA 運算),但記憶體頻寬需求極大。業界最佳化方向包括:

• 分塊(Tiling) :將大矩陣切分為適配多級快取(GPU Shared Memory、暫存器)的小塊,減少全域性記憶體訪問。
• 張量核心/矩陣引擎:針對低精度(FP16/INT8/FP8)定製的硬體單元,單週期可完成小矩陣乘加,如 NVIDIA 的 Tensor Cores。
• 融合核心:將 GEMM 與啟用函式、Dropout、偏置加法融合,消除反覆寫入/讀取視訊記憶體的開銷。

產業玩家圍繞 GEMM 的效率展開算力軍備競賽,從單純的硬體堆疊走向更精密的資料流排程;而 Transformer 架構的興起又帶來了“動態形狀”、“混合專家(MoE)路由批處理”等新挑戰,驅動 GEMM 從標準的稠密運算擴充套件到稀疏、塊稀疏、可變長度批處理等變體。

15 分鐘專家深入

對於一個大規模 Transformer 的訓練或推論,GEMM 的時間可以佔端到端延遲的 60%~85%。理解 GEMM 與它“上家”(架構層運算元)和“下家”(硬體執行)的對映關係,是系統最佳化工程師的基本功。

算力需求與矩陣形狀

以 GPT 類模型的單層前饋網路為例(隱藏維度 d_model,前饋層擴張係數通常為 4),它包含兩個權重矩陣:

• 第一層 W1: (d_model, 4*d_model),輸入 X: (batch*seq_len, d_model) → Y = XW1,計算量為 2 * batch*seq_len * d_model * 4*d_model FLOPs。
• 第二層 W2: (4*d_model, d_model) → 輸出同樣形狀。

注意力計算中,Q,K,V 分別通過三個(d_model, d_model)權重獲得,之後的 Scores = QK^T 本身也是一次矩陣乘法,形狀(batch*num_heads, seq_len, d_head) 矩陣相乘,算力消耗不可忽視。在大型模型長序列場景下,Scores 矩陣的 GEMM 經常成為記憶體和計算的爆發點。

並行策略中的 GEMM 通訊對映

深度學習分散式訓練的不同並行策略,本質上是把 GEMM 的輸入/權重/輸出的某一維切分到不同裝置上,然後插入通訊運算元在 GEMM 前後。

• 資料並行(DP):每個裝置有一份完整模型副本,各自做相同的 GEMM,僅需在反向後對梯度進行 AllReduce(所有權重梯度求和)。GEMM 本身無跨裝置依賴。
• 張量並行(TP,如 Megatron-LM):對單層權重矩陣沿列或行切分。例如 W1 沿列切成 [d_model, 4*d_model/tp],每個裝置拿一份 X 做區域性的 GEMM,然後對輸出拼接或求和。前向時通訊常見 AllReduce 對啟用求和,反向時需 ReduceScatter 對梯度累加。注意:此處通訊原語是 AllReduce/ReduceScatter,不是 All-to-All,後者多見於 MoE 的 expert dispatch。
• 流水線並行(PP):按層切分,GEMM 在單個裝置上完整執行,裝置間只傳遞啟用和梯度,本質是 P2P 通訊,不涉及集合通訊。

MoE 層的路由 Gating 會先將 token 按 top-k 分發到不同專家,此時輸入 token 被重排,形成非連續的 batch,每個專家內部仍是標準的稠密 GEMM,但矩陣形狀變得不規則(小 batch 數)。這對 GEMM 庫的 batch 維度最佳化提出要求,也催生了專用的稀疏 GEMM 核。

技術原理

通用矩陣乘法的數學定義與關鍵引數

GEMM 計算 C = α·A·B + β·C,其中 A 為 M×K 矩陣,B 為 K×N 矩陣,C 為 M×N 矩陣。深度學習場景下 α=1, β=0 或 β=1(偏置融合)最常見。FLOPs = 2*M*N*K。決定效能的關鍵引數:

• 算力利用率:實際達到的 TFLOPS 除以晶片理論峰值 TFLOPS,通常受資料搬運、流水線停頓、波尾效應(最後一批塊尺寸小)限制。
• 算術強度:計算量 / 訪存量 = 2*M*N*K / ((M*K + K*N + M*N) * sizeof(element)) FLOPs/byte (若僅計元素個數則為 2*M*N*K / (M*K + K*N + M*N) FLOPs/element)。當算術強度高於機器平衡點(算力/頻寬)時,運算才是計算受限,否則為訪存受限。

GPU 上的典型實現流程(以 NVIDIA 生態為例)

現代 GPU 將 GEMM 做成層次化分塊,以暫存器、Shared Memory 和 L2 Cache 構成三級快取體系:

1. 全域性→分塊:把 A、B 按 M 和 N 方向切為多個宏塊(Thread Block Tile),每個 Thread Block 負責 C 的一個子塊。
2. Shared Memory 緩衝:每個 Thread Block 將 A 和 B 的對應條帶分批從 Global Memory 載入 Shared Memory(例如一次取 Tm×Tk, Tk×Tn)。這裡的關鍵是使用雙緩衝,掩蓋資料載入延遲。
3. Warp 級乘加:Warp(32 執行緒)從 Shared Memory 取更小的塊放入暫存器,然後通過 mma.sync(Tensor Core 指令,FP16/INT8 等)或傳統 FMA 指令完成矩陣乘加。Tensor Core 一次可完成 M=16,N=16,K=16(形狀隨資料型別變化)的小矩陣乘法,並在單週期內輸出累加結果。
4. 結果寫回:子塊 C 最終寫入全域性記憶體。

簡化的分塊 GEMM 資料流:
Global Memory ──▶ Shared Mem Tile A ──▶ Reg Fragment A ──┐
                   (雙緩衝)                              │  FMA / mma
Global Memory ──▶ Shared Mem Tile B ──▶ Reg Fragment B ──┘
                                                        │
Global Memory ◀──  Epi. (啟用/Scale) ── C 累加器  ◀─────┘

精度與資料格式

• FP32:老標準,高精度但功耗大,通常只用於數值敏感的主權重副本。
• FP16/BF16:目前訓練主流。BF16 的指數位與 FP32 相同,不易溢位,混合精度訓練用 BF16 時一般不需損失縮放(loss scaling)。FP16 則需要。
• INT8/FP8:推論或低精度訓練前沿。NVIDIA Hopper 架構引入 FP8 的 Transformer Engine,配合微觀尺度縮放因子,將 GEMM 的吞吐翻倍。此時 GEMM 核需內嵌縮放和反量化邏輯。

與卷積、注意力運算元的關係

• 卷積→GEMM:通過 im2col 或直接卷積演算法(Winograd/FFT)將三維特徵圖和四維卷積核重排為二維矩陣相乘。im2col 會造成視訊記憶體膨脹,但讓 cuBLAS 等通用矩陣庫可以直接加速。
• 自注意力→GEMM:Q,K,V 投影是典型 GEMM;Scores = QK^T 是形狀 (B*H, S_q, d_head) 與 (B*H, d_head, S_kv) 的批次 GEMM(即把 B*H 視為批次維度,每個頭獨立進行 (S_q, d) × (d, S_kv) 的矩陣乘法,而不是將 S 直接拼入 M 維形成 (B*H*S, d) 的單次乘法);Attention Weight * V 同樣。FlashAttention 的核心思路正是將這些 GEMM 與 softmax 進行分塊融合,避免中間結果寫成完整 S×S 矩陣,從而將 O(N²) 的視訊記憶體需求降為 O(N)。但每個分塊內部仍大量呼叫 GEMM 核。
• 多層感知機(MLP)與啟用函式融合:將 GELU/SiLU 等逐元素操作合入 W1 的 GEMM 的後處理階段,減少 kernel launch 開銷和資料搬運。這是許多競速最佳化(如 FasterTransformer、vLLM)的常規操作。

技術演進史

1. 純 CPU 時代:BLAS 標準(Fortran 介面)定義了幾十年的 GEMM 黃金實現,如 GotoBLAS、OpenBLAS 基於手調分塊和暫存器分派達到峰值 80-90%。
2. GPU 興起(2006 CUDA):NVIDIA 提供 CUBLAS 庫,最初為稠密矩陣設計,但早期分塊策略偏保守,離峰值較遠。社群湧現 MAGMA 等專案,用異構架構將 CPU 與 GPU 的 GEMM 結合。
3. Volta Tensor Core(2017):硬體首次引入專用的矩陣乘加單元,cuBLAS 隨之推出 cublasGemmEx 介面支援混合精度。訓練精度從 FP32 主導向 FP16 混合精確轉移,訓練算力需求降低約一半,直接推動 BERT、GPT-2 等大型模型可行。
4. Ampere 與結構稀疏(2020):Tensor Core 支援 2:4 結構化稀疏,理論吞吐翻倍(實際受稀疏模式率限制)。BF16 和 TF32 資料型別加入,大幅降低混合精度訓練的調參難度。
5. Hopper FP8 與 Transformer Engine(2022):FP8 資料格式硬體原生支援,用微觀的縮放因子動態調整量化範圍,可在不犧牲收斂性的前提下將 GEMM 吞吐再翻倍。同時,wgmma 等非同步 warp 組 GEMM 指令允許與資料搬運更深層次覆蓋。
6. FlashAttention 等融合演算法(2022+):雖然不是直接創新 GEMM 本身,但改變了自注意力中 GEMM 的排布與融合方式,導致上游架構對 GEMM 呼叫的形狀與頻率發生顯著變化,反向推動 BLAS 庫針對 M、N 維度小但 K 較大或 batch 很大的場景進行特化。

技術路線對比(量化表)

維度 / 方案	傳統通用 GPU GEMM (cuBLAS)	專用加速器 GEMM (TPU/昇騰等)	稀疏 GEMM 核	GEMM 融合 (FlashAttention 等)
核心計算單元	Tensor Core / CUDA Core	脈動陣列(Systolic Array)	Tensor Core + 掩碼	手工發射的 CUDA kernel
優點	軟體棧成熟,適應性極強	能效比高,峰值算力集中	特定稀疏模式下算力翻倍	消除中間張量,視訊記憶體與頻寬收益巨大
缺點	通用性帶來 overhead,部分批次/形狀效率偏低	程式設計靈活度受限,生態碎片化	對非結構化稀疏、動態稀疏無用	實現複雜度高,每場景需手寫 kernel
典型算術強度	30~150 (依矩陣形狀)	50~200+ (小批次仍高)	與稠密類似,但有效算力較軟	依賴融合運算元的全域性資料流
規格標註	基於 GPU 型號,TF32/FP16 峰值見官方資料 [廠商財報]	TPU v5p 等,具體規格 [未充分揭露]	Ampere 結構稀疏,理論 2× [NVIDIA白皮書]	軟體特性,不提供硬體規格

上下游

上游:

• 硬體平台:GPU (NVIDIA/AMD)、TPU (Google)、Habana Gaudi (Intel)、自研 ASIC 等提供底層 FMA/脈動陣列單元;互連(NVLink/PCIe/IPU)頻寬決定多機 GEMM 拆分時通訊成本。
• 編譯器與底層庫 :cuBLAS、cuBLASLt、CUTLASS(模板庫)、rocBLAS、oneDNN。CUTLASS 解構了 GEMM 抽象,讓使用者可組合 tile 大小、資料移動、計算路徑,是研究 GEMM 最佳化的黃埔軍校。
• 並行架構與通訊庫:NCCL/RCCL 泛集的 AllReduce/ReduceScatter 實現影響張量並行中 GEMM 的通訊疊加效率。

下游:

• 深度學習架構:PyTorch 的 torch.nn.Linear、torch.nn.functional.linear、torch.baddbmm 等最終呼叫 cuBLAS GEMM。TensorFlow/JAX 通過 XLA 自動將 dot_general 對映到對應 GEMM。
• 推論引擎:TensorRT、ONNX Runtime、OpenVINO 會進行層融合,將卷積、BN、啟用重排為連續的 GEMM 呼叫,並應用 INT8/FP8 校準。
• 大型模型服務:vLLM 的 PagedAttention、SGLang 等將 KV 快取管理後,自注意力計算仍本質是多個自定義形狀的 GEMM,由廠商定製 kernel。

關鍵指標

• 有效 TFLOPS/TFLOPS 利用率:某 GEMM 例項實際達到的算力(通常通過 2*M*N*K / 時間 計算)與硬體標稱峰值的比值。利用率<40% 說明 kernel 設計不佳或形狀太差;>80% 為優秀。
• 算術強度與頻寬利用率:理論分析給出記憶體事務和 L2/視訊記憶體頻寬需求,實際頻寬利用率 = 實際資料搬運 / 理論頻寬上限,與 latency hiding 效果強相關。
• 延遲與批次吞吐:對於線上推論,更關心小 batch 下的延遲(μs級)。這時 kernel launch 開銷、執行緒塊規模不足引起的尾延時成為主導。
• 功耗與能效比:GFLOPS/W,大晶片通常在 0.5-2 GFLOPS/W 之間 [基於行業報告估算],小晶片或端側 MCU 可能更低。

供需與市場資料

• 算力需求:在大型模型訓練中,GEMM 佔總浮點操作的 70-90% 不等。每訓練 1 exaFLOP 的 GPT 模型,接近 0.7-0.9 exaFLOP 是 GEMM。全球 AI 訓練晶片在 2024 年的出貨量所對應的總算力估計已超過數百 exaFLOPs (FP16/BF16) [行業估算],GEMM 是算力的主要消費者。
• 軟體棧市場 :NVIDIA cuBLAS 生態佔據主導,capture 了超 90% 的 AI 訓練 GEMM 呼叫 [市場估算];開源 CUTLASS 是幾乎一切新互聯硬體廠商移植堆疊的首要目標;ROCm 的 rocBLAS 正在快速追趕。
• 新硬體格局:由於 GEMM 是明確的殺手應用,ASIC 公司(如 Groq、Cerebras、d-Matrix)或 IP 廠商(CEVA、Imagination)的第一款 AI 晶片第一件事就是復現高效能 GEMM,然後才新增其他運算元。

代表公司與資本對映

• NVIDIA:硬體(Tensor Core、FP8 支援)與軟體(cuBLAS、CUTLASS、cuDNN 隱式 GEMM)雙管齊下,直接定義了通用 GPU GEMM 的標準。股價的長期驅動力之一是大型模型軍備競賽帶來的 GEMM 算力需求。
• AMD:CDNA 架構矩陣核心 + rocBLAS,正通過 HIP 語言降低遷移門檻,在部分超算專案中與 NVIDIA 直接競爭 GEMM 效率。
• Google:TPU 的脈動陣列對 GEMM 特別高效,其 XLA 編譯器針對 TPU 上的 GEMM 進行極細顆粒度排程,但硬體不外售,僅以雲端服務提供。
• 初創與垂直整合公司:Tenstorrent、Graphcore、Chiplet 公司 d-Matrix、光計算公司 Lightmatter/Luminous 等,無一不以“比 GPU 更高效率的 GEMM”為賣點。
• 軟體工具鏈 :CUTLASS 本身不產生直接營收,但建置了強大的護城河;提供效能除錯工具的廠商(Snellius、Augtera)可切入 GEMM 效率監測市場。

投資邏輯

1. GEMM 專用硬體和軟體是AI 基建的“賣鏟人”:大型模型訓練成本中,GEMM 消耗的電費、晶片折舊佔比極高。任何能在同等精度下將 GEMM 能效提高 2x 的技術,就有條件獲取超額獲利。
2. 軟體生態鎖定才是真正壁壘 :NVIDIA 的 CUDA 和 cuBLAS 積累了幾十年(相比 AI 時代而言)的手工調優與自動調優資料,僅 cublasGemmEx 就包含數千個啟發式規則,這才是競爭對手即便硬體峰值算力追上也無法短期超越的“暗知識”。
3. 稀疏化和低精度是未來趨勢:結構化稀疏(如 2:4)和 FP8/FP4 會深刻改變 GEMM 資料通路和編譯器流程,相關 IP 和 EDA 工具(E.g,Cadence Tensilica)可能會有新的增長點。
4. 關注GEMM 的新場景:邊緣裝置上,小矩陣 GEMM (M,N 為幾十到幾百)的吞吐和延遲最佳化催生了 TinyML 引擎(CMSIS-NN、XNNPACK)的需求,這類最佳化是連線端側 AI 模型與廉價 MCU 的關鍵橋樑。

常見誤讀糾偏

1. 誤讀:“MoE 張量並行通訊主要是 All-to-All”
   • 事實:All-to-All 確實是 MoE 路由(dispatch & combine)時 token 重分佈的通訊原語。然而在 MoE 內部,每個 Expert 的 MLP 層做張量並行時,其 GEMM 前後的通訊仍然是 AllReduce/ReduceScatter(與標準 Transformer 一致)。All-to-All 只是 token 搬運,不改變矩陣乘法本身的並行通訊模式。
2. 誤讀:“GEMM 佔用視訊記憶體大小就是權重矩陣的大小”
   • 事實:一次典型的梯度更新中,視訊記憶體佔用包括輸入啟用(通常比權重大若干倍)、中間 GEMM 緩衝區(如重排矩陣)、以及用於雙緩衝的 Shared Memory 分配。尤其在訓練中,前向啟用需要儲存用於反向傳播,導致視訊記憶體佔用數倍於單純儲存矩陣權重。啟用檢查點(Activation Checkpointing)就是為了緩解因 GEMM 前後儲存啟用而導致的視訊記憶體爆炸。
3. 誤讀:“用 Tensor Core 就自動能到峰值算力”
   • 事實:Tensor Core 需要滿足許多約束:矩陣維度對齊(如 16x16 子塊要求 M,N,K 都為 16 倍數)、資料型別匹配、Shared Memory 無 bank conflict、Warp 排程規避停滯等。大多數現實網路中的 GEMM 形狀(如 M=13,N=2048)無法完美對齊,需 padding 或重排,利用率可能只有 40~60%。

學習路徑

1. 基礎 :理解矩陣乘法演算法複雜度與快取命中的關係,閱讀《深入理解計算機系統》Memory Hierarchy 章節,然後用 C 手寫一個簡單分塊 GEMM 並對比 OpenBLAS 效能。
2. 進階 :學習 CUDA 程式設計,執行官方 gemm 示例,逐步引入 Shared Memory 和雙緩衝;閱讀 CUTLASS 文件及程式碼,理解 device_kernel、threadblock_swizzle、Epilogue 等概念;在 NVIDIA Nsight Compute 中分析自己的 kernel 的 roofline 圖。
3. 硬體級最佳化 :研究 Tensor Core 的 mma.sync 指令和資料通路,用 CUTLASS 實現一個支援 BF16 輸出 + GELU 融合的 GEMM。嘗試用 CUTLASS Profiler 自動調優獲取最佳 tile 配置。
4. 系統級聯 :結合 Megatron-LM 或 vLLM 原始碼,找出 ColumnParallelLinear / RowParallelLinear 中 GEMM 呼叫點,用 nsys 觀察張量並行通訊與 GEMM 的重疊情況。
5. 前沿跟進 :追蹤 NVIDIA 每年的 GTC 關於 cuBLAS 和 CUTLASS 的 session;研究 FlashAttention 系列如何通過分塊 GEMM 和手寫 warp shuffle 替代標準 cuBLAS 呼叫。

一句話總結

GEMM 是深度學習計算之核,所有上層演算法創新最終都會轉化為對這一古老運算元的極致壓榨;掌握了 GEMM 在硬體-軟體棧各層的最佳化,才算真正觸達 AI 算力的第一性原理。

延伸閱讀與來源

• NVIDIA cuBLAS 文件:https://docs.nvidia.com/cuda/cublas/
• CUTLASS 開源庫:https://github.com/NVIDIA/cutlass
• “Roofline Model for Performance” by Lawrence Livermore National Lab [經典效能建模論文,無公開連結]
• 《Programming Tensor Cores in CUDA 9》- NVIDIA Developer Blog
• 《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》- Shoeybi et al. (ArXiv 1909.08053)
• FlashAttention 系列論文:Dao et al., 2022, 2023
• 注:本文未依賴聯網檢索,所有技術要點源於行業公開知識,具體硬體引數(算力峰值、頻寬等)應查閱對應廠商最新資料手冊,文中以定性描述為主。