GEMM

GEMM（General Matrix-Matrix Multiplication，通用矩阵乘法）是大语言模型中线性层、注意力投影与 MoE 专家计算的核心算子。模型的大部分参数都参与 GEMM，因此它往往决定推理的计算吞吐、显存带宽利用率与低精度加速效果。

本文先建立 GEMM 的计算模型，再从执行方式、稀疏性和实现生态三个相互独立的维度组织相关概念。

1. 计算模型

1.1 M、N、K

BLAS 将 GEMM 定义为：

D = α op (A) op (B) + β C

其中：

$op (A) \in R^{M \times K}$ ；
$op (B) \in R^{K \times N}$ ；
$C, D \in R^{M \times N}$ ；
$op (X)$ 表示 $X$ 或 $X^{T}$ ；
$α$ 、 $β$ 是缩放系数。

$M$ 和 $N$ 决定输出矩阵的形状， $K$ 是点积的归约维度。若一次乘法和一次加法分别计为一次浮点运算，矩阵乘积约包含：

2 M N K FLOPs

1.2 线性层映射

深度学习框架通常将线性层权重保存为 $W \in R^{N \times K}$ 。若输入为 $X \in R^{M \times K}$ ，则：

Y = X W^{T} + b, Y \in R^{M \times N}

这里的 $M$ 通常是本轮参与计算的 Token 数， $K$ 是输入特征维度， $N$ 是输出特征维度。权重布局虽然是 $[N, K]$ ，但 GEMM 看到的右矩阵是 $W^{T}$ ，因此底层接口经常将其记为 NT 布局。

Transformer 中常见的映射包括：

模块	输入形状	权重形状	输出形状
Q、K、V、O 投影	$[M, H]$	$[N, H]$	$[M, N]$
MLP Gate、Up 投影	$[M, H]$	$[I, H]$	$[M, I]$
MLP Down 投影	$[M, I]$	$[H, I]$	$[M, H]$
LM Head	$[M, H]$	$[V, H]$	$[M, V]$

其中 $H$ 是隐藏维度， $I$ 是中间维度， $V$ 是词表大小。模型结构可参考大语言模型，张量并行对本地 $M / N / K$ 的影响可参考张量并行。

2. 推理工作负载

2.1 Prefill

Prefill 一次处理提示词中的多个 Token，连续批处理还可能将多个请求拼接在同一轮中，因此 $M$ 通常较大。较大的 $M$ 可以让同一块权重被更多输入行复用，提高算术强度，并提供更多并行 Tile。

2.2 Decode

Decode 阶段每个请求通常只生成一个新 Token，但调度器可以把多个活跃请求合并到同一轮，所以 GEMM 的 $M$ 更接近本轮调度的 Token 总数，而不是固定等于 $1$ 。相比 Prefill，Decode 的 $M$ 通常较小，权重读取、Kernel 启动和并行度不足更容易成为瓶颈。

因此，“Prefill 使用 GEMM、Decode 使用 GEMV”只能作为单请求场景下的直觉，不能作为服务系统中的严格划分。

2.3 MoE

MoE 路由会把 Token 分配给不同专家。第 $e$ 个专家接收 $M_{e}$ 个 Token，并执行：

Y_{e} = X_{e} W_{e}^{T}

各专家的 $M_{e}$ 随路由结果变化，而 $N$ 、 $K$ 通常相同。将多个专家问题交给一次 Grouped GEMM，可以减少逐专家启动 Kernel 的开销，并让实现统一调度不同大小的专家任务。具体的数据组织见 GEMM 执行方式，MoE 架构见混合专家模型。

3. 执行方式

Single、Batched 与 Grouped GEMM 使用相同的矩阵乘法定义，但一次接口调用提交的问题数量和参数组织方式不同：

执行方式	问题数量	形状与参数	典型场景
Single GEMM	1	一组固定参数	普通线性层
Batched GEMM	多个	通常共享参数	规则的小矩阵批次
Grouped GEMM	多个	可以不同，约束取决于接口	MoE 专家计算

三者的数学关系、cuBLAS 示例接口，以及 Grouped GEMM 的 Top-2 MoE、Packed 布局和 SGLang/DeepGEMM 实现见 GEMM 执行方式。

4. 稀疏性

稀疏性与执行方式是两个正交维度：单 GEMM、Batched GEMM 和 Grouped GEMM 都可以继续区分为稠密或稀疏计算。

4.1 Dense GEMM

