Grouped GEMM 为什么适合 MoE？从碎片小矩阵到通信友好的算子设计
不只是少一次 kernel launch，而是让碎片化专家计算、L2 复用与 dispatch/combine 通信进入同一条流水线
适合：懂 Transformer/MoE 基础的工程读者，阅读约 9 分钟。Grouped GEMM 译为分组矩阵乘。
MoE 让大模型拥有更多参数，但也把原本规整的 FFN 计算，变成了许多大小不一、负载不均的小矩阵乘。
在 Dense FFN 里，一个 batch 通常对应少数几个大 GEMM。GPU 喜欢这种形态：矩阵足够大，tile 足够多，Tensor Core 容易喂饱。
到了 MoE，router 会把 token 分给不同 expert。每个 expert 拿到的 token 数不同，于是每个 expert 要做的 GEMM 的 M 维也不同。有的 expert 很忙，有的 expert 只有几个 token。NVIDIA 对 MoE 通信的描述中，dispatch 负责把 attention 输出 token 路由到对应 expert，combine 负责把 expert 输出再路由回 attention 输出。换句话说，MoE 的性能问题不是单纯“矩阵乘够不够快”，而是碎片化专家计算如何和 dispatch/combine 通信共存。
看图 1 时，可以重点看中间的专家区域：每个 expert 的 GEMM 形状不完全一样。如果用最直接的写法，就是对 expert 做一个 for-loop：Expert 0 启一个 GEMM kernel，Expert 1 再启一个 GEMM kernel……这样做的缺点很明显：kernel launch 多、单个小 GEMM 的并行度不够、专家之间负载不均，GPU 很难一直保持高利用率。
Grouped GEMM 的直觉就是：不要把每个 expert 当成一个孤立 GEMM，而是把一组 GEMM 交给同一个 grouped kernel 统一调度。
CUTLASS 的 grouped GEMM 示例把这类 workload 描述为一批 GEMM 操作，每个 GEMM 可以有不同 problem size，矩阵指针、leading dimension 和 problem size 以数组形式传入；这也正是它和普通 batched GEMM 的关键区别。Triton 的 Group GEMM 教程也把它描述为：启动固定数量的 CTA 来计算一组 GEMM，并在 device 侧完成调度。
Grouped GEMM 到底解决什么问题？
Grouped GEMM 不是一个神秘的新数学算子。数学上，它仍然是在算：
C_i = A_i × B_i, i = 0, 1, ..., G-1区别在于，这里的每个i可以代表一个 expert、一个 shard 或一个独立的 GEMM problem。它们的M/N/K、矩阵地址、stride 都可以不同。Grouped GEMM 关心的是：如何把这些独立问题放进一个 kernel 的调度空间里，让 CTA 持续拿 tile 干活。
这件事在 MoE 里尤其重要。因为 MoE 的 token 分布天然不均匀：热门 expert 的M大，冷门 expert 的M小。直接逐 expert 调 GEMM，热门 expert 会拖尾，冷门 expert 又不够大，最后很容易变成“有些 SM 很忙，有些 SM 等活”。
Grouped GEMM 把专家小 GEMM 变成一张 tile 队列。调度器关心的不再是“先做完 Expert 0，再做 Expert 1”，而是“下一个 CTA 应该拿哪个 expert 的哪个 tile”。CUTLASS 文档中也强调，grouped kernel scheduler 的核心职责就是把一组 problem 中的 tile 分配给 threadblock；threadblock 会持续查询下一个 tile，执行 MMA 和 epilogue，再推进调度状态。
这带来三个直接收益：减少 launch，提升小 GEMM 的并行度，缓解 expert 间负载不均。
只做成 grouped 还不够：真正难的是流水线和内存
很多人介绍 Grouped GEMM 时，会停在“多个 GEMM 合成一个 kernel”这一层。但在 MoE 训练或推理里，这还不够。
MoE 层同时有三类压力：expert 计算要读 activation 和 weight，dispatch/combine 要搬 token，下一层或其他并行通信也可能在抢 HBM、NVLink 或 RDMA 资源。因此，一个真正可用的 Grouped GEMM 设计，通常要同时看四个旋钮：tiling、double-buffer pipeline、L2 cache conflict 缓解、HBM 带宽控制。
1. Tiling：把专家不均匀变成 tile 队列
Tiling 的第一层作用，是把不同 expert 的 GEMM 切成统一粒度的 tile。
比如 Expert 1 有 160 个 token，Expert 2 只有 32 个 token。直接按 expert 调度，会看到一个很大的任务和一个很小的任务；切成 tile 后，它们都变成调度器可以分发的工作单元。这样热门 expert 不会独占整个 kernel，冷门 expert 也不会因为太小而浪费一次 launch。
Tiling 的第二层作用，是让 kernel 可以调节寄存器、shared memory、occupancy 和 Tensor Core 使用效率。tile 太小，调度开销和访存开销占比会上升；tile 太大，又容易让长尾 expert 拖住 SM。MoE 场景中，tile size 不是越大越好，而是要配合 token 分布、expert 数、输出维度和通信节奏一起选。
2. Double-buffer pipeline：别让 Tensor Core 等数据
GEMM kernel 的主循环通常在做两件事：从 HBM/L2 把下一块数据搬进 shared memory，同时让 Tensor Core 对当前块做 MMA。
Double-buffer pipeline 的基本做法是准备两套 buffer：当前 buffer 用来 compute，另一个 buffer 用来 load 下一块。下一轮交换角色。理想情况下，数据加载被计算覆盖，Tensor Core 不会频繁饿死。
在 Grouped GEMM 中，这个问题更复杂。因为不同 expert 的矩阵地址不连续，tile 形状和剩余 K 维可能不同。pipeline 不仅要解决单个 GEMM 内部的 load/compute overlap，还要处理跨 expert 的 tile 切换：什么时候预取下一个 expert 的 weight？什么时候切换 activation slice？什么时候停止过度预取，避免把 HBM 带宽打满？这些都属于 grouped kernel 里的工程细节。
3. L2 cache conflict 缓解：不要让 tile 调度互相拆台
L2 cache 在这里像一个共享中转站。Grouped GEMM 想复用 expert weight、activation tile 和 epilogue 相关数据；dispatch/combine 也会产生大量 token 搬运。调度顺序如果太随意，多个 CTA 可能在相近时间访问互相冲突的数据区域，导致 L2 hit rate 下降，最终表现为 HBM 压力增加。
因此，很多优化会围绕 locality 做文章：让相近 tile 更容易复用 L2；让同一 expert 的 weight 不要刚进 cache 就被别的访问模式挤掉；让 tile 排布不要集中冲击同一类地址。PyTorch 的 MoE locality-aware GEMM 文章就展示过，通过改变调度顺序来改善数据局部性，可以显著提升 MoE GEMM 的硬件利用率。
这里的关键不是记住某一种固定顺序，而是理解原则：Grouped GEMM 的调度器不仅是在分配计算，也是在塑造缓存访问模式。
4. HBM 带宽控制：给通信留水位
MoE 里的 GEMM 和通信不是两条互不相干的路。
如果 Grouped GEMM 过度预取，把 HBM 带宽打满，dispatch/combine 可能被挤压；如果为了通信把 GEMM 做得太保守，Tensor Core 又会吃不饱。真正难的是在计算吞吐和通信带宽之间找水位。
工程上可以调的旋钮包括：CTA 数量、pipeline stage 数、prefetch distance、tile size、persistent block 数量、是否按 chunk 分段执行，以及通信 stream 和 compute stream 的重叠策略。
这也是为什么 MoE kernel 不能只看单 kernel TFLOPS。一个单独跑起来很漂亮的 GEMM，如果上线后把 dispatch/combine 挤到尾部，端到端 latency 仍然可能变差。
和 dispatch/combine 一起看：Grouped GEMM 是 MoE 管线的一段
MoE 层可以粗略看成三步：
dispatch tokens -> expert GEMMs -> combine outputs但高性能实现不会真的把这三步完全串行化。更常见的思路是 chunk 化：当前 chunk 做 expert GEMM，下一 chunk 做 dispatch，上一 chunk 做 combine。这样计算和通信可以在时间线上重叠。
这时 Grouped GEMM 的角色会变化：它不只是“尽可能快地完成 expert GEMM”，还要“以不会破坏整体流水线的方式完成 expert GEMM”。
这解释了为什么 HBM 控制、L2 conflict 缓解和 tile 调度会变得重要。MoE 的瓶颈往往不是某一个 kernel，而是多个 kernel、通信库、缓存层级和 HBM 带宽共同形成的系统瓶颈。
Fused GemmAdd：把 FP32 加法塞进 epilogue
最后说 Fused GemmAdd。
一个 GEMM kernel 通常可以分成 mainloop 和 epilogue。mainloop 负责 MMA，把乘加结果累积到 accumulator；epilogue 负责把 accumulator 做缩放、类型转换、写回等后处理。CUTLASS 文档也把 GEMM operator 描述为 MMA 后接 epilogue operation。
如果 GEMM 后面还有一次 FP32 加法，例如：
Y = GEMM(X, W) + Add朴素做法是先让 GEMM 写回C，再启动一个 Add kernel，把C和Add读出来相加，再写回Y。这会多一次 kernel launch，也会多一次 HBM 读写。
Fused GemmAdd 的思路很直接：既然 GEMM accumulator 本来就在 epilogue 里准备写回，那就在写回前完成 FP32 加法：
Y = cast(Acc_fp32 + Add_fp32)或者更一般地写成：
Y = cast(alpha * Acc_fp32 + beta * Add_fp32)Fused GemmAdd 的收益主要来自两点。
第一，减少一次中间结果写回和再读取。对小 GEMM 或 memory-bound 后处理来说，这往往比省一点算术指令更重要。
第二，减少 kernel launch 和调度缝隙。MoE 本来就有很多 expert、很多通信和很多小 kernel，如果能把简单后处理融合进 epilogue，就能让执行图更短、更连续。
需要注意的是，“FP32 加法融合”并不等于最终输出必须保持 FP32。常见做法是 accumulator 和 add 在 FP32 中完成，再按输出要求 cast 到 BF16、FP16 或其他格式。这里要明确数值路径：加法在哪个精度做，什么时候 cast，是否需要 saturation 或 rounding。
什么时候 Grouped GEMM 不一定赚？
Grouped GEMM 很适合 MoE，但不是万能药。
如果 expert 数很少、每个 expert 的 GEMM 已经足够大，普通 GEMM 可能已经能跑满 GPU。此时 grouped 调度带来的收益会变小。
如果 token 分布极端不均，某个 expert 远大于其他 expert，Grouped GEMM 仍然会遇到长尾。此时需要排序、拆分大 expert、调整 tile 或做 persistent scheduling。
如果端到端瓶颈主要在 all-to-all、dispatch/combine 或跨节点通信，单独优化 GEMM 只能改善一部分。你会看到 kernel profile 变漂亮，但 step time 改善有限。
如果 epilogue 里融合太多逻辑，也可能带来寄存器压力、occupancy 下降和更复杂的边界处理。Fused GemmAdd 适合简单、规则、访存代价高的后处理；不是所有后处理都应该塞进 epilogue。
实战 profiling 看什么？
评估 Grouped GEMM，不建议只看 TFLOPS。至少要同时看以下几类指标：
维度 | 重点观察 |
计算 | Tensor Core 利用率、SM occupancy、tile 尾部空转 |
调度 | expert 间负载均衡、CTA 分配、长尾 tile |
缓存 | L2 hit rate、weight/activation 复用情况 |
带宽 | HBM throughput、dispatch/combine 是否被挤压 |
图执行 | kernel launch 数、GEMM 与通信重叠比例 |
数值 | FP32 accumulate、epilogue add、输出 cast 路径 |
一个好的 Grouped GEMM 优化，最终应该体现为端到端收益：MoE 层耗时下降，通信尾部减少，GPU 时间线更紧凑，而不是只有孤立 GEMM microbenchmark 变好。
小结：Grouped GEMM 的认知地图
可以用一句话总结：Grouped GEMM 是 MoE 里把碎片化专家计算重新组织成可调度整体的关键算子。
更细一点，它包含五层设计：
- Grouped GEMM 解决的是“很多大小不一的小 GEMM 如何合并调度”；
- Tiling 解决的是“专家负载不均如何变成 tile 粒度的工作队列”；
- Double-buffer pipeline 解决的是“Tensor Core 如何不等 HBM 数据”；
- L2 conflict 缓解和 HBM 带宽控制解决的是“计算如何不把通信挤死”；
- Fused GemmAdd 解决的是“简单 FP32 后处理如何少一次 launch、少一次读写”。
所以，MoE 里的 Grouped GEMM 不只是一个 GEMM 优化，而是一个小型系统工程：它连接 router 之后的专家计算，也连接 dispatch/combine 通信；它既要追求 Tensor Core 吞吐，也要尊重缓存和 HBM 的边界。
真正的优化目标不是“某个 GEMM kernel 跑得最猛”，而是MoE 这一层整体跑得最稳、最短、最少互相干扰。