GEMM 执行方式

Single、Batched 与 Grouped GEMM 描述一次接口调用提交多少个矩阵乘法问题，以及这些问题如何共享参数。三者使用相同的 GEMM 数学定义，区别在于问题数量、形状约束和地址组织方式。

GEMM 的计算模型、Kernel 分块与实现生态见 GEMM。

1. 统一模型

将一个 GEMM 问题的参数记为：

θ = (M, N, K, {op}_{A}, {op}_{B}, α, β)

对形状匹配的 $A$ 、 $B$ 、 $C$ ，其计算为：

D = GEMM (A, B, C; θ) = α {op}_{A} (A) {op}_{B} (B) + β C

其中 ${op}_{A} (A) \in R^{M \times K}$ 、 ${op}_{B} (B) \in R^{K \times N}$ ，输出 $D \in R^{M \times N}$ 。三种执行方式只改变一次调用包含的问题集合：

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

Single、Batched 与 Grouped 是 GPU 数学库中常见的 API 组织方式，不是 BLAS 规范定义的一套统一分类。不同后端对批次内形状、布局和缩放参数的约束应以具体接口为准。

2. Single GEMM

Single GEMM 一次提交一个矩阵乘法问题，使用一组确定的 $θ$ ：

D = GEMM (A, B, C; θ)

普通线性层 $Y = X W^{T} + b$ 通常属于这一形式。cuBLAS 的 cublasSgemm 是典型 FP32 接口：

cpp

cublasSgemm(
    handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    m, n, k,
    &alpha,
    A, lda,
    B, ldb,
    &beta,
    C, ldc);

m/n/k 描述问题形状，lda/ldb/ldc 描述矩阵相邻列之间的存储跨度。接口中的 C 是输入输出指针：调用前对应公式中的 $C$ ，调用后保存 $D$ 。

cuBLAS 的传统接口默认使用列主序。LLM 的低精度与复杂 Epilogue 场景通常使用 cublasGemmEx 或 cublasLtMatmul，但核心参数仍是形状、布局、输入指针与缩放系数。

3. Batched GEMM

Batched GEMM 一次提交多个参数一致、彼此独立的 GEMM。设批量大小为 $B_{batch}$ ，所有问题共享参数 $θ$ ：

D^{(b)} = GEMM (A^{(b)}, B^{(b)}, C^{(b)}; θ), 0 \leq b < B_{batch}

不同 $b$ 之间没有数据依赖。Batched GEMM 在数学上等价于 $B_{batch}$ 次形状相同的 Single GEMM，常见寻址方式包括：

Pointer Array：传入每个矩阵的指针数组；
Strided Batch：以基地址和固定步长定位每个矩阵。

当各矩阵连续存储且间距固定时，可以调用 cublasSgemmStridedBatched：

cpp

cublasSgemmStridedBatched(
    handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    m, n, k,
    &alpha,
    A, lda, stride_a,
    B, ldb, stride_b,
    &beta,
    C, ldc, stride_c,
    batch_count);

第 $b$ 个问题从 A + b * stride_a、B + b * stride_b 和 C + b * stride_c 开始，Stride 以元素为单位。若矩阵地址不满足固定间距，可以使用接收 Aarray/Barray/Carray 指针数组的 cublasSgemmBatched。

注意力中的 $Q K^{T}$ 和 $P V$ 在代数上可以表示为 Batched GEMM，但生产级推理通常由融合 Attention Kernel 完成，不一定直接调用 Batched GEMM 接口。

4. Grouped GEMM

4.1 问题列表

Grouped GEMM 一次提交 $P$ 个相互独立的问题，并允许第 $i$ 个问题使用自己的参数 $θ_{i}$ ：

D^{(i)} = GEMM (A^{(i)}, B^{(i)}, C^{(i)}; θ_{i}), 0 \leq i < P

它在数学上仍等价于分别执行 $P$ 次 Single GEMM。统一提交可以减少逐问题的 Host 调用与 Kernel 启动开销，并让后端共同调度不同大小的问题。

MoE 是典型场景：Router 将 Token 分给不同专家，每个专家形成一个 GEMM，而各专家收到的 Token 数通常不同。

4.2 Top-2 路由

