feat: add sm90 megamoe phase1 interfaces

MEGAMOE_SM90_DESIGN.md

@@ -808,24 +808,48 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
 
 ---
 
-### Phase 2: Dispatch + 单 tile L1 GEMM (Day 11-20)
+### Phase 2: Dispatch + L1 Pool 填充 (Day 11-15)
 
-**目标**：实现 dispatch + L1 WGMMA 的正确 TMA pipeline，**不含 epilogue**（accumulator 直接丢弃）。验证 dispatch pull、arrival count、TMA A/B 加载、WGMMA FP8×FP8 累积 的端到端正确性。
+**目标**：只实现 dispatch 数据路径，不进入 TMA/WGMMA。验证 expert count、token routing、source metadata、跨 rank pull、`l1_arrival_count` 等 dispatch 产物正确，为后续 L1 GEMM 提供稳定输入。
 
 | 步骤 | 交付物 | 验证标准 |
 |------|--------|----------|
-| 2.1 Kernel 骨架 | `sm90_fp8_mega_moe.cuh` 中的 warp role 分配：WG0 dispatch+TMA, WG1/WG2 math | 编译通过，launch_bounds(384,1) |
-| 2.2 Barrier 初始化 | 所有 mbarrier (full/empty/dispatch/combine) 初始化 + `__syncthreads` | 无 launch failure |
-| 2.3 Dispatch phase | 64-thread dispatch: expert count、source metadata、grid sync、NVLink barrier、token pull | 单 rank 模式下 `l1_token_buffer` 中数据正确 |
-| 2.4 TMA A/SFA load | w2 TMA A warp：wait `l1_arrival_count`，TMA load L1 acts + float SFA to SMEM | dump SMEM A 区域，与 global 数据逐字节对比 |
-| 2.5 TMA B load | w3 TMA B warp：TMA load L1 FP8 weights to SMEM | dump SMEM B 区域对比 |
-| 2.6 WGMMA pipeline | WG1 math warp：`warpgroup_arrive`, 4× K-step `WGMMA::wgmma`, `warpgroup_commit_batch`, `warpgroup_wait<0>` | 单 tile (1 expert, 1 block) 的 accumulator 与 PyTorch FP32 GEMM 结果对比（未 scale） |
-| 2.7 SF scaling | accum *= sfa * sfb (sfb via `__ldg`) | Scale 后的 accumulator 与 reference scaled GEMM 一致 |
-| 2.8 Empty barrier | WGMMA 完成后 arrive empty_barrier → TMA 可重用 SMEM | 多 K-block 的 pipeline 不死锁 |
+| 2.1 Dispatch-only kernel 骨架 | `sm90_fp8_mega_moe.cuh` 中 WG0 的 64-thread dispatch 角色分配；TMA/math WG 可 idle | 编译通过，launch_bounds(384,1)，launch 不崩溃 |
+| 2.2 Dispatch barrier / workspace 初始化 | dispatch mbarrier、grid sync counter、expert count、arrival counter 初始化 | 连续两次调用无残留 workspace 污染 |
+| 2.3 Expert count | 统计本 rank token/top-k 到 local experts 的 recv count，并写入 workspace | expert token counts 与 PyTorch routing reference 一致 |
+| 2.4 Source metadata | 写入 `token_src_metadata[pool_token_idx] = {rank, token_idx, topk_idx}` | metadata 与 routing reference 逐项一致 |
+| 2.5 Token/SF/top-k pull | 将输入 token、float SFA、top-k weight 拉入 local L1 pool | 单 rank 下 `l1_acts` / `l1_acts_sf` / `l1_topk_weights` 与输入逐字节或逐值一致 |
+| 2.6 Cross-rank dispatch smoke | NVLink barrier + remote pull 路径打通 | 2-rank / 8-rank 下无 hang，pool 内容与 reference 一致 |
+| 2.7 `l1_arrival_count` release | 每个 pool block dispatch 完成后 release-add arrival count | arrival count 达到 block 内 token 数，TMA consumer 可安全等待 |
+
+**验证里程碑**：单 rank 和多 rank 下 dispatch-only 输出 `l1_acts`、`l1_acts_sf`、`l1_topk_weights`、`token_src_metadata`、`l1_arrival_count` 与 PyTorch routing reference 一致；不要求任何 GEMM 计算。
+
+**依赖**：Phase 1 完成
+**关键挑战**：
+- pool token index / expert block offset 计算必须与 scheduler 后续消费保持一致
+- symmetric memory `sym_buffer.map` 与 NVLink barrier 的 phase 计数必须可连续调用
+- debug dump 只能临时使用，必须按 AGENTS.md 加 `// DEBUG` / `# DEBUG` 并 stash
+
+---
+
+### Phase 3: L1 TMA + 单 tile WGMMA (Day 16-20)
+
+**目标**：在 Phase 2 dispatch 产物上实现 L1 TMA/WGMMA pipeline，**不含 epilogue**（accumulator 可直接丢弃或写 debug buffer）。验证 TMA A/SFA/B 加载、WGMMA FP8×FP8 累积、SF scaling、empty barrier 的正确性。
+
+| 步骤 | 交付物 | 验证标准 |
+|------|--------|----------|
+| 3.1 TMA/math kernel 骨架 | WG0 w2/w3 作为 TMA A/B warp，WG1/WG2 作为 math warpgroup | 编译通过，launch_bounds(384,1) |
+| 3.2 Full/empty barrier 初始化 | full/empty mbarrier per pipeline stage 初始化 + `__syncthreads` | 无 launch failure / deadlock |
+| 3.3 TMA A/SFA load | w2 TMA A warp：wait `l1_arrival_count`，TMA load L1 acts + float SFA to SMEM | dump SMEM A/SFA，与 global pool 数据逐字节/逐值对比 |
+| 3.4 TMA B load | w3 TMA B warp：TMA load L1 FP8 weights to SMEM | dump SMEM B 区域对比 |
+| 3.5 WGMMA descriptor | 构造 A/B GMMA descriptor；cooperative 下 WG1/WG2 指向不同 M offset | descriptor 地址和 swizzle 与 SM90 GEMM reference 一致 |
+| 3.6 WGMMA pipeline | math WG：`warpgroup_arrive`, 4× K-step `WGMMA::wgmma`, `warpgroup_commit_batch`, `warpgroup_wait<0>` | 单 tile accumulator 与 PyTorch FP32 GEMM 结果对比（未 scale） |
+| 3.7 SF scaling | accum *= sfa * sfb (sfb via `__ldg`) | Scale 后 accumulator 与 reference scaled GEMM 一致 |
+| 3.8 Empty barrier | WGMMA 完成后 arrive empty_barrier → TMA 可重用 SMEM | 多 K-block pipeline 不死锁 |
 
 **验证里程碑**：单 expert、BLOCK_M=128 的 L1 GEMM（无 epilogue）输出与 PyTorch `(dequant(x) @ dequant(w).T) * sfa * sfb` 在 FP32 accumulator 精度内 `allclose(atol=1e-2, rtol=0.05)`。
 
