MoE Grouped GEMM与重排：把零碎专家计算组织起来

01 · 一张卡上也可能有许多零碎小矩阵

MoE的每个专家接到不同数量token。即使专家布局已平衡，一张卡上的专家也可能分别只有几行输入；逐专家启动两个Linear及激活，启动与调度成本可能很突出。所有专家统一补到同样行数又会产生空位。

Grouped GEMM把多个矩阵问题交给成组执行方案。不同专家的权重保持独立、有效行数可以不同；它不是把所有专家权重求和，也不是改变router选出的专家。重排则把对应输入组织成内核容易读取的布局。

依据：NVIDIA CUTLASS · Grouped Kernel Scheduler ↗ · NVIDIA cuDNN Frontend 1.18 · MoE Grouped Matmul ↗

02 · 四个token，top-2会产生八次专家调用

本例有四个逻辑专家E0–E3，调用数[3,4,1,0]。总计8次调用不是8个独立token，原输出仍只有A/B/C/D四行。若已有EP通信，以下操作可以在发送前或接收后的合适位置配套；不要求一次全局重排全部设备。

依据：DeepSeek · DeepEP dispatch/combine通信实现 ↗

03 · 排到一起，需要保存什么元数据

按专家排列的输入token：[A,C,D, A,B,C,D, B]
各行逻辑专家：         [0,0,0, 1,1,1,1, 2]
累计边界offsets：      [0,3,7,8,8]
还原token索引：        [0,2,3, 0,1,2,3, 1]
对应gate：             [.7,.5,.2, .3,.6,.5,.8, .4]

offsets[0:2]表示E0输入在[0,3)，E1在[3,7)，E2在[7,8)，E3为空[8,8)。位置顺序改变，但token身份和路由权重不变。不同API可能用每组长度、累计边界或指针数组，不能只看都叫group_list就假设单位相同。

空专家应按接口合同跳过或表示为空，不给它假造一条有效输入。top-k的多份输入需要共享读取或复制的安排；这些额外字节和索引成本也要计入。

依据：NVIDIA cuDNN Frontend 1.18 · MoE Grouped Matmul ↗ · Ascend cann-ops-adv · GroupedMatmulV2接口与分组语义 ↗

04 · Grouped不是必须补成同一shape

每个专家的第一层可以写作Xₑ[Mₑ,K]Wupₑ[K,H]，第二层为activation输出[Mₑ,H]Wdownₑ[H,D]。本例Mₑ=3/4/1/0。普通固定shape的batched GEMM常用统一M；Grouped用各问题的shape和地址选择工作tile。

CUTLASS的一类Grouped实现用持久kernel管理多个问题：工作块完成当前tile后继续领取其他tile。成组执行减少零散启动并组织负载，但具体后端可有不同调度，不保证整层MoE只启动一个kernel。第一Linear、激活、第二Linear及combine仍有各自依赖。

专家大小不同时，tile数量相同也不表示工作量相同；K方向长短会影响每tile耗时。局部调度需要关注工作量分配，这与M030跨设备专家布局是不同尺度的均衡。

依据：NVIDIA CUTLASS · Grouped Kernel Scheduler ↗

05 · 八个有效调用，为什么会算到十六个槽

如果四专家统一补到最大M=4，则有4×4=16行槽，其中8行有效。Grouped按真实行数组织输入可只保存8行，但实际内核仍按固定tile执行边界。M方向tile=2时，覆盖行数为4+4+2+0=10；不是看到“Grouped”就完全没有边界开销。

用教学K=D=2、H=3的两层MLP，忽略激活和其他操作，有效数学工作量为2×8×2×3+2×8×3×2=192 FLOP。补成16行的稠密坐标工作对应384 FLOP；mask与Tensor Core边界如何执行由后端决定。

06 · 算完还要按原token合并

for each packed output row r:
    t = restore_token_index[r]
    output[t] += gate[r] * expert_output[r]