设本轮有 $T$ 个 Token，每个 Token 选择 $k_{r}$ 个专家，共有 $E$ 个专家。Router 输出：

topk_ids[T,k_r]：每个 Token 选择的专家编号；
topk_weights[T,k_r]：每条 Token–Expert 路由的权重。

在不丢弃路由结果的情况下，Token–Expert assignment 总数为：

R = T k_{r} = \sum_{e = 0}^{E - 1} M_{e}

其中 $M_{e}$ 是专家 $e$ 收到的 assignment 数。以 $T = 5$ 、 $k_{r} = 2$ 、 $E = 4$ 、 $K = 8$ 为例。这里的 $K = 8$ 只用于展示数据排列，实际模型的特征维度通常更大，并需满足具体后端的对齐约束。

text

x = [
  [ 1,  2,  3,  4,  5,  6,  7,  8],  # x0
  [11, 12, 13, 14, 15, 16, 17, 18],  # x1
  [21, 22, 23, 24, 25, 26, 27, 28],  # x2
  [31, 32, 33, 34, 35, 36, 37, 38],  # x3
  [41, 42, 43, 44, 45, 46, 47, 48],  # x4
]

topk_ids = [
  [0, 2],  # x0
  [1, 3],  # x1
  [0, 2],  # x2
  [1, 3],  # x3
  [0, 2],  # x4
]

topk_weights = [
  [0.70, 0.30],  # x0 → Expert 0、Expert 2
  [0.60, 0.40],  # x1 → Expert 1、Expert 3
  [0.80, 0.20],  # x2 → Expert 0、Expert 2
  [0.55, 0.45],  # x3 → Expert 1、Expert 3
  [0.65, 0.35],  # x4 → Expert 0、Expert 2
]

Top-2 将每个 Token 展开为两个 assignment，因此 5 个 Token 产生 10 个专家输入行：

text

Expert 0：[x0, x2, x4] → M0 = 3
Expert 1：[x1, x3]     → M1 = 2
Expert 2：[x0, x2, x4] → M2 = 3
Expert 3：[x1, x3]     → M3 = 2

4.3 Packed 布局

Contiguous 布局将 assignment 按专家聚合。假设同一专家内保持 Token 顺序，不考虑对齐时：

text

x_packed = [
  [ 1,  2,  3,  4,  5,  6,  7,  8],  # row 0：x0 → Expert 0
  [21, 22, 23, 24, 25, 26, 27, 28],  # row 1：x2 → Expert 0
  [41, 42, 43, 44, 45, 46, 47, 48],  # row 2：x4 → Expert 0
  [11, 12, 13, 14, 15, 16, 17, 18],  # row 3：x1 → Expert 1
  [31, 32, 33, 34, 35, 36, 37, 38],  # row 4：x3 → Expert 1
  [ 1,  2,  3,  4,  5,  6,  7,  8],  # row 5：x0 → Expert 2
  [21, 22, 23, 24, 25, 26, 27, 28],  # row 6：x2 → Expert 2
  [41, 42, 43, 44, 45, 46, 47, 48],  # row 7：x4 → Expert 2
  [11, 12, 13, 14, 15, 16, 17, 18],  # row 8：x1 → Expert 3
  [31, 32, 33, 34, 35, 36, 37, 38],  # row 9：x3 → Expert 3
]

m_indices = [0, 0, 0, 1, 1, 2, 2, 2, 3, 3]

output_index = [
  [0, 5],  # x0 的两个结果位于 row 0、row 5
  [3, 8],  # x1 的两个结果位于 row 3、row 8
  [1, 6],  # x2 的两个结果位于 row 1、row 6
  [4, 9],  # x3 的两个结果位于 row 4、row 9
  [2, 7],  # x4 的两个结果位于 row 2、row 7
]

四类张量描述两个方向的映射：

张量	形状	映射方向
`topk_ids`	`[T,k_r]`	Token 与路由位置 → 专家编号
`topk_weights`	`[T,k_r]`	Token 与路由位置 → 路由权重
`m_indices`	`[M_packed]`	Packed 行 → 专家编号
`output_index`	`[T,k_r]`	Token 与路由位置 → Packed 行