-**依赖**：Phase 1 完成
+**依赖**：Phase 2 完成
 **关键挑战**：
 - WGMMA descriptor 构造（参考 `deep_gemm/mma/sm90.cuh` 中的 `make_gmma_desc`）
 - Cooperative mode 下两个 WG 的 desc 指向 smem_a 的不同 64-row 偏移
@@ -833,27 +857,27 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
 
 ---
 
-### Phase 3: L1 Epilogue — SwiGLU + FP8 Quant (Day 21-30)
+### Phase 4: L1 Epilogue — SwiGLU + FP8 Quant (Day 21-30)
 
-**目标**：在 Phase 2 的 WGMMA 结果上完成 L1 epilogue 全流程，输出正确的 FP8 `l2_acts` 和 float `l2_acts_sf`。
+**目标**：在 Phase 3 的 WGMMA 结果上完成 L1 epilogue 全流程，输出正确的 FP8 `l2_acts` 和 float `l2_acts_sf`。
 
 | 步骤 | 交付物 | 验证标准 |
 |------|--------|----------|
-| 3.1 Top-k weight load | 从 `l1_topk_weights_buffer` load weight per row | 值与 dispatch 写入一致 |
-| 3.2 SwiGLU 计算 | gate/up deinterleave + `silu(gate) * up * topk_weight` in registers | 与 reference SwiGLU 一致（精度：atol=0.02） |
-| 3.3 Activation clamp | `kActivationClamp` 模板参数控制 clamp | clamp 后无超范围值 |
-| 3.4 fast_math SwiGLU | `__expf` + `fast_rcp` 快速路径 | 与精确版相差 < 0.5% |
-| 3.5 Per-row amax | `warp_reduce<4>(max(abs(...)))` 跨 WG 内 4 warps reduce | amax 值与 PyTorch `abs().max(dim=-1)` 一致 |
-| 3.6 Cooperative amax | 两个 WG 独立 amax（各自 64 行），SMEM scratch 区域用于跨 warp reduce | 每行 amax 正确 |
-| 3.7 FP8 quantize | `sf = amax / 448.0`，`fp8 = cast(result / sf)` | 反量化后与 reference 误差 < 1% |
-| 3.8 Float SF write | 写入 `l2_acts_sf` (MN-major, per-64-K layout) | SF 值正确、layout 与 L2 TMA load 兼容 |
-| 3.9 STSM + double-buf TMA store | STSM 到 `smem_cd[stage]`，TMA store 到 `l2_acts` | `l2_acts` 内容与 reference FP8 quant output 一致 |
-| 3.10 `l2_arrival_mask` 设置 | `red_or_rel_gpu(mask, 1 << n_block_idx)` | 单 expert 全部 L1 N-blocks 完成后 mask == expected |
-| 3.11 Cooperative sync | `sync_aligned(256)` 等两个 WG 都 TMA store 完成后设置 mask | 无 race condition |
+| 4.1 Top-k weight load | 从 `l1_topk_weights_buffer` load weight per row | 值与 dispatch 写入一致 |
+| 4.2 SwiGLU 计算 | gate/up deinterleave + `silu(gate) * up * topk_weight` in registers | 与 reference SwiGLU 一致（精度：atol=0.02） |
+| 4.3 Activation clamp | `kActivationClamp` 模板参数控制 clamp | clamp 后无超范围值 |
+| 4.4 fast_math SwiGLU | `__expf` + `fast_rcp` 快速路径 | 与精确版相差 < 0.5% |
+| 4.5 Per-row amax | `warp_reduce<4>(max(abs(...)))` 跨 WG 内 4 warps reduce | amax 值与 PyTorch `abs().max(dim=-1)` 一致 |
+| 4.6 Cooperative amax | 两个 WG 独立 amax（各自 64 行），SMEM scratch 区域用于跨 warp reduce | 每行 amax 正确 |
+| 4.7 FP8 quantize | `sf = amax / 448.0`，`fp8 = cast(result / sf)` | 反量化后与 reference 误差 < 1% |
+| 4.8 Float SF write | 写入 `l2_acts_sf` (MN-major, per-64-K layout) | SF 值正确、layout 与 L2 TMA load 兼容 |
+| 4.9 STSM + double-buf TMA store | STSM 到 `smem_cd[stage]`，TMA store 到 `l2_acts` | `l2_acts` 内容与 reference FP8 quant output 一致 |
+| 4.10 `l2_arrival_mask` 设置 | `red_or_rel_gpu(mask, 1 << n_block_idx)` | 单 expert 全部 L1 N-blocks 完成后 mask == expected |
+| 4.11 Cooperative sync | `sync_aligned(256)` 等两个 WG 都 TMA store 完成后设置 mask | 无 race condition |
 
 **验证里程碑**：单 expert 的 `l2_acts` (FP8) 和 `l2_acts_sf` (float) 与 reference `SwiGLU → FP8_quant` 一致。Multi-expert 场景下 `l2_arrival_mask` 正确触发。
 
