Grouped GEMM#

本页内容8 个部分

查看概览图

MOE VARIABLE-LENGTH MATMUL

一次 kernel 算完
所有 expert 的变长 matmul

既不像 for-loop 那样反复 launch 小矩阵，也不像 padding 那样浪费算力——靠 permute 分组 + 运行时 tile 调度处理负载不均。

PER EXPERTY_e = X_eW_e·SHAPESm_e × K · K × N

一批 K、N 相同但 M（行数）逐个不同 的 GEMM。m_e 各不相同、数据相关，sum_e m_e = num_tokens × top_k —— 常规 GEMM / batched GEMM 都假设形状统一，套不上。

✗朴素做法 ①

for-loop 逐 expert 调 GEMM

E 次 kernel launchlaunch 开销累加
每个 GEMM 的 M 很小SM 占用低、tensor core 喂不饱
expert 之间无法重叠GPU 大量空转

✗朴素做法 ②

padding 到 capacity + batched GEMM

padding 是纯浪费的算力C 通常要留余量
超过容量的 token 被 drop影响效果
C 难调调小掉 token，调大浪费

ROUTED乱序

PERMUTE按 expert 连续

TILE 调度block_m × block_n

UN-PERMUTE还原 + gating 加权

permute / un-permute 的索引是同一份，反着用即可

1
拿到自己的全局 tile id 一次 launch 覆盖所有 tile 的 CTA
2
用 group_offsets 二分 / 前缀查找 group_sizes = [m_0, m_1, …, m_{E-1}]group_offsets = [0, m_0, m_0+m_1, …] tile → expert 的映射是数据相关的、运行时算出来的
3
取对应的 W_e 与 X_sorted 那一段 算这个 tile 的 matmul 每个 m_e 向上对齐到 block_m（如 128），tile 不跨越两个 expert

Layout场景做法

Contiguous
（连续）训练 / prefilltoken 沿 M 维拼接，每个 expert 的段按 block 对齐；靠 m_indices / offsets 定位分组

Masked
（掩码）decode每个 expert 预留固定上限 M，用 mask 标记有效行；形状静态，可被 CUDA Graph 捕获

为什么快

负载在 tile 粒度上自动均衡

只有一次 launch，开销摊薄；大 expert 自然分到更多 tile，小 expert 少分；persistent CTA 连续处理 tile，可静态步进或从队列领取；不是抢占式执行，收益需实测。

permute → 单 kernel → un-permutetile → expert 运行时决定MoE 高效训练 / 推理的地基

Grouped GEMM（分组矩阵乘，也叫 ragged / variable-shape GEMM）是一类**用单次 kernel launch 完成「一批形状各不相同的独立 GEMM」**的算子。它是 MoE（混合专家）训练与推理的核心算子。

一句话版本：

一堆 token 被路由到不同 expert，每个 expert 拿到的 token 数 m_e 不一样。Grouped GEMM 把这 E 个变长 matmul 合到一个 kernel 里，用 tile 级调度处理负载不均——既不像 for-loop 那样反复 launch 小矩阵，也不像 padding 那样浪费算力。

1. 它要解决的问题：MoE 里的变长 matmul#

MoE 一层的计算长这样：每个 token 的隐藏态 $h_t$ 先过 router，被分给 top-k 个 expert；expert $e$ 是一个 FFN，权重 $W_e$ 。设 expert $e$ 收到了 $m_e$ 个 token：

Y_e = X_e\, W_e, \qquad X_e \in \mathbb{R}^{m_e \times K},\ \ W_e \in \mathbb{R}^{K \times N}

难点在于：

m_e 各不相同（负载不均，热门 expert 拿到的 token 多）
sum_e m_e = num_tokens * top_k     # 统计所有专家且不丢 token 时成立

这是一批 K、N 相同但 M（行数）逐个不同 的 GEMM。常规 GEMM/batched GEMM 都假设形状统一，套不上。

2. 朴素做法为什么不好#

① for-loop 逐 expert 调 GEMM

for e in range(E):
    Y[e] = X[e] @ W[e]   # 每个 expert 一次 cuBLAS 调用

E 次 kernel launch，launch 开销累加；
每个 GEMM 的 M 很小，SM 占用低、tensor core 喂不饱；
expert 之间无法在一个 kernel 里重叠，GPU 大量空转。

② padding 到 capacity + batched GEMM

把每个 expert 填到固定容量 C（不足补零；若采用固定容量丢弃策略，超出部分丢弃），形状统一后用 batched GEMM：

padding 部分是纯浪费的算力（C 通常要留余量）；
超过容量的 token 被 drop，影响效果；
C 难调：调小掉 token，调大浪费。

Grouped GEMM 就是为了既避免反复 launch，又避免 padding 浪费。

3. 先补 baseline：一个普通 GEMM 在 GPU 上怎么算#

要看懂 grouped GEMM 的「tile 级调度」，得先知道 一个普通 matmul 是怎么被切块映射到 GPU 上的——grouped 只是在这套机制上改了一处。

算 $C_{[M,N]} = A_{[M,K]} \times B_{[K,N]}$ ，GPU 不会一个元素一个元素地算，而是 分块（tiling）：

