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

1.9 KiB

Raw Blame History

PR347 基础设施/SM100 重构代码review报告

范围

Worktree: pr-347
HEAD: 2b8dfd0e8
审查方式: 代码review
主题: SM100 MegaMoE 重构，SM90 MegaMoE kernel 已被移除

实现概述

核心变化

pr-347 是一个中间重构 PR。它不包含 SM90 MegaMoE kernel（相比 pr-323 已移除 sm90_fp8_mega_moe.cuh），重点关注:

SM100 FP8/FP4 MegaMoE kernel 重构 — 代码格式、变量命名、注释统一
Heuristics 模块重组 — mega_moe.hpp 从 211 行扩展到 276 行，增加了更多 shape 分支
Scheduler 增强 — mega_moe.cuh 从 221 行扩展到 273 行，支持更多 block 分配策略
API 层清理 — 移除 SM90 专用的 csrc/apis/sm90_mega.hpp

SM100 FP8/FP4 MegaMoE Fused Kernel

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

线程和寄存器布局与 pr-304 一致:

Role	Threads	Warps	所属 Warpgroup	Reg/thread
Dispatch	128	4	WG0	48
TMA (A+SFA, B+SFB, MMA issue, idle)	128	4	WG1	40
Math/epilogue + Combine	256	8 (2 warpgroups)	WG2, WG3	208
总计	512	16	4 WGs	64,512/64K

代码review发现

低: 这是一个重构 PR，kernel 逻辑变化小

主要改动是代码组织（命名、注释、文件结构调整），对运行时行为影响有限。风险评估较低。

低: SM90 MegaMoE 能力被移除

相比 pr-323，sm90_fp8_mega_moe.cuh 和相关 JIT host/heuristics 文件不再存在。如果这是故意的（为后续 PR 清理），则合理；如果是不小心丢失，需要注意。

正面评价

代码风格和注释质量有显著提升
Heuristics 模块增加了更多 shape 分支覆盖
为后续 PR352 的分层 MegaMoE 打下基础

建议检查清单

确认 SM90 MegaMoE 移除是故意的，非 merge 错误
验证重构后 SM100 MegaMoE 的 bitwise correctness 未退化

1.9 KiB Raw Blame History Unescape Escape

PR347 基础设施/SM100 重构 代码review报告

范围