-**依赖**：Phase 2 完成
+**依赖**：Phase 3 完成
 **关键挑战**：
 - WGMMA register layout 到 gate/up 列的映射关系（参考 SM100 的 `SM100_TMEM_LOAD_16dp256b1x` 读取模式，SM90 的 register layout 不同）
 - Double-buffered TMA store 的 `tma_store_wait<1>` 与 smem_cd 复用
@@ -861,26 +885,26 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
 
 ---
 
-### Phase 4: L2 GEMM + NVLink Scatter (Day 31-40)
+### Phase 5: L2 GEMM + NVLink Scatter (Day 31-40)
 
 **目标**：L2 WGMMA pipeline + per-64 dual-half SF + BF16 epilogue + NVLink scatter 到 remote combine buffer。
 
 | 步骤 | 交付物 | 验证标准 |
 |------|--------|----------|
-| 4.1 L2 TMA wait | TMA A warp wait `l2_arrival_mask == expected` | 正确等待 L1 所有 N-blocks 完成 |
-| 4.2 L2 TMA load | TMA load `l2_acts` + 两个 per-64 SFA halves | SMEM 数据正确 |
-| 4.3 L2 WGMMA dual-half | K-step 0-1 累积 half0，K-step 2-3 累积 half1 | 两组 partial accum 正确 |
-| 4.4 Dual-half SF apply | `final = half0 * sfa0 * sfb + half1 * sfa1 * sfb` | 与 reference L2 dequant GEMM 一致 |
-| 4.5 Weight SF `__ldg` | Math WG prefetch L2 weight SF from global | 值正确 |
-| 4.6 BF16 cast + STSM | FP32 → BF16 cast，STSM 到 `smem_cd_l2` (BF16 swizzle) | SMEM 内容正确 |
-| 4.7 Source metadata read | `token_src_metadata[pool_token_idx + row]` → {rank, token, topk} | 值与 dispatch 写入一致 |
-| 4.8 NVLink scatter | `ld_shared` + `sym_buffer.map` + remote store | 单 rank 下本地 combine buffer 数据正确 |
-| 4.9 Cooperative smem_cd guard | `sync_aligned(256)` 防止下一 tile 覆盖 smem_cd | 多 tile 序列无 data corruption |
-| 4.10 Multi-expert persistent loop | Scheduler `for_each_block` 遍历多 expert 的 L1→L2 cycle | 全量 expert 输出正确 |
+| 5.1 L2 TMA wait | TMA A warp wait `l2_arrival_mask == expected` | 正确等待 L1 所有 N-blocks 完成 |
+| 5.2 L2 TMA load | TMA load `l2_acts` + 两个 per-64 SFA halves | SMEM 数据正确 |
+| 5.3 L2 WGMMA dual-half | K-step 0-1 累积 half0，K-step 2-3 累积 half1 | 两组 partial accum 正确 |
+| 5.4 Dual-half SF apply | `final = half0 * sfa0 * sfb + half1 * sfa1 * sfb` | 与 reference L2 dequant GEMM 一致 |
+| 5.5 Weight SF `__ldg` | Math WG prefetch L2 weight SF from global | 值正确 |
+| 5.6 BF16 cast + STSM | FP32 → BF16 cast，STSM 到 `smem_cd_l2` (BF16 swizzle) | SMEM 内容正确 |
+| 5.7 Source metadata read | `token_src_metadata[pool_token_idx + row]` → {rank, token, topk} | 值与 dispatch 写入一致 |
+| 5.8 NVLink scatter | `ld_shared` + `sym_buffer.map` + remote store | 单 rank 下本地 combine buffer 数据正确 |
+| 5.9 Cooperative smem_cd guard | `sync_aligned(256)` 防止下一 tile 覆盖 smem_cd | 多 tile 序列无 data corruption |
+| 5.10 Multi-expert persistent loop | Scheduler `for_each_block` 遍历多 expert 的 L1→L2 cycle | 全量 expert 输出正确 |
 
 **验证里程碑**：多 expert、单 rank 下 `combine_token_buffer` 内容与 reference L2 GEMM BF16 输出一致。`l2_arrival_mask` 正确控制 L1→L2 依赖。
 