# 例如A = 0.7*E0(A) + 0.3*E1(A)
# 同一个token有多个贡献，要处理并发归并和正确初始化。

可以在适配内核中融合gate乘法、重排或输出写入，减少一次中间物化或额外启动；有些实现通过索引直接读取输入，避免显式生成全部重排副本。但不应把“Grouped GEMM”自动理解为dispatch和combine也已经融合。

多个贡献并发写同一token，需归约或正确的原子操作/分阶段方案；直接覆盖会丢结果。gate放在哪一步还要匹配模型定义，特别是含bias或非线性时，不能随意把输出权重移到输入。

依据：NVIDIA cuDNN Frontend 1.18 · MoE Grouped Matmul ↗ · DeepSeek · DeepEP dispatch/combine通信实现 ↗

07 · 反向也要保持分组与还原关系

对每个专家，输入梯度由该专家输出梯度反传，权重梯度累积该专家所有有效输入的贡献。一个token被top-2复制后，两条反向贡献仍要归并回同一原输入；router/gate梯度则按训练计算图另外处理。

重排的反向通常涉及逆映射和重复索引的累加。只做前向等价检查不够；mask的空行不能进入权重梯度，有效Loss归一化也不能由打包后的行数擅自改变。

dPackedY[r] = gate[r] * dOutput[token_index[r]]
每专家：dW = Xᵀ dPackedY；dX = dPackedY Wᵀ
按token_index累加各份dX回原输入
# 此处展示线性专家且gate固定的反向；完整MLP/router还需链式求导。

08 · CPU例子验证什么，设备接口再研究什么

下载脚本先对每token分别运行小ReLU专家MLP，再按offsets成组执行，核对加权还原输出；包含空专家和top-2贡献。另以固定gate线性专家检查输入梯度，并用有限差分作独立对照，展示错还原索引的反例。

统计每专家有效行数并构造offsets
准备packed输入、token索引与gate
Grouped GEMM(up) → 激活 → Grouped GEMM(down)
按原token归并专家输出
训练时保存反向需要的元数据与激活
# 实际可融合若干步骤；接口先核对shape、layout、dtype与workspace。

NVIDIA可从CUTLASS调度器和cuDNN MoE Grouped Matmul的索引输入阅读；昇腾可从GroupedMatmulV2文档核对groupList、splitItem、groupType及布局。接口名字相近不代表输入边界、量化或输出组织相同。

依据：NVIDIA CUTLASS · Grouped Kernel Scheduler ↗ · NVIDIA cuDNN Frontend 1.18 · MoE Grouped Matmul ↗ · Ascend cann-ops-adv · GroupedMatmulV2接口与分组语义 ↗

09 · 把重排和通信算在总账里

大专家批量已能高效独立执行时，Grouped额外调度未必划算；热点在远端网络时，本地GEMM变快也可能仍等待。要分别开关Grouped、重排融合和专家布局，识别各项收益。

10 · 自查与扩展实验

• 为什么本例4个token最终还是4行输出，而Grouped输入有8行？
• offsets最后两项都是8，表示E3空还是表示数据缺失？
• 把tile从1改成4，有效调用数和边界覆盖行数分别怎样变化？
• 为什么不一定能把gate从专家输出移到专家输入？考虑含bias的非线性MLP。

流程位置与关联方法

这个方法可能出现在多个位置；计算线包含模型计算与算子实现，执行线包含运行时、内存和通信。

资料与核查范围

原理依据以原始论文和官方文档为主；算例与时间线是教学推导。文档核查日期：2026-10-10。引用 main / latest 文档时，实际接口以安装版本为准。

• NVIDIA CUTLASS · Grouped Kernel Scheduler ↗
• NVIDIA cuDNN Frontend 1.18 · MoE Grouped Matmul ↗
• Ascend cann-ops-adv · GroupedMatmulV2接口与分组语义 ↗
• DeepSeek · DeepEP dispatch/combine通信实现 ↗

代码为接入片段或教学伪代码；本页没有声称完成 NVIDIA / 昇腾硬件实测。