Dense GEMM 按稠密矩阵参与计算，不根据元素值跳过乘加。即使权重中恰好存在零值，只要没有使用稀疏格式和对应的 Sparse Kernel，它仍会消耗普通乘加指令与存储带宽。

Dense GEMM 接口成熟、适用形状广，是当前 LLM 线性层的基础路径。量化会改变输入、权重、累加器和输出的数据类型，但不会自动把 Dense GEMM 变成 Sparse GEMM。

4.2 Sparse GEMM

Sparse GEMM 利用矩阵中的非零模式减少存储或计算，可分为：

非结构化稀疏：非零位置任意，需要索引或压缩格式描述；
结构化稀疏：在固定分组内满足规定模式，硬件更容易高效执行。

NVIDIA Sparse Tensor Core 支持的 2:4 结构化稀疏要求每四个元素中保留两个非零值。使用时不仅需要权重中出现零值，还要完成剪枝、压缩，并由支持该模式的 Kernel 读取元数据。cuSPARSELt 提供了相关的矩阵乘法、剪枝、压缩与 Epilogue 接口。

稀疏指令的理论吞吐提升不等于端到端同比加速，实际结果还取决于压缩开销、元数据访问、矩阵形状和非 GEMM 算子的占比。底层稠密与稀疏 MMA 指令可参考NVIDIA MMA。

5. Kernel 执行结构

5.1 分块矩阵乘法

为简化表达，先忽略 $α$ 、 $β$ 和转置，令 $P = A B$ 。输出元素 $P_{i j}$ 是 $A$ 的第 $i$ 行与 $B$ 的第 $j$ 列的点积：

P_{i j} = \sum_{k = 0}^{K - 1} A_{i k} B_{k j}

不同 $(i, j)$ 的输出元素彼此独立，天然可以并行；同一个输出元素内部则必须沿 $K$ 维完成归约。GEMM Kernel 的核心任务，就是同时提供 $M / N$ 方向的并行度和 $K$ 方向的数据复用。

设线程块负责一个 $M_{b} \times N_{b}$ 的输出 Tile，并将 $K$ 维按 $K_{b}$ 分段。为简化边界，先假设 $M / N / K$ 都能被对应的 Tile 尺寸整除；实际 Kernel 会对末尾的不完整 Tile 使用 Mask 或 Padding。第 $(p, r)$ 个输出 Tile 可以写成分块矩阵乘法：

P^{(p, r)} = \sum_{q = 0}^{⌈ K / K_{b} ⌉ - 1} A^{(p, q)} B^{(q, r)}

其中：

$A^{(p, q)} \in R^{M_{b} \times K_{b}}$ ；
$B^{(q, r)} \in R^{K_{b} \times N_{b}}$ ；
$P^{(p, r)} \in R^{M_{b} \times N_{b}}$ 。

这与逐元素点积完全等价，只是重新组织了计算顺序。每轮加载 $M_{b} K_{b} + K_{b} N_{b}$ 个输入元素，却完成约 $2 M_{b} N_{b} K_{b}$ 次浮点运算：一个 $A$ 元素被同一 Tile 的 $N_{b}$ 个输出复用，一个 $B$ 元素被 $M_{b}$ 个输出复用。忽略写回时，该轮的理想算术强度为：

I_{tile} \approx \frac{2 M_{b} N_{b} K_{b}}{b_{A} M_{b} K_{b} + b_{B} K_{b} N_{b}}

因此 Tiling 的数学目的不是减少乘加次数，而是让相同输入参与更多乘加后再离开片上存储，从而减少每次浮点运算对应的全局显存流量。

5.2 Mainloop 与 MMA

线程块先把当前的 $A^{(p, q)}$ 和 $B^{(q, r)}$ 从全局显存搬入共享内存或 Tensor Memory，再由 Warp 读取更小的 Fragment，通过 MMA 指令更新寄存器中的累加器：

{Acc}_{m \times n} \leftarrow {Acc}_{m \times n} + A_{m \times k} B_{k \times n}

这里的小写 $m / n / k$ 是单条或一组 MMA 指令处理的微型 Tile，满足 $m \leq M_{b}$ 、 $n \leq N_{b}$ 、 $k \leq K_{b}$ 。从完整矩阵到硬件指令，实际上是在多个层级重复同一个分块恒等式：

text

矩阵 P[M, N]
  └─ Thread Block Tile P[M_b, N_b]
       └─ Warp Tile
            └─ MMA Tile A[m, k] × B[k, n]

一个线程块的 Mainloop 按以下顺序执行：

