# 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_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` 已恢复为 base `torch 2.9.1+cu130` 可 import 状态，Torch 2.12 构建产物另存为 `.torch212cu130.so`。
- 为适配 PyTorch 2.12 的 tensor metadata 限制，本地验证版本在 `csrc/apis/mega.hpp` 的 `slice_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-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 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 |
| 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) 形成一个完整的 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 |
| 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 同构，风险降低]

初步审阅时怀疑 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`。