把输出 C 切成一堆 block_m × block_n 的小块（tile）
每个 tile 交给「一个线程块 CTA」负责，CTA 内部：
    acc = 0
    for k in range(0, K, block_k):           # 沿 K 维分段累加
        a = load A[tile_row, k:k+block_k]     # 一小块搬进 shared memory
        b = load B[k:k+block_k, tile_col]
        acc += a @ b                          # tensor core 做 block 级 MMA
    C[tile] = acc                             # 写回这一个输出 tile

           N 方向 →
        ┌────┬────┬────┐
   M    │ T0 │ T1 │ T2 │   每个 Tk = 一个 block_m×block_n 输出 tile
   ↓    ├────┼────┼────┤   由一个 CTA 负责；沿 K 维循环累加
        │ T3 │ T4 │ T5 │
        └────┴────┴────┘
   tile 总数 = ⌈M/block_m⌉ × ⌈N/block_n⌉

要点：

一次 launch 就发出「覆盖所有 tile」的 CTA，硬件调度器把它们分批（wave）铺到各个 SM 上跑；
分块是为了 数据复用（A/B 的小块进 shared memory 反复用）和 喂饱 tensor core（MMA 指令天然要块状输入）；
关键观察：在一个普通 GEMM 里，每个 tile 用的 A、B 是同一对。

grouped GEMM 的全部改动，就在最后这句话上——让不同 tile 去用不同的 (X_e, W_e)。下一节的 permute + tile 调度，本质就是「把上面这套分块，扩展到 tile 属于哪个 expert 是运行时才知道的情形」。

4. Grouped GEMM 的做法#

围绕 Grouped GEMM 的完整 MoE 数据流是：permute（按 expert 分组） → GEMM kernel 跑各组 tile → un-permute（还原 + gating 加权）。置换和加权通常是独立算子，不是任意 Grouped GEMM kernel 都包含它们。

4.1 Permute：按 expert 排序成连续块#

token 路由后是乱序的（图最左）。先做一次 gather，把属于同一 expert 的行排到一起，得到分组连续的 X_sorted，同时记录每组大小 m_e 与前缀和偏移：

group_sizes   = [m_0, m_1, ..., m_{E-1}]
group_offsets = [0, m_0, m_0+m_1, ...]      # 前缀和，标记每个 expert 在 buffer 里的起止

为了让 tile 不跨越两个 expert，通常把每个 m_e 向上对齐到 block_m（如 128）——区域之间留一点对齐 padding，但远小于 capacity 方案的浪费。

4.2 单 kernel + tile 级调度#

把所有 expert 的输出切成统一的 block_m × block_n tile。一个 kernel 启动足够覆盖所有 tile 的线程块（CTA），每个 CTA 在运行时：

1. 拿到自己的全局 tile id
2. 用 group_offsets 二分/前缀查找，判断这个 tile 落在哪个 expert
3. 取对应的 W_e 和 X_sorted 的那一段，算这个 tile 的 matmul

关键点：tile → expert 的映射是数据相关的、运行时算出来的（因为 m_e 直到 router 之后才知道）。这样：

只有一次 launch，开销摊薄；
大 expert 自然分到更多 tile，小 expert 少分，负载在 tile 粒度上自动均衡；
persistent kernel 让 CTA 连续处理多个 tile，可用静态步进或工作队列分配；这不是对正在运行的 tile 做抢占，占用率仍需实测。

4.3 Un-permute：还原顺序 + gating 加权#

把分组输出 Y_sorted scatter 回原 token 顺序，并按 router 的 gating 权重对 top-k 结果加权求和，得到最终输出。permute/un-permute 的索引是同一份，反着用即可。

5. 两种常见 layout（以 DeepGEMM 为例）#

DeepSeek 的 DeepGEMM（FP8 MoE GEMM）提供两种 grouped 模式，对应不同场景：

Layout	场景	做法
Contiguous（连续）	训练 / prefill	token 沿 M 维拼接，每个 expert 的段按 block 对齐；靠 `m_indices` / offsets 定位分组
Masked（掩码）	decode	每个 expert 预留固定上限 M，用 mask 标记有效行；形状静态，可被 CUDA Graph 捕获

decode 阶段每步 token 很少、对启动开销敏感，masked layout 让 shape 固定从而能用 CUDA Graph 消除 launch 开销；训练/prefill token 多，contiguous layout 更省。

6. 实现路径#

CUTLASS GroupedGemm   # 设备端传入 problem_sizes[] 数组，工业级模板
Triton grouped GEMM   # 官方 tutorial 有 group-gemm；MoE 框架（如 Megablocks）也用 Triton 写
DeepGEMM              # DeepSeek FP8，contiguous + masked 两种 grouped 接口
Megablocks            # 把 MoE 表达成 block-sparse matmul，是 grouped GEMM 的另一种视角

题外话：Megablocks 用 block-sparse 矩阵乘 来看 MoE——把 (token, expert) 的分配看成一个块稀疏矩阵，避免 padding 和 drop。和 grouped GEMM 是殊途同归的两种表述。

7. 要点回顾#

Grouped GEMM = 把一批「行数各异、权重各异」的 matmul 合进单次 kernel，靠 permute 分组 + 运行时 tile 调度处理负载不均。 它让 MoE 既不用反复 launch 小 GEMM，并减少固定容量 padding 的浪费（仍可能存在 tile 对齐开销），是 MoE 高效训练/推理的地基。

实现参考：Triton Group GEMM 教程。