Files
DeepGEMM/megamoe-research-reports/pr360_sm90_pingpong_coop_megamoe_code_review.md
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

22 KiB
Raw Blame History

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 列表

  • f4f8473cd sm90 megamoe: add pingpong and cooperative FP8 kernels
  • 3f10f93a1 sm90 megamoe: token-count dispatch and Python bindings
  • 89e3d351e sm90 megamoe: correctness test and performance benchmark
  • 7b65ae9b4 chore: ignore profiling artifacts (*.ncu-rep, *.nsys-rep, profile_logs/)
  • ba270a419 fix mul2
  • 0e6549b60 add sm90 megamoe bench
  • f27fd5611 Revert "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_moefp8_mega_moe_pingpongfp8_mega_moe_cooperativetransform_weights_for_mega_moe_sm90

8×H200 实验验证补充

本节为后续补测结果,目的是验证 PR 描述中的性能表是否能在单节点 8×H200 上复现。报告更新本身未重新访问远端机器。

运行环境与准备:

  • 机器: g0033 单节点 8×H2008 ranks每 rank 1 GPU。
  • 代码路径: /mnt/beegfs/lxy/sm90-bench/pr-360
  • Python 环境: /mnt/beegfs/lxy/venvs/pr360-torch212-cu130torch 2.12.0+cu130tilelang 0.1.9
  • PR360 _C/sgl-workspace/DeepEP/deep_ep_cpp 均按 Torch 2.12/CUDA 13.0 临时重建;运行后默认 .so 已恢复为 base torch 2.9.1+cu130 可 import 状态Torch 2.12 构建产物另存为 .torch212cu130.so
  • 为适配 PyTorch 2.12 的 tensor metadata 限制,本地验证版本在 csrc/apis/mega.hppslice_input_buffers lambda 中加入 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-contigv1-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 Flash tokens=256 明显偏慢(2213.3 us vs 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
47 4 (128 threads) 128 Math WG0: m64 WGMMA → pingpong epilogue → TMA store L1 / BF16 scatter WG1
811 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) 形成一个完整的 WG0
  • kNumEpilogueWarpgroups == 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
47 4 (128 threads) 128 Math WG0: 处理 tile 的 rows [0, 64) WG1
811 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 适合大 MB/weight 仅 load 一次weight HBM 读数减半

Token-Count Auto Routing

fp8_mega_moe 按 token 数自动选择 kernel:

| num_tokens | 选择 kernel | 适用场景 | |---:---|---| | < DG_SM90_MOE_COOPERATIVE_THRESHOLD (默认 256) | pingpong | 小/中等 Mepilogue 延迟主导 | | ≥ threshold | cooperative | 大 Mweight 重复读取主导 |

也可通过 DG_SM90_MOE_KERNEL=pingpong|cooperative 强制选择。

代码review发现

严重: setmaxnreg 在 warpgroup 分歧分支中执行 [代码形态与 SM100 同构,风险降低]

初步审阅时怀疑 SM90 kernel 将 WG0 的 4 个 warp 拆分在不同 if/else if 分支中调用 warpgroup_reg_dealloc,可能违反 PTX 对 warpgroup 一致性的要求。经与主仓库 SM100 fused kernelsm100_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-105csrc/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 indexptx::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-137141-163 在完成一个 expert 后调用 advance_expert_idx(),包括波次最后一个 expert。
  • deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:61-63 为每个 (i * 32 + lane_idx) expert 缓存一个值。当 kNumExpertsPerRank == 32kNumExpertsPerLane == 1expert 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-106get_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-714DG_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-324450-453555-558 断言 cumulative_local_expert_recv_stats->numel() == num_experts_per_rank
  • deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh:602-603deep_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-318444-447549-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-867913-9181006-1010 将 weight scale 视为以 K-block 为最快维度的连续原始数组进行索引。
  • deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh:847-857903-908996-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=pingpongDG_SM90_MOE_KERNEL=cooperative