Files

Xinyi Liu 062cb160cf Phase 0: SM90 MegaMoE design doc, reference baseline, nsys script

- MEGAMOE_SM90_DESIGN.md: complete design document with finalized decisions
  (fused single kernel, cooperative + single-WG, dynamic BLOCK_M, etc.)
- tests/test_mega_moe_sm90.py: PyTorch FP32/BF16 reference implementation
  for dispatch → L1 GEMM → SwiGLU → L2 GEMM → combine pipeline
- scripts/run_nsys_mega_moe_sm90.sh: nsys profiling wrapper script
- megamoe-research-reports/: research analysis of PR304/323/347/352/357/360

2026-06-16 18:01:12 +08:00

3.6 KiB

Raw Permalink Blame History

PR323 SM90 Fused MegaMoE 代码review报告

范围

Worktree: pr-323
HEAD: 23f46aa68
审查方式: 代码review
主题: 首次引入 SM90/Hopper FP8 MegaMoE 单 kernel 实现

实现概述

SM90 FP8 MegaMoE Fused Kernel

deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh (~1935 lines)

warp 分工与线程布局（3 种配置）:

配置 1: 512-epilogue (dispatch+TMA 各自独占 warpgroup)

Warp Index	数量	所属 Warpgroup	Role
0–3	4 warps (128 threads)	WG0	Dispatch
4–7	4 warps (128 threads)	WG1	TMA A+SFA / B+SFB / MMA issue / idle
8–23	16 warps (512 threads)	WG2–WG5	Math WGMMA + epilogue + combine

精确 warp 统计: 24 warps = 768 threads = 6 warpgroups

配置 2: 256-epilogue (dispatch+TMA 共享 WG0, compact)

Warp Index	数量	所属 Warpgroup	Role
0–1	2 warps (64 threads)	WG0	Dispatch
2–3	2 warps (64 threads)	WG0	TMA A+SFA / B+SFB
4–11	8 warps (256 threads)	WG1, WG2	Math WGMMA + epilogue + combine

精确 warp 统计: 12 warps = 384 threads = 3 warpgroups

配置 3: ≤256-total (dispatch+TMA 共享 WG0, BLOCK_M=32)

Warp Index	数量	所属 Warpgroup	Role
0–1	2 warps (64 threads)	WG0	Dispatch
2–3	2 warps (64 threads)	WG0	TMA A+SFA / B+SFB
4–7	4 warps (128 threads)	WG1	Math WGMMA + epilogue + combine

精确 warp 统计: 8 warps = 256 threads = 2 warpgroups |

关键差异 vs SM100:

无 TMEM，WGMMA accumulator 在 register 中
SF 为 per-128 float（非 UE8M0 int）
cluster_size ≤ 2（无 2-CTA UMMA）
dispatch 最小仅 64 threads（SM100 要求 ≥128），意味着 dispatch 可以和 TMA 共享 warpgroup

寄存器分配（多配置）:

配置	Dispatch Threads	TMA Threads	Epilogue Threads	Total Threads	Disp Reg	TMA Reg	Epi Reg	Total Reg
512-epilogue	128	128	512	768	32	24	112	64,512
256-epilogue	64	64	256	384	48	40	168	48,640
≤256-total	64	64	128	256	48	40	256	38,400

kEpilogueRegisterBudget 模板参数允许调用方精确控制 math warpgroup 的 register 分配。

调度特点:

sm90_fp8_mega_moe_for_each_block_split 将 L1/L2 phase 拆分到不同 CTA，每个 CTA 只执行一种 phase
使用 kNonEpilogueWarpsInWarpgroup 确保 dispatch+TMA 恰好填满整数个 warpgroup

代码review发现

高: 512-epilogue 配置下 dispatch register 仅 32，可能成为瓶颈

dispatch warp 需要大量 smem_expert_count 和 route metadata 操作。32 register/warp 对于复杂的 read_topk_idx lambda 和 rank round-robin 选择可能不够（register spilling 到 local memory）。

中: `kEpilogueRegisterBudget` 默认 0 的自动推导逻辑复杂

sm90_fp8_mega_moe.cuh:427-432 的自动 register budget 推导依赖 thread 数和总线程数，外部调用者可能不知情使用次优配置。

中: 与 SM100 kernel 共享 scheduler，但 SM90 无 cluster 概念

mega_moe.cuh scheduler 包含 kClusterSize=2 相关逻辑，SM90 cluster_size=1 时跳过但增加了编译期复杂度。

低: 256-total-thread 配置下 epilogue 获得 256 reg/warp，几乎占满 budget

正面评价

首次将 MegaMoE 带到 SM90/Hopper，填补了架构空白
灵活的 register budget 控制允许调用方按场景调优
L1/L2 phase 拆分减少了单一 CTA 的复杂度

建议检查清单

在 H100/H200 上实测 512-epilogue 配置的 dispatch spilling 情况
简化或文档化 kEpilogueRegisterBudget 的默认行为

3.6 KiB Raw Permalink Blame History Unescape Escape