- 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
22 KiB
PR360 SM90 Pingpong/Cooperative MegaMoE 代码review报告
范围
- Worktree:
pr-360 - Remote 分支:
pr-360/sm90-mega-moe-pingpong-coop - HEAD:
f27fd561164919925bf49aa2d015a1de5ec281db - 审查 diff 范围:
88965b078..f27fd5611 - 变更规模: 13 个文件,+6200 / -7 行
- 审查方式: 代码 review + 后续 8×H200 编译/benchmark 补测;未运行 correctness test 或 NCU profile
Commit 列表
f4f8473cdsm90 megamoe: add pingpong and cooperative FP8 kernels3f10f93a1sm90 megamoe: token-count dispatch and Python bindings89e3d351esm90 megamoe: correctness test and performance benchmark7b65ae9b4chore: ignore profiling artifacts (*.ncu-rep, *.nsys-rep, profile_logs/)ba270a419fix mul20e6549b60add sm90 megamoe benchf27fd5611Revert "chore: ignore profiling artifacts (*.ncu-rep, *.nsys-rep, profile_logs/)"
架构概述
- 为 Hopper/SM90 架构新增 FP8 MegaMoE kernel,与现有 SM100 FP8/FP4 MegaMoE 路径独立。
pingpong使用BLOCK_M=64,两个 math warpgroup 分别处理不同 tile,通过 order barrier 交替执行 MMA 和 epilogue。cooperative使用BLOCK_M=128,两个 math warpgroup 按 M 行协作切分同一 tile,共享一次 B/weight TMA load,将权重 HBM 读取量减半。- 两者均采用 persistent CTA-per-SM launch,共 64 dispatch threads + 64 TMA/non-epilogue threads + 256 math/epilogue threads。
- 新增 public API:
fp8_mega_moe、fp8_mega_moe_pingpong、fp8_mega_moe_cooperative、transform_weights_for_mega_moe_sm90。
8×H200 实验验证补充
本节为后续补测结果,目的是验证 PR 描述中的性能表是否能在单节点 8×H200 上复现。报告更新本身未重新访问远端机器。
运行环境与准备:
- 机器: g0033 单节点 8×H200,8 ranks,每 rank 1 GPU。
- 代码路径:
/mnt/beegfs/lxy/sm90-bench/pr-360。 - Python 环境:
/mnt/beegfs/lxy/venvs/pr360-torch212-cu130,torch 2.12.0+cu130,tilelang 0.1.9。 - PR360
_C与/sgl-workspace/DeepEP/deep_ep_cpp均按 Torch 2.12/CUDA 13.0 临时重建;运行后默认.so已恢复为 basetorch 2.9.1+cu130可 import 状态,Torch 2.12 构建产物另存为.torch212cu130.so。 - 为适配 PyTorch 2.12 的 tensor metadata 限制,本地验证版本在
csrc/apis/mega.hpp的slice_input_bufferslambda 中加入torch::NoGradGuard no_grad;。这是验证补丁,不属于原始 PR360 diff。 - 环境变量:
NVSHMEM_IBGDA_ENABLE=0 NVSHMEM_DISABLE_IBGDA=1 EP_DISABLE_GIN=1 NCCL_NVLS_ENABLE=0,并用PYTHONPATH=/sgl-workspace/DeepEP:$PYTHONPATH加载重建后的 DeepEP。 NCCL_NVLS_ENABLE=0是必要项;否则 Torch 2.12/NCCL 2.29 会尝试 NVLS multicast,并报Failed to bind NVLink SHARP (NVLS) Multicast memory ... CUDA error 401。- PyTorch symmetric memory multicast 仍报
the operation cannot be performed in the present state,但 Torch 2.12 会Gracefully skipping multicast initialization后走 fallback;因此本次验证不代表 CUDA multicast/NVLS 路径可用。
Baseline 覆盖限制:
v1-contig和v1-ll已验证。v2 ElasticBuffer未能验证,因为当前容器内 DeepEP API 没有deep_ep.ElasticBuffer;所有 PR 描述中关于 v2 的加速比本次未独立确认。
DeepSeek-V4 Flash (hidden=4096, intermediate=2048, experts=256, topk=6)
| tokens | fused (us) | TFLOPS | HBM GB/s | v1-contig (us) | v1-ll (us) |
|---|---|---|---|---|---|
| 16 | 241.9 | 19.4 | 3232 | 1685.7 | 264.3 |
| 64 | 251.4 | 70.5 | 3226 | 1698.2 | 290.6 |
| 256 | 317.6 | 245.0 | 2616 | 2213.3 | 349.7 |
| 512 | 336.8 | 464.4 | 2542 | 1670.8 | 459.3 |
| 1024 | 586.8 | 527.0 | 1544 | 2125.0 | 711.5 |
| 4096 | 1799.1 | 684.6 | 670 | 4235.9 | skipped |
| 8192 | 3579.6 | 689.2 | 449 | 6951.4 | skipped |
相对 PR 描述表中的 fused 时间,偏差范围约 -3.0% ~ +0.5%,可视为复现通过。
DeepSeek-V4 Pro (hidden=7168, intermediate=3072, experts=384, topk=6)
| tokens | fused (us) | TFLOPS | HBM GB/s | v1-contig (us) | v1-ll (us) |
|---|---|---|---|---|---|
| 16 | 743.7 | 17.2 | 3912 | 2674.9 | 809.4 |
| 64 | 818.6 | 59.4 | 3886 | 2782.4 | 902.4 |
| 256 | 1015.4 | 197.5 | 3164 | 2688.5 | 1002.7 |
| 512 | 1063.0 | 387.1 | 3064 | 2750.7 | 1181.7 |
| 1024 | 1622.7 | 510.6 | 2061 | 3286.3 | 1604.4 |
| 4096 | 4684.3 | 689.4 | 821 | 7250.6 | skipped |
| 8192 | 10048.5 | 644.7 | 450 | 12512.8 | skipped |
相对 PR 描述表中的 fused 时间,偏差范围约 -6.1% ~ -0.3%。趋势与 PR 描述一致:大 batch 下 fused 仍明显快于 v1-contig,但本次不能确认其相对 v2 的回退幅度。
MiMo-V2.5 (hidden=4096, intermediate=2048, experts=256, topk=8)
| tokens | fused (us) | TFLOPS | HBM GB/s | v1-contig (us) | v1-ll (us) |
|---|---|---|---|---|---|
| 16 | 246.7 | 25.1 | 3273 | 1684.0 | 271.4 |
| 64 | 248.7 | 97.5 | 3269 | 1637.8 | 299.6 |
| 256 | 323.7 | 326.4 | 2594 | 1597.7 | 386.6 |
| 512 | 469.9 | 443.2 | 1858 | 1900.1 | 541.6 |
| 1024 | 731.6 | 561.1 | 1283 | 2861.3 | 889.9 |
| 4096 | 2322.2 | 708.9 | 578 | 4933.5 | skipped |
| 8192 | 4677.9 | 703.4 | 401 | 8446.8 | skipped |
相对 PR 描述表中的 fused 时间,偏差范围约 -2.6% ~ +0.7%,复现度较高。
MiMo-V2.5-Pro (hidden=6144, intermediate=2048, experts=384, topk=8)
| tokens | fused (us) | TFLOPS | HBM GB/s | v1-contig (us) | v1-ll (us) |
|---|---|---|---|---|---|
| 16 | 471.1 | 20.5 | 3371 | 2210.8 | 518.8 |
| 64 | 502.5 | 73.9 | 3628 | 2375.6 | 578.2 |
| 256 | 642.6 | 242.4 | 2892 | 2331.6 | 677.7 |
| 512 | 675.7 | 454.0 | 2817 | 2232.6 | 857.7 |
| 1024 | 1226.9 | 502.9 | 1627 | 3110.1 | 1365.8 |
| 4096 | 3663.4 | 680.5 | 698 | 6787.8 | skipped |
| 8192 | 7482.0 | 663.3 | 440 | 11709.8 | skipped |
相对 PR 描述表中的 fused 时间,偏差范围约 -3.7% ~ -1.1%,复现度较高。
实验结论:
- PR360 fused kernel 的主性能结论基本成立:四个公开 shape 的 fused 时间均能在 8×H200 上复现,偏差通常在几个百分点内。
v1-ll与 PR 描述基本吻合;v1-contig波动较大,部分点偏差超过 10%,其中 DeepSeek-V4 Flashtokens=256明显偏慢(2213.3 usvs PR 表1525.0 us)。这不影响 fused kernel 复现结论,但说明 baseline 数值对 DeepEP/TileLang 版本和运行环境更敏感。- 由于
ElasticBuffer缺失,本次不能验证 PR 描述中“相对 DeepEP V2”的结论;合入前若要以 v2 对比作为依据,需要在包含deep_ep.ElasticBuffer的 DeepEP 版本上重跑。 - 本次所有 fused/baseline 结果均是在 symmetric-memory multicast fallback 条件下测得;如果目标环境能正确启用 NVLS/multicast,性能可能不同。
实验日志:
- DeepSeek-V4 Flash:
/mnt/beegfs/lxy/logs/pr360_h200_deepseek_flash_h4096_i2048_e256_topk6_baseline_v1_torch212.log - DeepSeek-V4 Pro:
/mnt/beegfs/lxy/logs/pr360_h200_deepseek_pro_h7168_i3072_e384_topk6_baseline_v1_torch212.log - MiMo-V2.5:
/mnt/beegfs/lxy/logs/pr360_h200_mimo_v25_h4096_i2048_e256_topk8_baseline_both_torch212.log - MiMo-V2.5-Pro:
/mnt/beegfs/lxy/logs/pr360_h200_mimo_v25_pro_h6144_i2048_e384_topk8_baseline_v1_torch212.log
实现概述
Pingpong Kernel (BLOCK_M=64)
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh (1418 lines)
warp 分工与 WGMMA 流水线:
| Warp Index | 数量 | Threads | Role | 所属 Warpgroup |
|---|---|---|---|---|
| 0, 1 | 2 (64 threads) | 64 | Dispatch: NVLink pull token+SF, topk weight copy, TMA store to L1 pool | WG0 |
| 2 | 1 | 32 | TMA A (acts) + SFA loader | WG0 |
| 3 | 1 | 32 | TMA B (weights) loader | WG0 |
| 4–7 | 4 (128 threads) | 128 | Math WG0: m64 WGMMA → pingpong epilogue → TMA store L1 / BF16 scatter | WG1 |
| 8–11 | 4 (128 threads) | 128 | Math WG1: m64 WGMMA → pingpong epilogue → TMA store L1 / BF16 scatter | WG2 |
OrderedSequenceBarrier pingpong 时序:
WG1: [MMA tile 0] [epilogue tile 0] [MMA tile 2] [epilogue tile 2] ...
WG2: [MMA tile 1] [epilogue tile 1] [MMA tile 3] [epilogue tile 3] ...
同一个 tile 的 MMA 和 epilogue 在不同 WG 上交替执行,通过 order_arrive/order_wait 门控 stage 0 (MMA) 和 stage 1 (epilogue) 的访问权。
寄存器分配:
| Role | Reg/thread | Threads | 总 Reg |
|---|---|---|---|
| Dispatch | 48 | 64 | 3,072 |
| TMA (A loader) | 48 | 32 | 1,536 |
| TMA (B loader) | 48 | 32 | 1,536 |
| Math WG0 | 224 | 128 | 28,672 |
| Math WG1 | 224 | 128 | 28,672 |
| 总计 | 384 | 63,488/64K |
关键约束:
kNumDispatchThreads + kNumNonEpilogueThreads == 128→ dispatch (64) + TMA (64) 形成一个完整的 WG0kNumEpilogueWarpgroups == 2→ 恰好两个 math warpgroup 参与 pingpong- WG0 内所有 warp 使用相同的 dealloc count (48),与 SM100 生产代码一致的模式
- BLOCK_M=64,每个 math WG 拥有整个 tile(一次 m64 WGMMA 完成)
Cooperative Kernel (BLOCK_M=128)
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh (1419 lines)
warp 分工:
| Warp Index | 数量 | Threads | Role | 所属 Warpgroup |
|---|---|---|---|---|
| 0, 1 | 2 (64 threads) | 64 | Dispatch (同 pingpong) | WG0 |
| 2 | 1 | 32 | TMA A (acts) + SFA loader | WG0 |
| 3 | 1 | 32 | TMA B (weights) loader | WG0 |
| 4–7 | 4 (128 threads) | 128 | Math WG0: 处理 tile 的 rows [0, 64) | WG1 |
| 8–11 | 4 (128 threads) | 128 | Math WG1: 处理 tile 的 rows [64, 128) | WG2 |
WG 协作模式:
同一 tile (BLOCK_M=128):
WG1: owns rows [0, 64) → m64 WGMMA on smem_a[0:64]
WG2: owns rows [64, 128) → m64 WGMMA on smem_a[64:128]
B/weight: 共享一次 TMA load → 两个 WG 各自读取同一 smem_b
两个 WG 各自执行自己的 epilogue,通过 256-thread barrier (kEpilogueFullBarrierIdx) 在 L2 phase 同步,防止一个 WG 的下一 tile 写入覆盖另一个 WG 仍在 scatter 的本 tile 数据。
寄存器分配 (与 pingpong 完全一致):
| Role | Reg/thread | Threads | 总 Reg |
|---|---|---|---|
| Dispatch | 48 | 64 | 3,072 |
| TMA (A+B loader) | 48 | 64 | 3,072 |
| Math WG0 | 224 | 128 | 28,672 |
| Math WG1 | 224 | 128 | 28,672 |
| 总计 | 384 | 63,488/64K |
关键设计:
- WG 通过
row_block_offset = epilogue_wg_idx * WG_BLOCK_M偏移 smem_a/smem_sfa/smem_cd 的读写位置 - L1 epilogue: 两个 WG 各自 TMA store 自己 64-row 带,使用
block_m / 2 = 64的 descriptor box - L2 epilogue: 跨 WG 256-thread barrier 防止 smem_cd 别名冲突
- Cooperative 适合大 M:B/weight 仅 load 一次,weight HBM 读数减半
Token-Count Auto Routing
fp8_mega_moe 按 token 数自动选择 kernel:
| num_tokens | 选择 kernel | 适用场景 |
|---:---|---|
| < DG_SM90_MOE_COOPERATIVE_THRESHOLD (默认 256) | pingpong | 小/中等 M,epilogue 延迟主导 |
| ≥ threshold | cooperative | 大 M,weight 重复读取主导 |
也可通过 DG_SM90_MOE_KERNEL=pingpong|cooperative 强制选择。
代码review发现
严重: setmaxnreg 在 warpgroup 分歧分支中执行 [代码形态与 SM100 同构,风险降低]
setmaxnreg 在 warpgroup 分歧分支中执行初步审阅时怀疑 SM90 kernel 将 WG0 的 4 个 warp 拆分在不同 if/else if 分支中调用 warpgroup_reg_dealloc,可能违反 PTX 对 warpgroup 一致性的要求。经与主仓库 SM100 fused kernel(sm100_fp8_fp4_mega_moe.cuh)交叉验证,确认该代码形态与生产代码同构,实际风险低于最初判断;但不应将其概括为“PC 地址无关”。
关键事实:
| SM100 fused(主仓库,生产验证) | SM90 PR360 | |
|---|---|---|
| WG 划分 | TMA warps (4个) 独占 WG1 | dispatch (2) + TMA (2) 共享 WG0 |
| dealloc 调用方式 | 4 个 TMA warp 分别从 4 个不同的 else if 分支调用 dealloc<40> |
4 个 warp 分别从不同的 if/else if 分支调用 dealloc<48> |
| 同一 WG 内目标指令与 operand 是否一致? | 是,均为 setmaxnreg.dec.sync.aligned.u32 40 |
是,均为 setmaxnreg.dec.sync.aligned.u32 48 |
PTX 规范要求 warpgroup 内所有 warp 在收敛的 warpgroup 上执行相同的 setmaxnreg 指令(相同 opcode + operand,且满足 .sync.aligned 的同步语义)。这里能降低风险的依据是:分支条件按 warp role 决定,所有相关分支最终发出相同 operand 的 setmaxnreg.dec.sync.aligned,并且 SM100 生产路径使用同构模式。严格来说,是否满足“同一动态 warpgroup 指令”的约束仍应以生成 SASS/实机测试为准,不能把“源码中不同分支但同 operand”泛化成任意 PC 都安全。当前 8×H200 benchmark 已覆盖该路径的运行稳定性,但未替代专门的 SASS 审核或 sanitizer 验证。
高: 波次大小启发式可能违反调度器的整除约束
SM90 调度器在编译期要求 kNumExpertsPerWave 精确整除 kNumExpertsPerRank。复用的 host 启发式按尾部比例搜索最佳值,但不限制候选值为除数。
证据:
deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:37包含DG_STATIC_ASSERT(kNumExpertsPerRank % kNumExpertsPerWave == 0, "Invalid wave config")。csrc/jit_kernels/heuristics/sm90_mega_moe.hpp:116-124将 SM90 波次选择委托给get_num_experts_per_wave_for_mega_moe。csrc/jit_kernels/heuristics/mega_moe.hpp:129-144在[min_num_experts_per_wave, min_num_experts_per_wave * 2]内按尾部比例搜索并返回best_num_experts_per_wave,不筛选除数。csrc/jit_kernels/impls/sm90_fp8_mega_moe_pingpong.hpp:101-105和csrc/jit_kernels/impls/sm90_fp8_mega_moe_cooperative.hpp:101-105将启发式结果传入 kernel 模板。
影响:
- 合法的 public API shape 可能因
Invalid wave config在 JIT 编译时失败。 - 现有测试多使用 2 的幂次 per-rank expert 数,可能掩盖问题。API 本身仅检查
num_experts == num_experts_per_rank * num_ranks,不限制 per-rank expert 为 2 的幂或其他与启发式兼容的值。
建议修复:
- 约束启发式只返回
num_experts_per_rank的除数,或放松调度器不变式以支持尾部非完整波次。 - 增加非 2 的幂
num_experts_per_rank的针对性测试。
中: 调度器在波次边界读取越界 expert count
advance_expert_idx() 递增当前 expert 并立即调用 get_num_tokens(current_local_expert_idx)。当 consume 完一个波次的最后一个 expert 后,会调用 get_num_tokens(wave_end)。如果 wave_end == kNumExpertsPerRank 且 per-rank expert 数是 32 的倍数,没有 lane 拥有该 expert index,ptx::exchange 的源值为未初始化。
证据:
deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:73-81声明uint32_t valid_value且未初始化,仅在 lane 的缓存 expert index 匹配expert_idx时才赋值。deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:94-98递增current_local_expert_idx并无条件调用get_num_tokens。deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:126-137和141-163在完成一个 expert 后调用advance_expert_idx(),包括波次最后一个 expert。deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:61-63为每个(i * 32 + lane_idx)expert 缓存一个值。当kNumExpertsPerRank == 32时kNumExpertsPerLane == 1,expert index 32 无人拥有。
影响:
- 在 32 个 local expert 等常见 shape 下,设备调度器使用未初始化寄存器值。
- 该值通常在波次切换后被覆盖,不一定会造成可见的错误输出,但仍是未定义行为,在 sanitizer 或编译器变化时可能暴露。
建议修复:
get_num_tokens中初始化valid_value = 0。advance_expert_idx中仅当current_local_expert_idx < kNumExpertsPerRank时调用get_num_tokens,否则将current_num_tokens设为 0。
中: 中间 FP8 量化 scale 语义在 kernel、测试、baseline 间不一致
fused kernel 对 L1/SwiGLU 输出使用 2 的幂次 scale 进行量化,但 correctness reference 和 benchmark baseline 使用精确的 amax / 448.0 float scale。导致 correctness 检查与性能 baseline 比较的是不同的数值管线。
证据:
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh:1125-1135通过math::get_e4m3_sf_and_sf_inv计算 L1 输出 scale。deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh:1115-1125同样如此。deep_gemm/include/deep_gemm/common/math.cuh:98-106将get_e4m3_sf_and_sf_inv实现为2 ** ceil(log2(amax / 448)),并静态断言 UE8M0 风格行为。tests/test_mega_moe_sm90.py:231-238在 reference L1 quant/dequant 路径中使用精确的sf2 = amax / 448.0。tests/bench_mega_moe_sm90.py:323-325在 V1 contiguous baseline 中使用use_ue8m0_scale=False。tests/bench_mega_moe_sm90.py:592-599在 V1 low-latency baseline 中使用use_ue8m0_scale=False。
影响:
calc_diff测量的并非 fused kernel 与数值等价 reference 的差异。- 如果非融合 baseline 使用不同的 L1 量化 scale 策略,benchmark 加速比不是严格同类对比。
- 较宽松的
diff_tol=0.01可能掩盖此问题,但会削弱测试信号并隐藏 scale 敏感的回归。
建议修复:
- 明确 SM90 中间 activation scale 应使用精确 float 还是 2 的幂次 float。
- 若采用精确 float,将 L1 epilogue 中的
get_e4m3_sf_and_sf_inv替换为精确的amax / 448.0及其倒数。 - 若采用 2 的幂次,更新 PyTorch reference 和所有 benchmark baseline 使用相同的幂次 scale,并文档说明 float tensor 中存储 UE8M0 等效值。
中: Phase-profiling benchmark 路径与 API 不兼容且 kernel 未实现
benchmark 中 DG_SM90_MOE_PHASE_PROFILE 路径在 per-expert stats 后追加 64 个 int,但所有 public SM90 wrapper 要求 stats tensor 长度精确等于 num_experts_per_rank。kernel 中也仅写入 per-expert recv count,不包含 phase metric 写入。
证据:
tests/bench_mega_moe_sm90.py:712-714在DG_SM90_MOE_PHASE_PROFILE被设置时分配num_experts_per_rank + 64个 int。tests/bench_mega_moe_sm90.py:810-828从追加的 64-int 区域读取 phase metrics。csrc/apis/mega.hpp:321-324、450-453、555-558断言cumulative_local_expert_recv_stats->numel() == num_experts_per_rank。deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh:602-603和deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh:610-611仅累加 per-expert recv count。- 在整个 PR 中搜索,未发现
DG_SM90_MOE_PHASE_PROFILE或追加 phase slot 的 kernel 侧写入。
影响:
- 启用文档中描述/环境变量驱动的 phase profile 路径会在 host API 断言处失败。
- 即使放宽断言,打印的 phase metrics 也会是零或过期数据,因为 kernel 从未写入这些值。
建议修复:
- 删除 benchmark phase-profile 路径,或实现独立的 stats tensor/API 合约及 kernel 侧 instrumentation。
中: SM90 weight scale layout 检查允许 kernel 无法正确索引的 layout
API 接受两种 weight scale 布局: 连续的 (E, N/128, K/128) 或最后两维转置后连续。SM90 kernel 使用原始指针算术,假定 K 是连续的最内层维度。
证据:
csrc/apis/mega.hpp:315-318、444-447、549-552使用check_sf_layout(..., sm90_sfb_check=true, torch::kFloat)检查 weight scale。csrc/utils/layout.hpp:109-114同时接受(stride(-1) == 1 and stride(-2) == size(-1))和(stride(-1) == size(-2) and stride(-2) == 1)。deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh:858-867、913-918、1006-1010将 weight scale 视为以 K-block 为最快维度的连续原始数组进行索引。deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh:847-857、903-908、996-1000使用相同索引方式。deep_gemm/mega/__init__.py:135-158未对 SM90 scale tensor 做任何变换,测试仅覆盖自然连续布局。
影响:
- 转置 stride 的 scale tensor 可以合法通过校验,但在 kernel 中被按错误线性顺序解释。
- 将产生错误的 dequantization scale,且不会触发 host 断言失败。
建议修复:
- 对 SM90 原始指针 weight scale,要求严格自然连续性: group stride 为
size(-2) * size(-1),stride(-2) == size(-1),stride(-1) == 1。 - 或者传入并在 kernel 中使用显式 stride,但对此路径可能得不偿失。
低: Benchmark HBM expert 计数在单 rank 或无过滤场景下可能少计
benchmark 将 unique expert 数减 1 以剔除插入的 -1 哨兵值。如果 tensor 中不存在 -1,则会少计一个 touched expert。
证据:
tests/bench_mega_moe_sm90.py:784-792将非本地 expert 过滤为-1,然后计算torch.unique(...).numel() - 1。- 在单 rank 运行或无 route 被过滤/屏蔽的场景下,
-1可能不存在。
影响:
- 这些场景下报告的 HBM GB/s 偏低。
- 不影响 kernel 正确性。
建议修复:
- 改为统计
torch.unique(gathered_topk_idx[gathered_topk_idx >= 0]).numel()。
正面评价
- PR 将 SM90 代码与现有 SM100 MegaMoE 路径分离,降低了跨架构回归风险。
- Public API 同时提供自动路由和强制
pingpong/cooperative变体,便于 A/B 测试。 - 测试按分层结构组织,覆盖 smoke、启发式、shape、边界条件和随机压力场景。
- cooperative kernel 在共享
smem_cd重用处显式加入了跨 warpgroup barrier,正确处理了一类行带别名的数据竞争风险。
建议合入前检查清单
- 使
num_experts_per_wave与调度器合约一致,并增加非 2 的幂 expert 数覆盖。 - 修复调度器波次边界越界 token count 读取。
- 统一 L1 中间 FP8 scale 语义(kernel / 测试 / baseline)。
- 删除或实现
DG_SM90_MOE_PHASE_PROFILE。 - 收紧 SM90 weight scale layout 检查。
- 若 PR 描述继续引用 DeepEP V2 加速比,需在包含
deep_ep.ElasticBuffer的 DeepEP 版本上重跑--baseline-version v2/both。 - 在 Fabric Manager / NVSwitch / NCCL NVLS multicast 正常的环境上复测一次,确认非 fallback symmetric-memory 路径下的性能变化。
- 修复后,在 SM90/H100 或 H200 上运行
tests/test_mega_moe_sm90.py的 layers 1-5,分别测试DG_SM90_MOE_KERNEL=pingpong和DG_SM90_MOE_KERNEL=cooperative。