-**依赖**：Phase 3 完成
+**依赖**：Phase 4 完成
 **关键挑战**：
 - Per-64 dual-half 需要在 K-loop 中区分 half0/half1（k_step 0-1 vs 2-3），需要两组 accumulator 或 intermediate scale → **224 reg budget 最紧张的地方**
 - NVLink scatter 的 `sym_buffer.map` 在 SM90 上的行为验证（与 SM100 一致）
@@ -888,24 +912,24 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
 
 ---
 
-### Phase 5: Combine Reduce + 端到端 Fused (Day 41-48)
+### Phase 6: Combine Reduce + 端到端 Fused (Day 41-48)
 
 **目标**：完成 combine reduce、NVLink barrier、workspace 清理。实现完整端到端 fused kernel。
 
 | 步骤 | 交付物 | 验证标准 |
 |------|--------|----------|
-| 5.1 NVLink barrier | Epilogue grid sync + cross-rank signal | 多 rank 下所有 rank 到达 barrier |
-| 5.2 Dispatch↔Epilogue handshake | `kDispatchWithEpilogueBarrierIdx` sync | Dispatch 和 epilogue 阶段正确交替 |
-| 5.3 Combine TMA load | Double-buffered TMA 从 combine buffer load BF16 chunks | Load 数据与 scatter 写入一致 |
-| 5.4 FP32 accumulate | 逐 top-k slot 累加到 float registers | 累加结果与 reference `sum(topk_weights * expert_outputs)` 一致 |
-| 5.5 BF16 cast + TMA store | Cast + TMA store 到 output `y` | `y` 与 reference 在 BF16 精度内一致 |
-| 5.6 Workspace cleanup | Dispatch warps 清理 expert counts、arrival counters、metadata | 下次 kernel 调用前 workspace 已归零 |
-| 5.7 NVLink cleanup barrier | 跨 rank 确认 workspace 清理完成 | 连续两次 kernel 调用无 data corruption |
-| 5.8 端到端 correctness | 完整 `fp8_mega_moe(...)` 调用 | `test_mega_moe_sm90.py` Smoke test 通过 (2 ranks, 4096×2048, topk=6) |
+| 6.1 NVLink barrier | Epilogue grid sync + cross-rank signal | 多 rank 下所有 rank 到达 barrier |
+| 6.2 Dispatch↔Epilogue handshake | `kDispatchWithEpilogueBarrierIdx` sync | Dispatch 和 epilogue 阶段正确交替 |
+| 6.3 Combine TMA load | Double-buffered TMA 从 combine buffer load BF16 chunks | Load 数据与 scatter 写入一致 |
+| 6.4 FP32 accumulate | 逐 top-k slot 累加到 float registers | 累加结果与 reference `sum(topk_weights * expert_outputs)` 一致 |
+| 6.5 BF16 cast + TMA store | Cast + TMA store 到 output `y` | `y` 与 reference 在 BF16 精度内一致 |
+| 6.6 Workspace cleanup | Dispatch warps 清理 expert counts、arrival counters、metadata | 下次 kernel 调用前 workspace 已归零 |
+| 6.7 NVLink cleanup barrier | 跨 rank 确认 workspace 清理完成 | 连续两次 kernel 调用无 data corruption |
+| 6.8 端到端 correctness | 完整 `fp8_mega_moe(...)` 调用 | `test_mega_moe_sm90.py` Smoke test 通过 (2 ranks, 4096×2048, topk=6) |
 
 **验证里程碑**：2-rank Smoke test (4 tokens, 8 experts) `allclose(atol=0.05, rtol=0.1)` 通过。
 