不考虑 Padding 时， $M_{packed} = R = 10$ 。因此 m_indices 是一维 [10]，而不是 [5,2]：topk_ids 说明一个 Token 被发给哪些专家，m_indices 说明 Pack 后每一行使用哪一个专家的权重。

4.4 专家计算与合并

把专家权重记为 $W_{e} \in R^{N \times K}$ ，每个非空专家分别执行：

Y_{e} = X_{e} W_{e}^{T}, Y_{e} \in R^{M_{e} \times N}

本例对应四个形状不同的专家问题：

text

X0[3, K] × W0ᵀ[K, N] → Y0[3, N]
X1[2, K] × W1ᵀ[K, N] → Y1[2, N]
X2[3, K] × W2ᵀ[K, N] → Y2[3, N]
X3[2, K] × W3ᵀ[K, N] → Y3[2, N]

Grouped GEMM 统一提交这四个问题。计算完成后，Gather 根据 output_index 找回每条路由的结果，并使用 topk_weights 合并同一 Token 的两个专家输出：

z_{t} = \sum_{r = 0}^{k_{r} - 1} p_{t, r} y_{t, r}

在未对齐的示例布局中，x0 的两个结果位于 y_packed[0] 和 y_packed[5]：

z_{0} = 0.70 y_{packed} [0] + 0.30 y_{packed} [5]

4.5 cuBLAS 接口

cublasSgemmGroupedBatched 要求同一 Group 内的问题共享形状、转置方式和缩放参数，不同 Group 之间可以不同。本例中 Expert 0、2 的 $M_{e} = 3$ ，Expert 1、3 的 $M_{e} = 2$ ，因此可以组成两个 Group：

cpp

int group_count = 2;
int group_size[] = {2, 2};
int m[] = {3, 2};
int n[] = {N, N};
int k[] = {K, K};

const float* Aarray[] = {X0, X2, X1, X3};
const float* Barray[] = {W0, W2, W1, W3};
float* Carray[] = {Y0, Y2, Y1, Y3};

cublasSgemmGroupedBatched(
    handle,
    trans_a, trans_b,
    m, n, k,
    alpha,
    Aarray, lda,
    Barray, ldb,
    beta,
    Carray, ldc,
    group_count,
    group_size);

m/n/k 等参数数组按 Group 保存，矩阵指针数组按 GEMM 问题保存。这里有两个 Group，但包含四个独立的 GEMM。

4.6 SGLang 与 DeepGEMM

以下以 SGLang 0b3bb0c 的 DeepEP Normal + JIT DeepGEMM 路径为例。DeepEP 将每个专家的行数对齐到 $A = 128$ ：

{\tilde{M}}_{e} = ⌈ \frac{M_{e}}{A} ⌉ A, M_{packed} = \sum_{e} {\tilde{M}}_{e}

本例有 4 个非空专家，因此对齐后的 $M_{packed} = 4 \times 128 = 512$ 。四个专家分别占据 [0,128)、[128,256)、[256,384) 和 [384,512)。若同一专家内保持前文的 Token 顺序，有效行和反向索引为：

text

Expert 0：row   0 = x0，row   1 = x2，row   2 = x4
Expert 1：row 128 = x1，row 129 = x3
Expert 2：row 256 = x0，row 257 = x2，row 258 = x4
Expert 3：row 384 = x1，row 385 = x3

m_indices = [0] * 128 + [1] * 128 + [2] * 128 + [3] * 128

output_index = [
  [  0, 256],  # x0 → Expert 0、Expert 2
  [128, 384],  # x1 → Expert 1、Expert 3
  [  1, 257],  # x2 → Expert 0、Expert 2
  [129, 385],  # x3 → Expert 1、Expert 3
  [  2, 258],  # x4 → Expert 0、Expert 2
]

实际 ep_scatter 通过原子加法分配 Packed 行，同一专家内部的顺序不保证与示例一致；output_index 会记录每条有效路由实际分配到的行号。完整数据流为：

text

x、topk_ids
  → SGLang ep_scatter
      ├─ x_packed：Packed 行中的 Token 数据
      ├─ m_indices：Packed 行 → Expert
      └─ output_index：Token 与 Top-k 位置 → Packed 行
  → DeepGEMM(x_packed, weights, m_indices)
      └─ y_packed：全部 Packed 行的专家输出
  → SGLang ep_gather(y_packed, output_index, topk_weights)
      └─ z：每个 Token 的加权专家输出