将累加器 $Acc$ 初始化为零；
加载第 $q$ 组 $A$ 、 $B$ Tile；
用若干次 MMA 计算 $A^{(p, q)} B^{(q, r)}$ ，并累加到 $Acc$ ；
沿 $K$ 维遍历所有 $q$ ，得到完整的 $P^{(p, r)}$ 。

多级流水和双缓冲不会改变上述数学结果，只是让第 $q + 1$ 组数据的搬运与第 $q$ 组数据的计算重叠。Tile 越大，理论复用越充分，但寄存器和片上存储占用也越高，可能减少可并发的线程块数量。

若使用 Split-K，第 $s$ 个线程块只计算一段 $K$ ：

P_{i j}^{(s)} = \sum_{k \in K_{s}} A_{i k} B_{k j}, P_{i j} = \sum_{s} P_{i j}^{(s)}

它用额外的 $K$ 方向并行度解决小 $M / N$ 时线程块不足的问题，但多个线程块会产生同一输出 Tile 的部分和，因此还需要原子操作或额外 Kernel 完成第二次归约。

浮点加法不满足严格结合律。分块、流水线或 Split-K 会改变累加顺序，所以不同 Kernel 即使计算同一个数学表达式，也可能产生不完全相同的末位结果。

5.3 Epilogue

Mainloop 结束后，累加器中保存的是 $P = op (A) op (B)$ 的一个输出 Tile。Epilogue 对每个输出元素执行：

D_{i j} = f (α P_{i j} + β C_{i j} + b_{j}; s)

其中 $b_{j}$ 表示可选的 Bias， $s$ 表示可选的量化 Scale， $f$ 表示激活、Clamp 或数据类型转换等逐元素操作。没有对应操作时可省略相关项。

Epilogue 除了完成基础的 $α A B + β C$ ，还可以融合：

Bias；
激活函数；
量化缩放与反量化；
数据类型转换；
Clamp 等逐元素操作。

融合让中间结果直接从累加器变换为最终输出，避免先写回 $P$ 、再由其他 Kernel 读出处理。其代价是增加寄存器压力，并限制算子的组合方式。cuBLASLt、CUTLASS 与 cuSPARSELt 都提供不同范围的 Epilogue 能力。

6. 实现生态

cuBLAS 是预构建数学库，CUTLASS 是 Kernel 开发框架，DeepGEMM 则是面向特定架构和低精度工作负载的专用 GEMM 库。

实现	定位	主要特点
cuBLAS / cuBLASLt	通用数学库	接口稳定、覆盖广，通过启发式选择算法
CUTLASS / CuTe DSL	Kernel 开发框架	显式组合布局、数据搬运、MMA、流水线与 Epilogue
DeepGEMM	专用 GEMM 库	面向 Hopper、Blackwell 与低精度 GEMM，提供 MoE M-Grouped 路径

6.1 cuBLAS 与 cuBLASLt

cuBLAS 提供 NVIDIA GPU 上的 BLAS 实现，包括 GEMM、Batched GEMM、Strided Batched GEMM 和 Grouped Batched GEMM。cuBLASLt 以矩阵乘法为中心，提供更灵活的布局、数据类型、计算类型、算法启发式与 Epilogue 配置。

应用通常描述问题并让库选择候选算法，适合需要广泛硬件和形状覆盖的通用路径。其代价是调用者对 Kernel 内部流水线和数据布局的控制相对有限。

6.2 CUTLASS 与 CuTe DSL

CUTLASS 提供构建高性能 CUDA Kernel 的 C++ 模板组件，并在同一生态中提供基于 Python 的 CuTe DSL。二者都允许开发者表达 Tensor Layout、Tiled Copy、Tiled MMA、流水线和 Epilogue；CuTe DSL 侧重以 Python 语法生成并 JIT 编译特定 Kernel。

因此，CuTe DSL 不宜与 DeepGEMM 并列为同一种“实现库”：它是 CUTLASS 生态中的 Kernel 编程方式，而 DeepGEMM 是可以被推理框架直接集成的专用算子库。CUTLASS 还提供 Operator API，用于发现、编译和执行预定义算子。

6.3 DeepGEMM

DeepGEMM 面向 NVIDIA Hopper 与 Blackwell GPU，提供普通 Dense GEMM，以及为 MoE 设计的 M-Grouped GEMM。M-Grouped 接口包含两种常见数据组织：

Contiguous：将各专家 Token 连续拼接，并用专家索引标识每一行，适合 Token 数量动态变化的 Prefill 或训练；
Masked：为每个专家预留固定上限，通过 Mask 指定有效行数，便于 Decode 场景复用 CUDA Graph。