-**依赖**：Phase 4 完成
+**依赖**：Phase 5 完成
 **关键挑战**：
 - Combine 逻辑直接移植自 SM100 (与 TMEM/UMMA 无关)，但需要确认 SM90 下 TMA 1D load/store 的行为
 - NVLink barrier 的 phase 计数在多次 kernel 调用间的正确性
@@ -913,65 +937,65 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
 
 ---
 
-### Phase 6: Single-WG 变体 (BLOCK_M=64/32) (Day 49-55)
+### Phase 7: Single-WG 变体 (BLOCK_M=64/32) (Day 49-55)
 
 **目标**：实现 `kCooperativeMode=false` 路径，支持小 BLOCK_M 场景。
 
 | 步骤 | 交付物 | 验证标准 |
 |------|--------|----------|
-| 6.1 模板分支 | `if constexpr (kCooperativeMode)` 控制 WG2 行为 | 两种模式均编译通过 |
-| 6.2 WG2 idle | WG2 在 `kCooperativeMode=false` 时 `warpgroup_reg_dealloc` 后 idle | 无 deadlock / launch failure |
-| 6.3 TMA A 调整 | TMA A 只加载 BLOCK_M 行（64 或 32）而非 128 | SMEM A 数据正确 |
-| 6.4 WGMMA 有效行 | `valid_m = min(tokens_in_block, BLOCK_M)` 控制 epilogue 写入范围 | 无 OOB 写入 |
-| 6.5 Scheduler 适配 | 动态 BLOCK_M 下 scheduler 的 block 分配和 pool_block_offset 正确 | 全量 expert 覆盖，无漏 tile |
-| 6.6 Heuristics JIT 分发 | 根据 `expected_tokens_per_expert` 选择 BLOCK_M，JIT 编译对应模板 | 正确选择 config 并编译 |
-| 6.7 Small M correctness | BLOCK_M=64: 单 WG 处理完整 64 行 | Smoke test 通过 |
-| 6.8 Extreme decode | BLOCK_M=32: token-per-expert=1-2 | 极端 decode 场景输出正确 |
+| 7.1 模板分支 | `if constexpr (kCooperativeMode)` 控制 WG2 行为 | 两种模式均编译通过 |
+| 7.2 WG2 idle | WG2 在 `kCooperativeMode=false` 时 `warpgroup_reg_dealloc` 后 idle | 无 deadlock / launch failure |
+| 7.3 TMA A 调整 | TMA A 只加载 BLOCK_M 行（64 或 32）而非 128 | SMEM A 数据正确 |
+| 7.4 WGMMA 有效行 | `valid_m = min(tokens_in_block, BLOCK_M)` 控制 epilogue 写入范围 | 无 OOB 写入 |
+| 7.5 Scheduler 适配 | 动态 BLOCK_M 下 scheduler 的 block 分配和 pool_block_offset 正确 | 全量 expert 覆盖，无漏 tile |
+| 7.6 Heuristics JIT 分发 | 根据 `expected_tokens_per_expert` 选择 BLOCK_M，JIT 编译对应模板 | 正确选择 config 并编译 |
+| 7.7 Small M correctness | BLOCK_M=64: 单 WG 处理完整 64 行 | Smoke test 通过 |
+| 7.8 Extreme decode | BLOCK_M=32: token-per-expert=1-2 | 极端 decode 场景输出正确 |
 
 **验证里程碑**：BLOCK_M=64 和 BLOCK_M=32 的 Smoke test 通过。动态 BLOCK_M 选择逻辑根据 token count 正确切换。
 
