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 (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
• 注:本文未依赖联网检索,所有技术要点源于行业公开知识,具体硬件参数(算力峰值、带宽等)应查阅对应厂商最新数据手册,文中以定性描述为主。