一次 kernel launch 算完一组形状可以各不相同的独立矩阵乘。通常 $N$、$K$ 固定，每组的 $M$ 不一样。最主要的用处是 MoE 层的 expert 计算。

为什么需要

MoE 的 expert 计算天然是「一堆小 GEMM」：token 按 router 结果分给各个 expert，每个 expert 拿自己那一段做 FFN。照着定义写就是一个循环：

1
2
3

for i, m in enumerate(tokens_per_expert):
    Y[off : off + m] = A[off : off + m] @ W[i]
    off += m

$E$ 上千时这是上千次 kernel launch，每次的启动开销和收尾的 tail effect 都要单独付一遍。更麻烦的是形状：单个 expert 只有几十上百行 token，这种又瘦又小的 GEMM 填不满 SM，算力大半是空转。

Grouped GEMM 把这层循环搬进 kernel——一次 launch，一个 grid 覆盖全部 group 的 tile：

1
2
3
4

# A: [sum(M_i), K]        按 expert 排好序的激活，连续存储
# W: [E, N, K]            expert 权重
# offsets: [M0, M0+M1, ..., sum(M_i)]   每段的边界
Y = grouped_gemm(A, W, offsets)   # [sum(M_i), N]，仍按 offsets 分段

输入要求同一个 expert 的 token 在显存里连续，这正是 MoE dispatch 阶段那次 permute 在做的事；输出保持同样的分段，再交给 combine 还原回 token 顺序。

和 Batched GEMM 的区别

Batched GEMM 要求所有子矩阵尺寸完全相同，A/B/C 是规整的高维 tensor，靠 stride 定位第 $b$ 个矩阵，总 tile 数编译期就能算出来。多头注意力属于这一类。

MoE 不满足这个前提：每个 expert 分到多少 token 由 router 在运行时决定，且彼此不等。padding 到最大值可以硬套 batched GEMM，但负载越不均衡，浪费越大。Grouped GEMM 直接接受变长——kernel 从一张描述表（每组的 $M$ 与指针/偏移）读形状，把各 group 的 tile 展平成一维 tile id，每个 block 按自己的 id 反查它属于哪个 group、落在哪个 tile。代价是形状不再是编译期常量，寻址和调度都得动态做。

实现

CUTLASS 从 2.x 起提供 grouped GEMM，Hopper 上配合 persistent kernel 与 TMA；Triton 也能写，MoE 框架里常见。

实践中还会按阶段分两种布局：

  • 训练和 prefill 用 contiguous 布局，token 已经排好序、$M$ 随 routing 变化，就是上面那个形式。
  • decode 时每个 expert 只摊到几个 token，且拿 tokens_per_expert 去 launch 会引入一次 device-host 同步，于是改用 masked 布局：$M$ 固定成一个上界，用 mask 标出哪些行有效，形状对 kernel 而言是静态的。
  • DeepGEMM 两种都提供。