-**依赖**：Phase 5 完成（cooperative 端到端已验证）
+**依赖**：Phase 6 完成（cooperative 端到端已验证）
 **关键挑战**：
 - BLOCK_M=32 时 WGMMA 计算 64 行但只用 32 行的 register 浪费
 - Single-WG mode 下 WG2 的 named barrier arrive 需要调整（WG2 不参与 epilogue barrier）
 
 ---
 
-### Phase 7: 多 Rank 集成测试 (Day 56-62)
+### Phase 8: 多 Rank 集成测试 (Day 56-62)
 
 **目标**：在 8×H200 集群上验证完整 multi-rank 正确性和性能。
 
 | 步骤 | 交付物 | 验证标准 |
 |------|--------|----------|
-| 7.1 Multi-rank launch | 8-rank torchrun 启动 MegaMoE kernel | 无 hang / crash |
-| 7.2 Correctness sweep | Token sweep (1-8192) × Shape sweep × BLOCK_M 三档 | 全部 `allclose` 通过 |
-| 7.3 Edge cases | 0 tokens, unbalanced routing, single-expert-per-token | 无 crash，输出正确 |
-| 7.4 Benchmark framework | `bench_mega_moe_sm90.py`：latency, TFLOPS, HBM GB/s | 可生成 performance table |
-| 7.5 nsys timeline | Phase profile: dispatch / L1 / L2 / combine 各阶段时间 | 识别主要 stall 来源 |
-| 7.6 DeepEP v1 baseline | 对比 DeepEP dispatch + grouped GEMM + combine | 输出 speedup ratio |
-| 7.7 Regression test | CI 集成到 `tests/test_mega_moe_sm90.py` | `pytest` 通过 |
+| 8.1 Multi-rank launch | 8-rank torchrun 启动 MegaMoE kernel | 无 hang / crash |
+| 8.2 Correctness sweep | Token sweep (1-8192) × Shape sweep × BLOCK_M 三档 | 全部 `allclose` 通过 |
+| 8.3 Edge cases | 0 tokens, unbalanced routing, single-expert-per-token | 无 crash，输出正确 |
+| 8.4 Benchmark framework | `bench_mega_moe_sm90.py`：latency, TFLOPS, HBM GB/s | 可生成 performance table |
+| 8.5 nsys timeline | Phase profile: dispatch / L1 / L2 / combine 各阶段时间 | 识别主要 stall 来源 |
+| 8.6 DeepEP v1 baseline | 对比 DeepEP dispatch + grouped GEMM + combine | 输出 speedup ratio |
+| 8.7 Regression test | CI 集成到 `tests/test_mega_moe_sm90.py` | `pytest` 通过 |
 
 **验证里程碑**：8-rank × DSV4 Flash shape (4096×2048, E256, topk6)：correctness 通过 + fused latency < 3000us (128 tokens)。
 
