M2 讲解：fused MoE grouped GEMM

① 先看全貌：一层 MoE 的五步

drillsmoe_layer（79–91 行）

W13 是每个专家的 gate 和 up 两个投影拼在一起（[E, 2I, H]，前 I 行是 gate，后 I 行是 up），和 F3 里 gate_up_proj 的布局一样。

② grid：每个 M 块只属于一个专家

drills28–34 行 · host 端 drillsgrid（71 行）

grid = (sorted_ids.numel() // block_m, cdiv(N, block_n))

pid_m, pid_n = tl.program_id(0), tl.program_id(1)
offs_m   = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
flat_ids = tl.load(sorted_ids_ptr + offs_m)            # 这一块的 BLOCK_M 个 (token, slot) 编号
row_mask = flat_ids < num_valid                        # padding 行的值是 T·K，被 mask 掉
expert   = tl.load(expert_ids_ptr + pid_m)             # 这一块属于哪个专家

M1 的例子（T = 4、K = 2、block_m = 2）里，sorted_ids 长 10，有 5 个 M 块：

每个 program 只读一个专家的权重 W[expert]。划掉的 8 是 padding，row_mask 为 False。GEMM1 读 x 的行号是 flat_id // K。

和 T4 的普通 GEMM 相比，只有两处不同：A 的行不连续（要通过 flat_ids 去 gather），B 的起点由 expert 决定（w_ptr + expert * stride_we）。

③ A 的 gather：GEMM1 和 GEMM2 不一样

drills36–46 行

if A_ROWS_ARE_TOKENS:          # GEMM1：A 是 x [T, H]
    a_rows = flat_ids // top_k
else:                          # GEMM2：A 是 h [T·K, I]
    a_rows = flat_ids
a_ptrs = a_ptr + a_rows[:, None] * stride_am + offs_k[None, :]
w_ptrs = w_ptr + expert * stride_we + offs_n[None, :] * stride_wn + offs_k[:, None]

A_ROWS_ARE_TOKENS 是 constexpr，所以两种 gather 是同一个 kernel 编译出来的两个版本。

W 的转置：W[e] 存成 [N, K] 的行优先（PyTorch Linear 的权重格式），GEMM 需要的是 Wᵀ。和 T4、T5 一样，用 stride 直接读成 [BLOCK_K, BLOCK_N] 的转置块：offs_n[None, :] * stride_wn + offs_k[:, None]。

为什么 gather 比连续读慢：同一个 M 块里的行在 x 里是分散的（t0、t2），每一行是连续的 H 个数，所以每行内部仍然是合并访存，但不同行之间无法连成更大的连续块，硬件预取和 TMA 之类的块拷贝也用不上。题目页提到这是和稠密 GEMM 差距的原因之一。

④ 主循环和 epilogue

drills48–61 行

主循环和 T4 一样：沿 K 方向每次取 BLOCK_K，tl.dot 累加到 fp32 的 acc。mask 多了一项 row_mask[:, None]，padding 行读 0。

if MUL_WEIGHT:                                  # 只在 GEMM2 打开
    tw = tl.load(topk_w_ptr + flat_ids, mask=row_mask, other=0.0)
    acc = acc * tw[:, None]
out_ptrs = out_ptr + flat_ids[:, None] * stride_om + offs_n[None, :]
tl.store(out_ptrs, acc.to(...), mask=row_mask[:, None] & (offs_n[None, :] < N))

• 输出写回到 flat_id 对应的行（scatter）：out 是 [T·K, N]，按 (token, slot) 的原始顺序排列，不是排序后的顺序。所以最后可以直接 view(T, K, H).sum(1) 把同一个 token 的 K 个结果加起来。
• 路由权重为什么在 GEMM2 里乘：最终输出是 Σ_k w_k · expert_k(x)，而 expert(x) = W2 · silu_and_mul(W13 · x)。权重是标量，可以挪到最后一步乘。放在 GEMM2 的 epilogue 里，几乎不花额外时间，也省掉一次单独的逐元素 kernel。如果放在 GEMM1 里，silu 是非线性的，结果就不对了。
• out 初始化为 0（drills69 行）：所有合法行都会被写到，不清零也可以；清零是为了让未被覆盖的情况更容易被发现。

⑤ 性能：为什么只有稠密 GEMM 的一半

题目页的实测：T = 4096、E = 64、top-8、H = 2048、I = 768（接近 Qwen3-30B-A3B），bf16。

差距的来源（题目页列出的几条）：

• 固定的 tile 配置：64×64×32 对这个形状不是最优。vLLM 为每种 (E, N, dtype, GPU) 准备了调好的 JSON 配置，仓库里有三百多个 vllmget_config_file_name 指向的配置文件，运行时由 vllmtry_get_optimal_moe_config 选择。
• padding：每个专家最多浪费 BLOCK_M − 1 行（M1 第 ③ 节）。
• gather：A 的行不连续。
• silu_and_mul 和最后的求和是单独的 PyTorch 操作：多了中间结果的写回和读取。

⑥ 对照 vLLM 的 fused_moe_kernel

参考解的结构就是从 vLLM 抄来的：vllmfused_moe_kernel

vLLM 的版本还处理了 FP8 / INT8 的 per-tensor 和 block-wise scale（M3 会讲 block scaling）、专家并行，以及按 GROUP_SIZE_M 分组的 program 顺序（T4 讲过的 L2 优化）。

⑦ 常见坑与自测

• GEMM1 用了 a_rows = flat_ids：读 x 越界（x 只有 T 行）。
• GEMM2 用了 flat_ids // top_k：读错了中间结果，每个 slot 都用了 slot 0 的 h。
• 路由权重乘了两次（GEMM1 和 GEMM2 各一次），或者乘在 silu 之前。
• out 按排序后的顺序写：最后的 view(T, K, H).sum(1) 就对不上 token 了。
• expert 没转 int64：expert * stride_we 在大模型上溢出（E = 256、每个专家几千万个参数）。

源码与资料索引

本仓库：drillssolutions/m2_fused_moe.py · drillsmine/m2_fused_moe.py · drillstests/test_m2_fused_moe.py · drillscontent/m2.md · drillsdrill.py（load() 选择你的 M1 还是参考解）

vLLM：vllmfused_moe_kernel · vllmtry_get_optimal_moe_config · vllmmoe_align_block_size

Triton：tritontutorial 08 grouped GEMM（另一种做法：每个 group 一个独立的 GEMM 描述）

论文：Mixtral · MegaBlocks（block-sparse 的 MoE 计算）