SGLang 的本地实现同时封装了普通 FP8/BF16 GEMM，以及 Contiguous、Masked 两类 Grouped GEMM，并在 MoE Runner 中根据执行路径调用对应接口。这说明 DeepGEMM 不是“Grouped GEMM”的同义词，而是同时实现多种 GEMM 执行方式的一个具体后端。MoE 路由与 M-Grouped 接口的完整数据流见 GEMM 执行方式。

7. 性能边界

7.1 算术强度

若暂时忽略缓存、量化 Scale、稀疏元数据与 $C$ 矩阵，并假设输入、权重和输出各访问一次，则 GEMM 的乐观最小数据量约为：

Q_{min} = b_{A} M K + b_{B} K N + b_{D} M N

其中 $b_{A}$ 、 $b_{B}$ 、 $b_{D}$ 是每个元素占用的字节数，对应的算术强度近似为：

I = \frac{2 M N K}{b_{A} M K + b_{B} K N + b_{D} M N}

较大的 $M$ 能让每块权重服务更多输入行，通常更容易接近计算受限；小 $M$ 对权重的复用较少，更容易受显存带宽、启动开销和并行度限制。低精度量化同时减少权重流量并提高 Tensor Core 吞吐，但还需考虑 Scale 加载、转换和累加精度。

7.2 形状与调度

相同 FLOPs 的 GEMM 可能具有完全不同的性能，主要影响因素包括：

$M / N / K$ 是否与硬件 Tile 和对齐要求匹配；
输入布局、转置方式与内存连续性；
数据类型、累加类型和量化粒度；
Thread Block 数量能否充分占用 GPU；
Grouped GEMM 中各问题的负载是否均衡；
Epilogue 融合收益是否超过额外资源开销；
Split-K 等并行策略是否引入过多归约成本。

因此不存在对所有模型、批量和 GPU 都最优的单一实现。推理框架通常依据硬件架构、形状、数据类型、布局和是否启用 CUDA Graph 进行 Dispatch，并通过 Benchmark 或启发式选择 Kernel。

8. 参考

Netlib, Basic Linear Algebra Subprograms
NVIDIA cuBLAS, GEMM、Strided Batched GEMM、Grouped Batched GEMM
NVIDIA, cuSPARSELt Documentation
NVIDIA, CUTLASS Efficient GEMM
NVIDIA, CuTe DSL Overview
DeepSeek-AI, DeepGEMM
SGLang, DeepGEMM Wrapper；DeepGEMM MoE Runner；DeepEP Dispatcher；EP Scatter/Gather

LLM

Decoder Layer

FFN

Qwen

DeepSeek V3

Attention

量化基础

量化配置

典型方案

Weight-Activation

PTQ

Weight-Only

Weight-Activation

Prefix Cache

KV Cache

Recurrent State Cache

DP（Data Parallelism）

GEMM

投机采样

方式

Draft Model

优化

Prefill CUDA graph

模型实现

架构

硬件结构

逻辑结构

通信

CUDA

指令体系

矩阵指令

CUDA编程

通信

特性

专题

GEMM ​

1. 计算模型 ​

1.1 M、N、K ​

1.2 线性层映射 ​

2. 推理工作负载 ​

2.1 Prefill ​

2.2 Decode ​

2.3 MoE ​

3. 执行方式 ​

4. 稀疏性 ​

4.1 Dense GEMM ​

4.2 Sparse GEMM ​

5. Kernel 执行结构 ​

5.1 分块矩阵乘法 ​

5.2 Mainloop 与 MMA ​

5.3 Epilogue ​

6. 实现生态 ​

6.1 cuBLAS 与 cuBLASLt ​

6.2 CUTLASS 与 CuTe DSL ​

6.3 DeepGEMM ​

7. 性能边界 ​

7.1 算术强度 ​

7.2 形状与调度 ​

8. 参考 ​

GEMM

1. 计算模型

1.1 M、N、K

1.2 线性层映射

2. 推理工作负载

2.1 Prefill

2.2 Decode

2.3 MoE

3. 执行方式

4. 稀疏性

4.1 Dense GEMM

4.2 Sparse GEMM

5. Kernel 执行结构

5.1 分块矩阵乘法

5.2 Mainloop 与 MMA

5.3 Epilogue

6. 实现生态

6.1 cuBLAS 与 cuBLASLt

6.2 CUTLASS 与 CuTe DSL

6.3 DeepGEMM

7. 性能边界

7.1 算术强度

7.2 形状与调度

8. 参考