# PR323 SM90 Fused MegaMoE 代码review报告

## 范围

- Worktree: `pr-323`
- HEAD: `23f46aa68`
- 审查方式: 代码review
- 主题: 首次引入 SM90/Hopper FP8 MegaMoE 单 kernel 实现

## 实现概述

### SM90 FP8 MegaMoE Fused Kernel

`deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` (~1935 lines)

**warp 分工与线程布局（3 种配置）:**

**配置 1: 512-epilogue (dispatch+TMA 各自独占 warpgroup)**

| Warp Index | 数量 | 所属 Warpgroup | Role |
|---:|---:|---|---|
| 0–3 | 4 warps (128 threads) | WG0 | Dispatch |
| 4–7 | 4 warps (128 threads) | WG1 | TMA A+SFA / B+SFB / MMA issue / idle |
| 8–23 | 16 warps (512 threads) | WG2–WG5 | Math WGMMA + epilogue + combine |

**精确 warp 统计: 24 warps = 768 threads = 6 warpgroups**

**配置 2: 256-epilogue (dispatch+TMA 共享 WG0, compact)**

| Warp Index | 数量 | 所属 Warpgroup | Role |
|---:|---:|---|---|
| 0–1 | 2 warps (64 threads) | WG0 | Dispatch |
| 2–3 | 2 warps (64 threads) | WG0 | TMA A+SFA / B+SFB |
| 4–11 | 8 warps (256 threads) | WG1, WG2 | Math WGMMA + epilogue + combine |

**精确 warp 统计: 12 warps = 384 threads = 3 warpgroups**

**配置 3: ≤256-total (dispatch+TMA 共享 WG0, BLOCK_M=32)**

| Warp Index | 数量 | 所属 Warpgroup | Role |
|---:|---:|---|---|
| 0–1 | 2 warps (64 threads) | WG0 | Dispatch |
| 2–3 | 2 warps (64 threads) | WG0 | TMA A+SFA / B+SFB |
| 4–7 | 4 warps (128 threads) | WG1 | Math WGMMA + epilogue + combine |

**精确 warp 统计: 8 warps = 256 threads = 2 warpgroups** |

**关键差异 vs SM100:**
- 无 TMEM，WGMMA accumulator 在 register 中
- SF 为 per-128 float（非 UE8M0 int）
- cluster_size ≤ 2（无 2-CTA UMMA）
- dispatch 最小仅 64 threads（SM100 要求 ≥128），意味着 dispatch 可以和 TMA 共享 warpgroup

**寄存器分配（多配置）:**

| 配置 | Dispatch Threads | TMA Threads | Epilogue Threads | Total Threads | Disp Reg | TMA Reg | Epi Reg | Total Reg |
|---|---|---|---|---|---|---|---|---|
| 512-epilogue | 128 | 128 | 512 | 768 | 32 | 24 | 112 | 64,512 |
| 256-epilogue | 64 | 64 | 256 | 384 | 48 | 40 | 168 | 48,640 |
| ≤256-total | 64 | 64 | 128 | 256 | 48 | 40 | 256 | 38,400 |

`kEpilogueRegisterBudget` 模板参数允许调用方精确控制 math warpgroup 的 register 分配。

**调度特点:**
- `sm90_fp8_mega_moe_for_each_block_split` 将 L1/L2 phase 拆分到不同 CTA，每个 CTA 只执行一种 phase
- 使用 `kNonEpilogueWarpsInWarpgroup` 确保 dispatch+TMA 恰好填满整数个 warpgroup

## 代码review发现

### 高: 512-epilogue 配置下 dispatch register 仅 32，可能成为瓶颈

dispatch warp 需要大量 `smem_expert_count` 和 route metadata 操作。32 register/warp 对于复杂的 `read_topk_idx` lambda 和 rank round-robin 选择可能不够（register spilling 到 local memory）。

### 中: `kEpilogueRegisterBudget` 默认 0 的自动推导逻辑复杂

`sm90_fp8_mega_moe.cuh:427-432` 的自动 register budget 推导依赖 thread 数和总线程数，外部调用者可能不知情使用次优配置。

### 中: 与 SM100 kernel 共享 scheduler，但 SM90 无 cluster 概念

`mega_moe.cuh` scheduler 包含 `kClusterSize=2` 相关逻辑，SM90 cluster_size=1 时跳过但增加了编译期复杂度。

### 低: 256-total-thread 配置下 epilogue 获得 256 reg/warp，几乎占满 budget

## 正面评价

- 首次将 MegaMoE 带到 SM90/Hopper，填补了架构空白
- 灵活的 register budget 控制允许调用方按场景调优
- L1/L2 phase 拆分减少了单一 CTA 的复杂度

## 建议检查清单

- 在 H100/H200 上实测 512-epilogue 配置的 dispatch spilling 情况
- 简化或文档化 `kEpilogueRegisterBudget` 的默认行为