-**依赖**：Phase 6 完成
+**依赖**：Phase 7 完成
 
 ---
 
-### Phase 8: 性能优化 (Day 63-75)
+### Phase 9: 性能优化 (Day 63-75)
 
 **目标**：基于 nsys timeline 分析，针对瓶颈进行优化迭代。
 
 | 步骤 | 优化方向 | 预期收益 | 依据 |
 |------|---------|---------|------|
-| 8.1 N-major scheduling tuning | 调整 `tokens/expert >= 256` 阈值 | 大 M 下 weight L2 cache 命中率提升 | NCU L2 sector miss 数据 |
-| 8.2 SFA 软件 prefetch | L2 per-64 SFA 的 `__ldg` prefetch 提前一个 K-block | 减少 math stall waiting on SF | nsys timeline 中 SF load 占比 |
-| 8.3 Combine chunk 大小调优 | 根据 hidden dim 选择 1 vs 2 chunks | 减少 combine 阶段 TMA round-trip | combine 占总时间比例 |
-| 8.4 Dispatch pull 优化 | 增加 dispatch warp 内的 token parallelism | 减少 dispatch phase 时间 | nsys dispatch 耗时 |
-| 8.5 Register spill 消除 | 检查 SASS 中的 local memory access，调整变量生命周期 | 消除 spill 带来的 DRAM 流量 | cuobjdump --dump-sass |
-| 8.6 wave_count 调优 | 根据实测 per-shape 最优 experts_per_wave | 减少 tail effect | benchmark sweep |
-| 8.7 Optional: cluster=2 探索 | TMA multicast B tile（需解决 amax 跨 CTA sync） | weight HBM 进一步减半 | 若 cooperative 的 HBM 仍是瓶颈 |
+| 9.1 N-major scheduling tuning | 调整 `tokens/expert >= 256` 阈值 | 大 M 下 weight L2 cache 命中率提升 | NCU L2 sector miss 数据 |
+| 9.2 SFA 软件 prefetch | L2 per-64 SFA 的 `__ldg` prefetch 提前一个 K-block | 减少 math stall waiting on SF | nsys timeline 中 SF load 占比 |
+| 9.3 Combine chunk 大小调优 | 根据 hidden dim 选择 1 vs 2 chunks | 减少 combine 阶段 TMA round-trip | combine 占总时间比例 |
+| 9.4 Dispatch pull 优化 | 增加 dispatch warp 内的 token parallelism | 减少 dispatch phase 时间 | nsys dispatch 耗时 |
+| 9.5 Register spill 消除 | 检查 SASS 中的 local memory access，调整变量生命周期 | 消除 spill 带来的 DRAM 流量 | cuobjdump --dump-sass |
+| 9.6 wave_count 调优 | 根据实测 per-shape 最优 experts_per_wave | 减少 tail effect | benchmark sweep |
+| 9.7 Optional: cluster=2 探索 | TMA multicast B tile（需解决 amax 跨 CTA sync） | weight HBM 进一步减半 | 若 cooperative 的 HBM 仍是瓶颈 |
 
