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.2 KiB

Raw Permalink Blame History

PR304 SM100/FP8/FP4/BF16 GEMM Suite 代码review报告

范围

Worktree: pr-304
HEAD: 211d2678d
审查方式: 代码review
这是 DeepGEMM 主分支的早期 snapshot，包含完整的 SM100 和部分 SM90 kernel 套件

实现概述

SM100 FP8/FP4 MegaMoE Fused Kernel

deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (~1364 lines)

warp 分工与线程布局 (硬编码 dispatch=128, TMA=128, epilogue=256):

Warp Index	数量	Threads	Role	所属 Warpgroup
0–3	4 warps	128	Dispatch: 扫描 topk_idx，写 expert send/recv count，NVLink pull token+SF	WG0
4	1 warp	32	TMA A + SFA loader	WG1
5	1 warp	32	TMA B + SFB loader	WG1
6	1 warp	32	MMA issue warp (仅 leader CTA)	WG1
7	1 warp	32	Idle (warpgroup 占位)	WG1
8–15	8 warps	256	Math UMMA + SwiGLU epilogue + BF16 scatter + Combine reduce	WG2, WG3

精确 warp 统计: 16 warps = 512 threads = 4 warpgroups |

寄存器分配 (典型配置: dispatch=128, TMA=128, epilogue=256):

Role	每 thread register 数	Threads	总 register 消耗
Dispatch	48	128	6,144
TMA/non-epilogue	40	128	5,120
Math/epilogue	208	256	53,248
总计		512	64,512 (恰好 64K budget)

关键约束:

kNumDispatchThreads % 128 == 0 → dispatch 独占整数个 warpgroup
kNumNonEpilogueThreads == 128 → TMA 恰好一个 warpgroup
kNumEpilogueThreads % 128 == 0 → math 独占整数个 warpgroup
2-CTA cluster MMA (SM100 UMMA)，cluster_size=2

其他 SM100 Kernel

Kernel	线程布局	备注
`sm100_fp8_gemm_1d1d`	dispatch≥128 + TMA=128 + math≥128	FP8 1D-1D grouped GEMM，TMEM accumulator
`sm100_fp8_fp4_gemm_1d1d`	同上	FP8/FP4 混合精度，UE8M0 SF packing
`sm100_bf16_gemm`	同上	BF16 GEMM, TMEM, stage merging
`sm100_fp8_mqa_logits`	cluster=2, 128 specialized + 256 math	FP8 MQA attention logits
`sm100_tf32_hc_prenorm_gemm`	cluster=2, dispatch+math	HyperConnection prenorm GEMM

代码review发现

中: SM100 MegaMoE 调度器波次边界越界读取

deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh:73-81 中 get_num_tokens() 声明的 valid_value 未初始化，在波次最后一个 expert 完成时 advance_expert_idx() 会在无人拥有该 expert index 的 warp 中调用 get_num_tokens(wave_end)。

中: SM100 kernel 的 register budget 在默认配置下已达上限

dispatch=12848 + TMA=12840 + epilogue=256*208 = 64512，恰好等于 64K reg budget（SM 寄存器总数 65536，扣除 CUDA runtime 开销后约 64512 可用）。任何 register 增加都会溢出。

低: TMA TMA load warp 中 `warp_idx == kNumDispatchWarps + 3` 的 idle warp 仅执行 dealloc，无实际工作

正面评价

完整覆盖 SM100 GEMM、MegaMoE、MQA attention、paged attention 场景
2-CTA cluster 充分利用 SM100 硬件能力
HEURISTICS 文件结构清晰，按架构分离

建议检查清单

修复 scheduler valid_value 未初始化
所有 kernel 的 warpgroup_reg_dealloc 参考 SM90 已验证模式，当前无问题

3.2 KiB Raw Permalink Blame History Unescape Escape