具体执行顺序如下：

分配对齐空间：SGLang 按 num_recv_tokens_per_expert 的对齐后计数分配 x_packed[M_packed,K]、m_indices[M_packed] 和 output_index[T,k_r]。x_packed 使用 torch.empty 分配，因此未写入行的数值没有定义。
Scatter 与索引生成：ep_scatter 先根据各专家的对齐后行数计算起始位置，并为整个专家区间填写 m_indices。对每条 expert_id >= 0 的有效路由，它通过原子加法取得目标行 dest，复制 Token，并同时写入反向索引：

python

dest = atomic_add(expert_start_loc[expert_id], 1)
x_packed[dest] = x[token_id]
output_index[token_id, topk_position] = dest

output_index 由 SGLang 的 Scatter Kernel 生成，不是 DeepGEMM 的返回值。Padding 行没有有效路由，因此不会写入 x_packed，也不会出现在 output_index 中；m_indices 仍会把这些行标记为其所在专家。

Grouped GEMM：SGLang 调用 DeepGEMM BF16 Contiguous 接口：

python

deep_gemm.m_grouped_bf16_gemm_nt_contiguous(
    x_packed,      # [M_packed, K]
    weights,       # [E, N, K]
    y_packed,      # [M_packed, N]
    m_indices,     # [M_packed]
)

DeepGEMM 没有接收 output_index 或有效行数，只能根据 m_indices[r] 为第 $r$ 行选择专家权重：

y_{packed} [r] = x_{packed} [r] W_{m_indices [r]}^{T}, 0 \leq r < M_{packed}

因此，有效行与 Padding 行都会参与计算。Padding 输入可能包含任意值，相应输出也没有语义；GEMM 沿 $K$ 维独立计算各输出行，不会把 Padding 行累加到有效行中。

有效结果合并：ep_gather 遍历真实 Token 的 Top-k 路由，根据 output_index[t,r] 只读取有效的 y_packed 行，再乘以 topk_weights[t,r]：

z_{t} = \sum_{r = 0}^{k_{r} - 1} p_{t, r} y_{packed} [output_index [t, r]]

Padding 输出没有对应的 output_index，因此不会进入最终结果。在本例的对齐布局中：

z_{0} = 0.70 y_{packed} [0] + 0.30 y_{packed} [256]

5. 性能边界

一次提交多个问题不保证一定比多次 Single GEMM 更快。实际性能取决于：

问题数量与每个问题的 $M_{i} / N_{i} / K_{i}$ ；
输入地址是否规则，以及构造指针数组的成本；
Grouped GEMM 中各问题的负载均衡；
MoE 的 Scatter、Gather 和跨设备通信开销；
对齐产生的无效计算；
数据类型、布局和后端调度策略。

Grouped GEMM 也不等同于 Sparse GEMM：前者组织多个独立问题，后者利用矩阵内部的稀疏模式减少存储或计算。cuBLAS Grouped Batched GEMM 面向一般分组问题；DeepGEMM M-Grouped 则利用 MoE 中 $N / K$ 固定、 $M_{e}$ 可变的约束提供专用路径。

6. 参考

NVIDIA cuBLAS, GEMM
NVIDIA cuBLAS, Strided Batched GEMM
NVIDIA cuBLAS, Grouped Batched GEMM
DeepSeek-AI, DeepGEMM
SGLang, DeepGEMM Wrapper；DeepGEMM MoE Runner
SGLang, 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. 统一模型 ​

2. Single GEMM ​

3. Batched GEMM ​

4. Grouped GEMM ​

4.1 问题列表 ​

4.2 Top-2 路由 ​

4.3 Packed 布局 ​

4.4 专家计算与合并 ​

4.5 cuBLAS 接口 ​

4.6 SGLang 与 DeepGEMM ​

5. 性能边界 ​

6. 参考 ​

GEMM 执行方式

1. 统一模型

2. Single GEMM

3. Batched GEMM

4. Grouped GEMM

4.1 问题列表

4.2 Top-2 路由

4.3 Packed 布局

4.4 专家计算与合并

4.5 cuBLAS 接口

4.6 SGLang 与 DeepGEMM

5. 性能边界

6. 参考