-**验证里程碑**：相对 Phase 7 baseline，主要 shape 上 latency 降低 10-20%。
+**验证里程碑**：相对 Phase 8 baseline，主要 shape 上 latency 降低 10-20%。
 
 ---
 
@@ -984,24 +1008,27 @@ Phase 0 (环境)
 Phase 1 (基础设施)
     │
     v
-Phase 2 (Dispatch + L1 GEMM)
+Phase 2 (Dispatch)
     │
     v
-Phase 3 (L1 Epilogue)
+Phase 3 (L1 TMA + WGMMA)
     │
     v
-Phase 4 (L2 GEMM + Scatter)
+Phase 4 (L1 Epilogue)
     │
     v
-Phase 5 (Combine + 端到端)
+Phase 5 (L2 GEMM + Scatter)
+    │
+    v
+Phase 6 (Combine + 端到端)
     │
     ├──────────────────┐
     v                  v
-Phase 6 (Single-WG)   Phase 7 (多 Rank 集成)
+Phase 7 (Single-WG)   Phase 8 (多 Rank 集成)
     │                  │
     └──────┬───────────┘
            v
-    Phase 8 (性能优化)
+    Phase 9 (性能优化)
 ```
 
 ### 时间线总结
@@ -1010,12 +1037,13 @@ Phase 6 (Single-WG)   Phase 7 (多 Rank 集成)
 |-------|------|------|---------|
 | 0 | 3 | 3 | 环境 + baseline |
 | 1 | 7 | 10 | 基础设施 (可 JIT 编译) |
-| 2 | 10 | 20 | Dispatch + L1 WGMMA 正确 |
-| 3 | 10 | 30 | L1 SwiGLU + FP8 quant 正确 |
-| 4 | 10 | 40 | L2 GEMM + NVLink scatter 正确 |
-| 5 | 8 | 48 | 端到端 fused kernel 正确 (2-rank) |
-| 6 | 7 | 55 | 动态 BLOCK_M 全变体正确 |
-| 7 | 7 | 62 | 8-rank 集成 + benchmark |
-| 8 | 13 | 75 | 性能调优完成 |
+| 2 | 5 | 15 | Dispatch + L1 pool 正确 |
+| 3 | 5 | 20 | L1 TMA + WGMMA 正确 |
+| 4 | 10 | 30 | L1 SwiGLU + FP8 quant 正确 |
+| 5 | 10 | 40 | L2 GEMM + NVLink scatter 正确 |
+| 6 | 8 | 48 | 端到端 fused kernel 正确 (2-rank) |
+| 7 | 7 | 55 | 动态 BLOCK_M 全变体正确 |
+| 8 | 7 | 62 | 8-rank 集成 + benchmark |
+| 9 | 13 | 75 | 性能调优完成 |
 
-**总工期**：约 75 工作日（10-11 周），含充分的验证和调试时间。关键路径是 Phase 2-5 的串行 kernel 开发（40 天）。
+**总工期**：约 75 工作日（10-11 周），含充分的验证和调试时间。关键路径是 Phase 2-6 的串行 kernel 开发。