# MegaMoE on SM100/SM90 研究现状报告

## 1. 概述

DeepGEMM 的 MegaMoE 是一个融合了跨节点 token dispatch、双线性层 GEMM（L1 SwiGLU + L2 Linear）、以及跨节点 combine reduce 的端到端 MoE 算子。最初在 SM100（Blackwell）上以 FP8/FP4 精度实现，目前社区正在积极探索将其移植到 SM90（Hopper/H100/H200/H20）架构，以及利用 CUDA 13.1 green context 实现 SM100 上的 split pipeline 并发。本报告对六个相关 worktree（pr-304、pr-323、pr-347、pr-352、pr-357、pr-360）进行了横向对比分析。

本报告将 **pr-304** 作为官方原始 SM100 fused MegaMoE 基线：它建立了 FP8 activation + FP4 weight、2-CTA cluster、TMEM/UMMA accumulator、UE8M0 scale factor 和 symmetric-memory workspace 的完整形态。后续 SM90 章节中的对比都以这个基线为参照，重点说明 Hopper 缺少 TMEM/UMMA 后需要重新设计的部分。

---

## 2. Git 历史与分支关系

### 2.1 关键基线 commit

| Commit | 描述 |
|--------|------|
| `7f2a703ed` | [Public release 26/04] Introducing Mega MoE, FP4 Indexer (#304) |
| `891d57b4d` | Add various optimizations and Mega MoE benchmarks (#316) |
| `714dd1a4a` | main 分支 HEAD: Update test_mega_moe.py |

### 2.2 各 worktree 概览

| Worktree | HEAD commit | 基于 | 独有 commit 数 | 核心主题 | 来源 |
|----------|------------|------|---------------|---------|------|
| **pr-304** | `211d2678d` | `d30fc36c8` (before #292) | 7 | 原始 MegaMoE 引入 (SM100 only) | 官方 |
| **pr-323** | `23f46aa68` | `714dd1a4a` (main HEAD) | 2 | SM90 MegaMoE **fused kernel** decode 支持 | 官方 |
| **pr-347** | `2b8dfd0e8` | `714dd1a4a` (main HEAD) | 1 | 大规模重构（无 SM90 mega kernel） | 官方 |
| **pr-352** | `655075ef3` | `891d57b4d` (before main) | 16 | SM90 MegaMoE **split L1/L2 kernel** 优化 | 官方 |
| **pr-357** | `bb837421b` | `714dd1a4a` (main HEAD) | 2 | SM100 **green-context split-kernel** MegaMoE + FP16 weights MQA | RayWang96 |
| **pr-360** | `f27fd5611` | `88965b078` | 7 | SM90 **pingpong/cooperative** FP8 MegaMoE 双 kernel | mpdfdfl |

### 2.3 血缘关系图

```
d30fc36c8 (#292)                                                     
    │                                                                
    ├── 7f2a703ed (#304: 引入 MegaMoE)                                 
    │       │                                                        
    │       ├── 891d57b4d (#316: 优化+benchmark)                       
    │       │       │                                                
    │       │       ├── 714dd1a4a (main HEAD)                        
    │       │       │       │                                        
    │       │       │       ├── pr-323: SM90 fused decode 支持         
    │       │       │       │   (继承 pr-347)                          
    │       │       │       │                                        
    │       │       │       ├── pr-347: 重构 refactoring               
    │       │       │       │                                        
    │       │       │       └── pr-357: SM100 green-context split    
    │       │       │           (RayWang96, 独立于 SM90 路径)             
    │       │       │                                                
    │       │       └── pr-352: SM90 split L1/L2 kernel              
    │       │           (独立开发，16 个 commit)                           
    │       │                                                        
    │       └── pr-304 worktree: (附加 #256 + README 更新)               
    │                                                                
└── 88965b078 (PR347 squash merge: Multiple updates and refactorings)
        │                                                            
        └── pr-360: SM90 pingpong/cooperative                        
            (mpdfdfl, 基于 PR347)                                      
                                                                     
```

**关键发现**：
- **pr-304 是官方原始 SM100 MegaMoE 基线**，入口为 `sm100_fp8_fp4_mega_moe_impl` / `fp8_fp4_mega_moe`，采用 FP8×FP4、2-CTA cluster、TMEM/UMMA 和单 fused kernel。
- **pr-323 和 pr-352 是两条平行的、独立的 SM90 MegaMoE 实现路径**。
- pr-323 在 pr-347（重构）的基础上构建，采用 **fused kernel** 架构。
- pr-352 从 `#316` 分叉独立开发，采用 **split kernel** 架构。
- pr-347 本身不包含 SM90 MegaMoE kernel，但为 SM90 路径提供了必要的调度器/基础设施重构。
- pr-304 无 SM90 支持，但其 scheduler、workspace、weight transform、SF layout 是 SM90 移植需要对齐或替代的参考实现。
- **pr-357 是 SM100 only 的 green-context split pipeline**，与 SM90 路径无关，目标是利用 CUDA 13.1 green context 实现 K1/K2 并发。
- **pr-360 是第三套独立的 SM90 MegaMoE 实现**，引入 pingpong（warpgroup 交替）+ cooperative（warpgroup 协作共享 weight load）双 kernel 策略。

---

## 3. MegaMoE Kernel 技术对比

### 3.1 核心架构差异

| 维度 | pr-304 (SM100 fused) | pr-323 (SM90 fused) | pr-352 (SM90 split) | pr-360 (SM90 pingpong/coop) |
|------|----------------------|---------------------|---------------------|------------------------------|
| **目标架构** | SM100 only | SM90 | SM90 | SM90 |
| **精度路径** | FP8 activation + FP4 weight | FP8 activation + FP8 weight | FP8 activation + FP8 weight | FP8 activation + FP8 weight，block-(128,128) float SF |
| **Kernel 数量** | 1 个 fused kernel（dispatch→L1→L2→combine） | 1 个 fused kernel（L1→L2 内部状态机切换） | 2 个独立 kernel（分别 JIT 编译） | 1 个 fused kernel，auto 路由到 pingpong 或 cooperative |
| **Kernel 入口** | `sm100_fp8_fp4_mega_moe_impl` | `sm90_fp8_mega_moe_impl` | `sm90_fp8_mega_moe_l1_impl` + `sm90_fp8_mega_moe_l2_impl` | `sm90_fp8_mega_moe_pingpong_impl` / `sm90_fp8_mega_moe_cooperative_impl` |
| **核心 MMA** | SM100 UMMA / `tcgen05`，accumulator 在 TMEM | SM90 WGMMA，accumulator 在 registers | SM90 `mma.sync` / WGMMA，accumulator 在 registers 或 SMEM staging | SM90 WGMMA，register accumulator；pingpong 重叠 MMA/epilogue，coop 两 WG M-split |
| **dispatch 执行** | 128-thread dispatch WG，负责计数、NVLink pull、workspace 清理 | dispatch warp 在 L1 阶段拉取 token | dispatch 仅在 L1 kernel 中执行，L2 kernel 跳过 | 64 dispatch threads，与 64 TMA/non-epi threads 共享 WG0 |
| **combine 执行** | epilogue warps 在 L2 后执行 top-k BF16 reduce | 在 L2 完成后执行 combine reduce | 仅在 L2 kernel 中执行 combine | math/epilogue warps 在 same fused kernel 尾部执行 combine，逻辑从 SM100 port |
| **代码规模** | ~1364 行 kernel 代码 | ~1935 行 kernel 代码 | ~2507 行 kernel 代码 | ~1418 + ~1419 行双 kernel 代码 |

**设计理念差异**：
- **pr-304 SM100 fused**: 单 kernel 内用 2-CTA cluster + TMEM/UMMA 表达整个 MoE pipeline。Dispatch、TMA/UMMA issue、TMEM epilogue、combine 共存在 512-thread CTA 中；2 个 CTA 通过 `Allocator2Sm` 共享 tensor memory pipeline，是后续实现的功能基线而非 SM90 可直接复用代码。
- **pr-323 fused**: 单次 kernel launch 完成全流程，减少 kernel launch overhead，适合 token-per-expert 较小的 decode 场景。通过 `kSplitPhaseHotPath` 模板参数控制是否在 kernel 内部快速切换 phase。
- **pr-352 split**: 两次独立的 kernel launch，允许 L1 和 L2 使用**完全不同的线程配置、register budget、cluster size**。L1 kernel 不需要 combine 逻辑，L2 kernel 不需要 dispatch 逻辑，各自优化更彻底。
- **pr-360 pingpong/cooperative**: 保留 single-launch fused 目标，但拆成两套 SM90 kernel。`pingpong` 用两个 math WG 处理不同 tile 并通过 `OrderedSequenceBarrier<2,2>` 交替 MMA/epilogue；`cooperative` 用两个 math WG M-split 同一 tile，共享一次 B/weight TMA load，减少大 M 下的 weight HBM 流量。

### 3.2 BLOCK 配置灵活性

| 配置项 | pr-304 (SM100) | pr-323 | pr-352 | pr-360 |
|--------|----------------|--------|--------|--------|
| **BLOCK_M** | 192（固定，`get_block_m_for_mega_moe` 直接返回） | 64（decode 场景），128（auto_split_mn） | 16, 32, 64, 96, 128, 192 | pingpong=64；cooperative=128 |
| **BLOCK_N** | 128（固定） | 128 或 256 | 64, 128, 256 | 128（固定） |
| **BLOCK_K** | 128（固定；UMMA K step 为 32） | 128（固定） | 128（固定） | 128（固定；WGMMA K step 32，L2 acts SF per-64） |
| **Load / Store M** | `load_block_m=96`, `store_block_m=32` | 依 kernel 配置 | 依候选配置 | load=`BLOCK_M`；pingpong L1 store box=`BLOCK_M`，coop store box=`BLOCK_M/2` |
| **Cluster Size** | 2（`LaunchArgs(..., cluster_dim=2)`，`cute::TMEM::Allocator2Sm`） | 1（固定，无 TMA multicast） | 1 或 2（支持 B/weight tile TMA multicast） | 1（源码注释说明 cluster=2 可行但延后，因 L1 amax/SF 需跨 CTA 同步） |
| **MMA 模式** | SM100 UMMA (`SM100_MMA_MXF8F6F4_2x1SM_SS`) + TMEM | 仅 WGMMA | 支持 `mma.sync` (M16/M32) + WGMMA | WGMMA m64n128k32，两个 math WGs 固定 |

**pr-352 的 BLOCK_M=16/32 支持** 是关键差异化 feature：
- 使用 `mma.sync` 指令替代 WGMMA，专门针对 token-per-expert 极小的 decode 场景（如 RL long-tail rollout）。
- `mma.sync` 的 FP32 accumulator 通过 SMEM staging 实现 epilogue logical row/column 映射。
- 此时 swizzle mode 设为 0（禁用 TMA swizzle，使用 32B 粒度加载）。

### 3.3 Warpgroup 拆分策略

| 维度 | pr-304 (SM100) | pr-323 | pr-352 | pr-360 |
|------|----------------|--------|--------|--------|
| **CTA threads** | 512 = 128 dispatch + 128 non-epilogue + 256 epilogue | 256/384/768 | 384/512 | 384 = 64 dispatch + 64 TMA/non-epi + 256 math/epi |
| **Dispatch WG** | WG0: warps 0-3，全 128 threads | 独立 dispatch/TMA 或简化配置 | compact 配置下 dispatch+TMA 可共享 WG | WG0 warps 0-1 dispatch，warps 2-3 TMA A/B，共享 `dealloc<48>` |
| **TMA / MMA issue WG** | WG1: warp4 A/SFA TMA, warp5 B/SFB TMA, warp6 UMMA issue, warp7 idle/reserved | TMA warps + math WGs | L1/L2 可独立配置 TMA/math WGs | TMA 在 WG0；WGMMA 由 WG1/WG2 math warpgroups 各自发出 |
| **Epilogue WGs** | WG2-3: 8 warps；2 个 WG 按 M 方向切 `BLOCK_M/2=96` 行，每 WG 内 4 warps 覆盖 N=128 | Split-N: 2 WG split N | Serial/Wide/Fused 变体 | WG1/WG2 固定 2 个；pingpong 按 tile 奇偶交替，coop 按 M rows [0,64)/[64,128) 协作 |
| **Split-N** | 不使用 SM90 式 split-N；SM100 通过 TMEM+2CTA UMMA 覆盖 tile | 支持（BLOCK_M=64, BLOCK_N=256, 2 WG） | 支持（同左） | 不使用；N 固定 128 |
| **Split-MN** | 不使用；M 方向由 2-CTA cluster 和 2 epilogue WG 分摊 | 支持（BLOCK_M=128, BLOCK_N=256, 512 epilogue threads / 4 math WGs） | 不支持 | cooperative 是 M-split 同一 tile，但不是 pr-323 的 Split-MN N=256 模式 |
| **Serial-N** | 不适用 | 不支持 | 理论支持（`kSerialNWarpgroups`，实际禁用） | 不支持 |
| **wide-N** | 不适用 | 不支持 | 支持（BLOCK_N=256, 1 WG） | 不支持 |

PR304 的 warp 拆分不是 SM90 上的“多 math WG 持有各自 accumulator”模型。UMMA issue 由 leader CTA 的单个 warp 发出，结果落在 TMEM；8 个 epilogue warp 再从 TMEM load，执行 L1 SwiGLU/FP8 quant 或 L2 BF16 scatter/combine。这一点是 SM90 移植难点的根源：Hopper 没有 TMEM，math/epilogue 必须围绕 register accumulator 重新组织。

PR360 则把 SM90 的 2 个 math WG 明确分成两种策略：pingpong 是“不同 tile 的时间交错”，用 order barrier 让一个 WG 的 MMA 与另一个 WG 的 epilogue 重叠；cooperative 是“同一 tile 的空间切分”，两个 WG 各做 64 行，共享同一 B tile，L2 epilogue 末尾用 256-thread barrier 防 `smem_cd` 别名覆盖。

### 3.4 调度器 (Scheduler) 扩展

pr-304 的原始 scheduler 已经围绕 2-CTA cluster 设计，pr-347 和 pr-352 又在此基础上对调度器进行了不同方向的扩展：

**pr-304 原始 SM100 scheduler**：
- `MegaMoEScheduler` 对每个 expert wave 先遍历 L1 blocks，再回到 wave 起点遍历 L2 blocks，单 kernel 内完成 L1→L2 状态机切换。
- 保留严格约束：`kNumExpertsPerRank % kNumExpertsPerWave == 0`，并要求 `kNumSMs`、`kNumL1BlockNs`、`kNumL2BlockNs` 都能满足 2-CTA cluster 的偶数配对。
- `block_idx = blockIdx.x`，每次 `block_idx += kNumSMs` 做 persistent scheduling；2 个相邻 CTA 依赖 cluster 语义保持相同 M block、不同 CTA rank 的协同关系。
- `fetch_expert_recv_count()` 等待所有 rank 的 expert count 高 32 位达到 `kNumSMs * kNumRanks`，然后每 lane cache 本地 expert token 数。

**pr-347 引入的变更**（被 pr-323 继承）：
- 修复了 `kNumExpertsPerWave` 的整除约束变为更宽松的 `> 0 && <= kNumExpertsPerRank`
- `get_wave_expert_end_idx` 增加 clamp 保护部分 wave

**pr-352 引入的额外变更**：
- 新增 `kClusterSize` 模板参数（1 或 2），`block_idx = blockIdx.x / kClusterSize` 实现 cluster 内的 CTA 配对
- 新增 `kL2NMajorSchedule` 和 `kL1NMajorSchedule` 参数，支持 N-major block 分配（先分 N block 再分 M block）
- 新增 `for_each_phase_block<Phase>()` 方法，为 split kernel 提供单 phase 遍历接口
- M block 使用 `ceil_div(num_m_blocks, kClusterSize)` 确保 2-CTA cluster 中每个 cluster 有完整的 M block pair

**注意**: PR352 分叉自 #316（早于 PR347），因此其 scheduler 保留了旧的严格整除约束（`kNumExpertsPerRank % kNumExpertsPerWave == 0`）。PR304 原始 scheduler 也是严格整除口径；PR347 的放宽约束（`> 0 && <=`）仅被 PR323 继承。PR360 的独立 scheduler (`sm90_mega_moe.cuh`) 同样使用严格整除约束。

**`valid_value` 初始化问题**: PR304 原始 `mega_moe.cuh:70` 为 `uint32_t valid_value;`，PR357 的 K2 scheduler (`l2_combine.cuh:51`) 已修复为 `uint32_t valid_value = 0`。但 K1 沿用的 `mega_moe.cuh` 和 PR360 的 scheduler 仍有同类未初始化写法。

### 3.5 SMEM 与 Pipeline 管理

| 维度 | pr-304 (SM100) | pr-323 | pr-352 | pr-360 |
|------|----------------|--------|--------|--------|
| **SFA per stage** | `SF_BLOCK_M * sizeof(uint32_t)`，UTCCP 128-aligned | `align(2 * BLOCK_M * sizeof(float), 128)` | `2 * kL2SFAHalfStride * sizeof(float)`（更精确的 L2 per-64-K 半部分 stride 计算） | `align(2 * BLOCK_M * sizeof(float), 128)`；L1 per-128、L2 per-64 |
| **SFB per stage** | `SF_BLOCK_N * sizeof(uint32_t)`，UE8M0 packed | float SF | float SF | 0；weight SF 由 math WG 从 global 直接 `__ldg`，并做软件 pipeline |
| **Accumulator** | TMEM columns: accum + SFA + SFB，`kNumEpilogueStages=2` | registers | registers / SMEM staging | registers；`final_accum[64]` 持有 WGMMA 累积结果 |
| **CD output** | `max(L1_FP8 double TMA-store, L2_BF16)` | 简单 max(L1_FP8, L2_BF16) | 分层：CD_ACCUM（mma.sync）+ CD_L1 + CD_L2 + CD_L1_ASYNC | 单个 `smem_cd` 区域在 L1 FP8 row-major 与 L2 BF16 staging 间复用，1024B align |
| **Async L1 TMA store** | 支持（L1 output double-buffered TMA store） | 不支持 | 支持（double-buffered CD region） | L1 epilogue用 TMA store；pingpong 每 WG 整 tile store，coop 每 WG store 自己 64-row band |
| **Direct L2 scatter** | 支持（L2 BF16 通过 NVLink 写 remote combine buffer） | 不支持 | 支持（L2 输出通过 NVLink 直接 scatter，跳过 SMEM staging） | 支持；L2 BF16 staging 后通过 `sym_buffer.map` scatter 到 remote combine buffer |
| **L2 Dual Accum** | TMEM epilogue-stage double buffer | 不支持 | 支持（双 accumulator 流水线） | 不使用 dual accum；L2 K=128 拆为两个 per-64 SFA half，各自 WGMMA + scale |
| **L1 Dual-K Accum** | 依赖 `kNumStages` 的 TMA/UMMA K pipeline | 不支持 | 支持（L1 K 维度双缓冲） | 不使用 dual accum；通过 TMA stages + WGMMA wait/commit overlap SFA/weight SF loads |
| **Combine** | 复用 pre-barrier SMEM 区域，1 或 2 chunk | 复用 pre-barrier SMEM 区域 | 同左，但更精确的 chunk 计算（1 或 2 chunk） | 复用 pre-barrier SMEM 区域，SM100 combine 逻辑 port，`kNumChunks` 取 1 或 2 |

### 3.6 L1 Epilogue 差异

三条 fused/split 路径的 L1 epilogue 都执行 SwiGLU + per-row amax + FP8 量化 + store，但实现位置和 scale 格式差别很大：

- **pr-304**: epilogue WG 从 TMEM 通过 `SM100_TMEM_LOAD_16dp256b1x` 读取 accumulator，在 registers 中执行 gate/up SwiGLU、top-k weight 乘法、跨 warp-pair amax reduce，生成 UE8M0 scale 并写入 `l2_acts_sf`，再用 STSM + TMA store 写 FP8 L2 acts。完成一个 L1 output block 后用 `red_or_rel_gpu` 设置 `l2_arrival_mask` 唤醒 L2。
- **pr-323**: 支持 shared-SF 机制 —— 当 split-N 的 `WG_L1_OUT_BLOCK_N < 64` 时，两个 split-N warpgroup 共享同一个 per-64-K 的 L2 acts SF slot，需要跨 warpgroup amax reduce。
- **pr-352**: 简化为每个 warpgroup 独立管理自己的 SF，不支持 shared-SF。但通过 `kAsyncL1TMAStore` 支持异步 TMA store（double-buffered CD region）。
- **pr-360**: L1 epilogue 在 math WG registers 中执行 SwiGLU、top-k weight 乘法、per-row amax、FP8 quant，然后写 float SF 到 `l2_acts_sf` 的 per-64 layout。pingpong 由拥有 tile 的 WG 自己设置 `l2_arrival_mask`；cooperative 必须等两个 WG 的 64-row band 都 TMA store 完成后，经 256-thread barrier 再设置 mask。

---

## 4. Warp 数据流图解 (pr-304 / pr-323 / pr-352 / pr-360)

### 通用符号说明

```
[WG0]  Warpgroup (4 warps = 128 threads)
  w0    单个 warp (32 threads)
  -->   数据流
  - ->  同步/等待信号 (barrier / arrival count)
  ~~~   TMA async copy 边界
```

### 4.1 pr-304: SM100 Fused MegaMoE

```
SM100 fused MegaMoE: one persistent 2-CTA cluster per SM
Each CTA has 512 threads = WG0 dispatch + WG1 TMA/UMMA issue + WG2/WG3 epilogue

Roles per CTA
  WG0  w0-w3    dispatch, NVLink pull, workspace cleanup
  WG1  w4       TMA A + SFA
       w5       TMA B + SFB
       w6       leader CTA issues tcgen05 UMMA and UTCCP SF -> TMEM
       w7       reserved / register-deallocated
  WG2  w8-wB    TMEM epilogue half 0
  WG3  wC-wF    TMEM epilogue half 1

Cross-rank dispatch path
  source rank input token / SF / top-k weight
    -> WG0 counts top-k, writes source metadata and expert recv counts
    -> grid sync + NVLink barrier
    -> WG0 pulls remote token/SF/top-k weight into local L1 pool
    -> l1_arrival_count[pool_block] release-add

Per scheduled tile, same TMA/UMMA/TMEM pipeline handles both phases

  Linear1 phase
    local L1 pool ready
      -> w4 waits l1_arrival_count, TMA loads L1 acts + SFA
      -> w5 TMA loads L1 FP4 weights + SFB
      -> full_barrier
      -> w6 issues tcgen05 UMMA, accum/SF live in TMEM
      -> WG2/WG3 load TMEM
      -> SwiGLU * top-k weight, amax, UE8M0 L2 act SF
      -> STSM + TMA store FP8 l2_acts / l2_acts_sf
      -> l2_arrival_mask bit release

  Linear2 phase
    L1 output ready
      -> w4 waits l2_arrival_mask for required L1 N blocks
      -> w4 TMA loads l2_acts + l2_acts_sf
      -> w5 TMA loads L2 FP4 weights + SFB
      -> full_barrier
      -> w6 issues tcgen05 UMMA into TMEM
      -> WG2/WG3 load TMEM
      -> BF16 stage in SMEM
      -> NVLink scatter to remote combine buffer

After all L2 scatter
  epilogue NVLink barrier
    -> combine TMA-loads top-k BF16 chunks from combine buffer
    -> FP32 reduce
    -> TMA store output y
```

**PR304 warp role map:**

| Warp(s) | WG | Role |
|---|---|---|
| 0-3 | WG0 | Dispatch: expert count, source metadata, NVLink token/SF pull, workspace cleanup |
| 4 | WG1 | TMA load activations and SFA; waits `l1_arrival_count` for L1 and `l2_arrival_mask` for L2 |
| 5 | WG1 | TMA load FP4 weights and SFB |
| 6 | WG1 | Leader CTA issues `tcgen05` UMMA and UTCCP copies SF to TMEM |
| 7 | WG1 | Register-deallocated reserved warp; no substantial work in original source |
| 8-15 | WG2/WG3 | TMEM epilogue, L1 FP8 quant/TMA store, L2 BF16 scatter；selected epilogue warps run final combine reduce |

**PR304 同步关键点:**

| 同步 | 机制 |
|---|---|
| 2-CTA cluster 初始化 | `cute::cluster_sync()` before/after `Allocator2Sm::allocate` |
| source metadata/count visible before pull | grid sync + cross-rank `nvlink_barrier` |
| dispatch pull-->L1 TMA | `l1_arrival_count` per pool block |
| L1 epilogue-->L2 TMA | `l2_arrival_mask` bit per L1 N block |
| TMA-->UMMA | `full_barriers[stage].wait(phase)` |
| UMMA-->TMA recycle | `empty_barriers[stage]` with 2-CTA UMMA multicast arrive |
| UMMA-->epilogue | `tmem_full_barriers[accum_stage]` |
| epilogue-->UMMA recycle | `tmem_empty_barriers[accum_stage]` |
| L2 scatter-->combine | cross-rank `nvlink_barrier` |

### 4.2 pr-323: SM90 Fused MegaMoE (256-epilogue 配置)

```
SM90 fused MegaMoE decode split-N path: one CTA, 384 threads
  WG0 w0-w1    dispatch / remote pull
  WG0 w2       TMA A + SFA
  WG0 w3       TMA B only; weight SF is loaded directly by math WGs from global
  WG1 w4-w7    WGMMA + epilogue for one N slice
  WG2 w8-wB    WGMMA + epilogue for the other N slice

Dispatch path
  input_topk_idx
    -> WG0 counts tokens per expert and writes source metadata to destination ranks
    -> grid sync + NVLink barrier
    -> WG0 pulls remote token, input SF, top-k weight
    -> local L1 token/SF/top-k-weight pool
    -> l1_arrival_count[pool_block] release-add

Linear1 phase
  local L1 pool ready
    -> w2 waits l1_arrival_count, TMA loads L1 acts + float SFA into SMEM
    -> w3 TMA loads L1 FP8 weights into SMEM
    -> WG1/WG2 load L1 weight SF via __ldg from global
    -> full_barrier
    -> WGMMA accumulates in registers
    -> L1 epilogue in registers: SwiGLU, top-k weight, per-row amax, FP8 quant
    -> optional cross-WG amax if split-N shares one per-64 SF group
    -> smem_cd_l1
    -> TMA store l2_acts and write float l2_acts_sf
    -> l2_arrival_mask/counter release

Linear2 phase
  L1 output ready
    -> w2 waits l2_arrival_mask/counter
    -> w2 TMA loads l2_acts and two per-64 SFA halves
    -> w3 TMA loads L2 FP8 weights into SMEM
    -> WG1/WG2 load L2 weight SF via __ldg from global
    -> full_barrier
    -> WGMMA accumulates in registers; K=128 is scaled as two per-64 halves
    -> BF16 stage in smem_cd_l2
    -> NVLink scatter to remote combine buffer

After all L2 scatter
  epilogue NVLink barrier
    -> combine TMA-loads top-k BF16 chunks from combine buffer
    -> FP32 reduce
    -> TMA store output y
```

**数据流关键路径:**
```
Remote ranks -> dispatch pull -> local L1 pool
local L1 pool --l1_arrival_count--> L1 TMA A/SFA
L1 weights -----------------------> L1 TMA B
weight SF global -----------------> math WG __ldg
TMA full_barrier -----------------> WGMMA registers
register epilogue ----------------> l2_acts/l2_acts_sf + l2_arrival release
l2_acts/l2_acts_sf ---------------> L2 TMA A/SFA
L2 weights -----------------------> L2 TMA B
L2 WGMMA registers ---------------> BF16 remote combine buffer
combine buffer -------------------> top-k reduce -> y
```

**Warp 间同步关键点:**
| 同步 | 机制 |
|---|---|
| WG0 内部 dispatch warps | sync_aligned(64) |
| source metadata/count visible before pull | grid_sync + NVLink barrier |
| dispatch pull-->L1 TMA | `l1_arrival_count` release/acquire |
| L1 epilogue-->L2 TMA | `l2_arrival_mask` bitmask or counter release/acquire |
| TMA-->math (SMEM ready) | `full_barriers[stage].wait(phase)` |
| math-->TMA (SMEM consumed) | `empty_barriers[stage].arrive()` / `wait()` |
| WG1<-->WG2 split-N | WG-local sync normally；full epilogue sync when shared-SF/shared tile is enabled |
| L2 scatter-->combine | epilogue `nvlink_barrier` |

### 4.3 pr-352: SM90 Split L1/L2 MegaMoE (Compact 配置)

```
+==========================================+                        
|   KERNEL 1: L1 (dispatch+GEMM+SwiGLU)    |                        
+==========================================+                        
|  dispatch=64    TMA=64    math=256       |                       
| +---------+  +---------+  +-------------+|                       
| |WG0  w0  |  |WG0  w2  |  | WG1    WG2  ||                       
| |     w1  |  |     w3  |  | w4..7 w8..B ||                       
| +----+----+  +----+----+  +------+------+|                       
|      | dispatch  | TMA load     | WGMMA  |                        
|      | pull      | A+SFA,B+SFB  | SiLU   |                        
|      | -->L1 pool| --> smem     | quant  |                        
|      |           |              | TMA    |                        
|      | l1_arrival|              | store  |                        
|      | _count    | full/empty   | l2_acts|                        
|      +-----------+---barrier----+--------+                        
+==========================================+                        
                          |                                         
             +------------v------------+                            
             |      Global Memory      |                            
             |   l2_acts (FP8)         | <-- K1 writes, K2 consumes 
             |   l2_arrival_mask       | <-- K1 sets bit per N-block
             +------------+------------+                            
                          |                                         
+==========================================+                        
|   KERNEL 2: L2 (GEMM+Combine Scatter)    |                        
+==========================================+                        
|  TMA=128               math=128/256      |                        
| +------------------+  +---------------+  |                        
| |WG0 w0 w1 w2 w3   |  |WG1 (WG2)      |  |                         
| |A+SFA B+SFB idle  |  |math WG: WGMMA |  |                        
| +--------+---------+  +-------+-------+  |                        
|          | TMA load           |          |                        
|          | <-- l2_arrival_mask| BF16     |                        
|          |   (busy-wait K1)   | scatter  |                        
|          +--------------------+----------+                        
+==========================================+                        
                                                                    
```

**K1-->K2 跨 kernel 同步:**
```
K1 dispatch 写完 pool block  --> l1_arrival_count release-add
K1 TMA load 等待 l1_arrival_count == num_tokens_per_block
K1 math epilogue 写完 L1 output block -->
    l2_arrival_mask <-- red_or_rel(bit n_block_idx)

K2 scheduler busy-wait:
    while(expert_recv_count_sum.high != kKernel1SMs*kNumRanks);
    while(l2_arrival_mask != expected_mask);
```

### 4.4 pr-360: Pingpong

```
+==============================================================+  
|         Single Fused CTA (384 threads, 12 warps)             | 
+==============================================================+  
| WG0 (dispatch+TMA)    WG1(math)           WG2(math)          | 
| +--------+--------+  +--------------+   +--------------+     | 
| |w0  w1  | w2  w3 |  |w4 w5 w6 w7   |   |w8 w9 wA wB   |     |   
| |dispatch| A    B |  |  Math WG0    |   |  Math WG1    |     |  
| | 64     | loader |  | BLOCK_M=64   |   | BLOCK_M=64   |     |  
| +---+----+---+----+  +------+-------+   +------+-------+     | 
|     |        |              |                  |             |
|     |dispatch| TMA load     | WGMMA(tile0)     |             |
|     |pull    | A+SFA B+SFB  |   |              |             |
|     |-->L1   | --> smem     | SiLU(tile0)      |             |
|     | pool   |              |   |              |             |
|     |   +----+--order_arrive->|(epi tile0)     | WGMMA(t1)   |
|     |   |    |              |                  |   |         |
|     |   | <--order_arrive---+------------------+ SiLU(t1)    | 
|     |   |    |              | WGMMA(tile2)     |   |         |
|     |   +----+--order_arrive->|  ....          |(epi t1)     | 
+=====+========+==============+==================+=============+
                                                                 
```

**Pingpong 时序展开:**
```
Time -->                                                          
WG1(Math0): [WGMMA t0]--[epi t0]--[WGMMA t2]--[epi t2]--          
               | order_arrive | order_arrive |                    
               v              v              v                    
WG2(Math1):    |[WGMMA t1]----[epi t1]----[WGMMA t3]--[epi t3]    
                +-- order_barrier --+                             
                                                                  
OrderedSequenceBarrier<2,2> state machine (per WG):               
  ord_stage 0 = MMA region, ord_stage 1 = Epilogue region         
  WG0 initial phase=1 (wait(1) passes immediately)                
  WG1 initial phase=0 (wait(0) blocks until WG0 arrives)          
  Loop: order_wait --> execute --> order_arrive --> ord_stage ^= 1
                                                                  
```

### 4.5 pr-360: Cooperative

```
+================================================================+             
|         Single Fused CTA (384 threads, 12 warps)               |            
+================================================================+             
| WG0 (dispatch+TMA)     WG1(math)            WG2(math)          |            
| +--------+--------+  +----------------+  +----------------+    |            
| |w0  w1  | w2  w3 |  |w4 w5 w6 w7     |  |w8 w9 wA wB     |    |              
| |dispatch| A    B  |  | rows [0, 64)  |  | rows [64,128)  |    |            
| +---+----+---+----+  +-------+--------+  +-------+--------+    |            
|     |        |               |                   |             |           
|     |        | TMA load  +---+-------------------+----+        |            
|     |        |           | smem_A[0:64] smem_A[64:128] |       |           
|     |        |           | smem_B ===== SHARED! =======|       |           
|     |        |           +---+-------------------+----+        |            
|     |        |               |  WGMMA(rows0)   WGMMA(rows1)    |            
|     |        |               |    SiLU(rows0)   SiLU(rows1)    |            
|     |        |               |       TMA store    TMA store    |           
|     |        |               |                   |             |           
|     |        |          256-thread barrier <-- wait both WGs   |            
+=====+========+===============+===================+==============+           
                                                                              
M-Split layout:                                                               
  smem_A (BLOCK_M=128 rows)     smem_B (128 cols)                             
  +---------------------+       +----------+                                  
  | rows [0, 64)  WG1  |       |  SHARED  | <-- TMA loads only ONCE           
  | rows [64,128) WG2  |       | WG1+WG2  |                                   
  +---------------------+       +----------+                                  
                                                                              
  smem_cd (128 rows x L1_OUT_BLOCK_N)                                         
  +---------------------+                                                     
  | rows [0, 64)  WG1  | <-- WG1 TMA store                                    
  | rows [64,128) WG2  | <-- WG2 TMA store                                    
  +---------------------+                                                     
                                                                              
  WARNING smem_cd aliasing hazard:                                            
    L2 epilogue: smem_cd_l2 (BF16) and smem_cd_l1 (FP8) alias same smem bytes.
    WG0 writing NEXT tile's smem_cd could overlap WG1 still reading           
    CURRENT tile's smem_cd_l2 in NVLink scatter.                              
    --> 256-thread barrier (kEpilogueFullBarrierIdx) prevents this.           
                                                                              
```

### 4.6 四实现横向对比

| 维度 | pr-304 SM100 | pr-323 SM90 | pr-352 SM90 | pr-360 SM90 |
|---|---|---|---|---|
| Kernel count | 1 | 1 | 2 | 1 |
| CTA / cluster | 512-thread CTA × 2-CTA cluster | single CTA | L1/L2 independent CTA configs | single 384-thread CTA |
| Dispatch --> L1 | same kernel, full dispatch WG | same CTA state machine | K1 independent kernel | same CTA |
| L1 --> L2 data | global `l2_acts` in symmetric buffer | same CTA/global l2_acts path | via global memory | same CTA/global l2_acts path |
| L1 --> L2 sync | `l2_arrival_mask` within fused kernel | implicit CTA order | `l2_arrival_mask` cross-kernel | implicit CTA order |
| Accumulator storage | TMEM | registers | registers / SMEM staging | registers |
| Math strategy | 2-CTA UMMA issue, epilogue reads TMEM | Split-N / Split-MN WGMMA | `mma.sync` or WGMMA variants | Pingpong alternation or cooperative M-split |
| Weight load optimization | FP4 halves weight bytes; 2SM TMA pipeline | none | cluster multicast for B/weight in selected configs | cooperative shares B-load across two math WGs |
| L2 --> Combine | epilogue tail in same kernel | same CTA tail | K2 tail | same CTA tail |

### 4.7 同步原语对比

| 同步点 | pr-304 | pr-323 | pr-352 | pr-360 |
|---|---|---|---|---|
| dispatch-->TMA | `l1_arrival_count` + NVLink barrier | grid_sync + NVLink barrier | same | sync_unaligned + NVLink barrier |
| TMA-->math (smem ready) | full/empty barrier | full/empty barrier | same | same |
| math-->TMA (smem consumed) | UMMA arrive on empty barrier | empty_barrier.arrive | same | same |
| UMMA/WGMMA-->epilogue | TMEM full/empty barriers | register-local | register-local / SMEM staging | register-local |
| L1-->L2 phase switch | `l2_arrival_mask` inside fused kernel | CTA internal order | `l2_arrival_mask` (cross-kernel) | CTA internal order |
| WG1<-->WG2 (Pingpong) | -- | -- | -- | OrderedSequenceBarrier<2,2> |
| WG1<-->WG2 (Coop smem_cd) | -- | -- | -- | sync_aligned(256 threads) |
| L2-->combine | NVLink barrier | NVLink barrier | NVLink barrier | NVLink barrier |

---

## 5. Heuristics 策略对比

### 5.1 配置选择器架构

| 维度 | pr-304 (SM100) | pr-323 | pr-352 | pr-360 |
|------|----------------|--------|--------|--------|
| **代码规模** | ~211 行 | ~180 行 | ~1025 行 | ~320 行 |
| **选择方式** | 固定主 tile + wave/stage 动态计算 | 确定性规则选择单个 config | 候选生成 + 评分排序 + 选择最优 | 两套固定 block config + token-count auto routing |
| **Token-per-expert 分级** | 无显式分档；`num_experts_per_wave` 由 L1 blocks/SM 覆盖度决定 | 2 档（`< 64` vs `>= 64`） | 6 档（`<=8.5, 16.5, 32.5, 64.5, 96.5, else`） | kernel 选择按 `num_tokens` 阈值；N-major 按 tokens/expert ≥256 |
| **设备适配** | SM100-only，使用 `SM100ArchSpec::smem_capacity` | 无 | H20（LowSm ≤80 SMs）vs H200（HighSm ≥100 SMs）vs Generic | SM90-only，`cluster_size=1`，使用 `SM90ArchSpec::smem_capacity` |
| **Profile 驱动** | 无 | 无 | 基于 `DG_SM90_MOE_DEVICE_PROFILE` 环境变量的 per-bucket 最佳实践 | 无 profile 表；提供 `DG_SM90_MOE_KERNEL`、`DG_SM90_MOE_NMAJOR`、threshold override |

### 5.2 pr-304 的 SM100 配置逻辑

PR304 的 heuristics 不是 profile-driven 搜索，而是把 SM100 fused kernel 的形状固定下来，然后按 workspace 和 SM 数动态补齐 wave/stage：

```
block_m = 192
block_n = 128
block_k = 128
load_block_m = block_m / 2 = 96
load_block_n = 128
store_block_m = 32
num_dispatch_threads = 128
num_non_epilogue_threads = 128
num_epilogue_threads = 256
```

关键规则：
- `get_block_m_for_mega_moe()` 直接返回 192，Python buffer 层会将 `num_max_tokens_per_rank` 对齐到该 block size。
- `num_experts_per_wave` 先找不超过 32 的最大 divisor，再按 `2 * num_sms / num_l1_blocks_per_expert` 估算需要多少 expert 才能填满 SM，最后向上取整到 `num_experts_per_rank` 的 divisor。
- `num_stages` 由 `SM100ArchSpec::smem_capacity` 扣除 dispatch region、CD region、barriers、amax reduction、TMEM pointer 后最大化，要求至少 2 stages。
- SF block size 通过 `SM100ArchSpec::get_sf_uttcp_aligned_block_sizes(..., MmaKind::MXFP8FP4)` 获取，保证 UTCCP 4x32 / 128-aligned scale layout。

### 5.3 pr-352 的 profile 驱动配置示例

对于主配置（`num_experts_per_rank=32, num_topk=8, intermediate_hidden=2048`），H200（HighSm）上的 per-bucket 最优配置：

| Tokens/Expert | Experts/Wave | Stages | Direct L2 Scatter | L2 N-Major | One-Warp Cleanup |
|---------------|-------------|--------|-------------------|------------|-----------------|
| ≤3 | 32 | 4 | ✓ | ✓ | ✗ |
| ≤6 | 32 | 4 | ✗ | ✓ | ✓ |
| ≤12 | 32 | 4 | ✓ | ✗ | ✓ |
| ≤24 | 32 | 4 | ✗ | ✓ | ✓ |
| ≤48 | 32 | 4 | ✓ | ✗ | ✓ |
| ≤64 | 32 | 4 | ✗ | ✓ | ✓ |
| ≤160 | 32 | 4 | ✗ | ✓ | ✗ |
| ≤240 | 32 | 4 | ✗ | ✓ | ✗ |
| ≤384 | 16 | 4 | ✗ | ✓ | ✗ |
| ≤640 | 32 | 4 | ✗ | ✓ | ✗ |
| ≤896 | 32 | 4 | ✗ | ✓ | ✗ |
| ≤1536 | 32 | 4 | ✗ | ✓ | ✓ |

### 5.4 pr-323 的配置逻辑

```
expected_tokens_per_expert >= 64 → auto_split_mn: BLOCK_M=128, BLOCK_N=256, epilogue=512 threads (4 WGs)
else → decode path:
    BLOCK_M=64, BLOCK_N=128 (或 256 如果 intermediate_hidden >= 3072
                             且 `expected_tokens_per_expert >= 0.25`)
    num_epilogue_threads = 256 (decode split-n) 或 128
```

### 5.5 pr-360 的 token-count 路由

PR360 不做 pr-352 式候选搜索，而是把 SM90 fused 路线固定为两套 kernel config：

| 条件 | Kernel | 主配置 | 设计动机 |
|------|--------|--------|----------|
| `num_tokens < DG_SM90_MOE_COOPERATIVE_THRESHOLD`（默认 256） | `pingpong` | `BLOCK_M=64, BLOCK_N=128, BLOCK_K=128`，384 threads | 小/中 M 下 epilogue/latency 更敏感，两个 math WG 交替执行不同 tile |
| `num_tokens >= threshold` | `cooperative` | `BLOCK_M=128, BLOCK_N=128, BLOCK_K=128`，384 threads | 大 M 下 weight 重复读取主导，两个 math WG M-split 同一 tile 并共享 B load |

两套 config 共享同一个 wave/stage 计算框架：`num_experts_per_wave` 复用 SM100 的 wave-balancing 逻辑，`num_stages` 根据 SM90 shared memory capacity、dispatch region、CD region、SFA region 和 barrier 数量最大化。PR360 当前把 `cluster_size` 固定为 1；源码注释明确 cluster=2 multicast 可行，但 L1 SwiGLU/FP8 quantize 的 per-128 activation SF 会需要跨 CTA amax 同步，因此暂缓。

L2 scheduling 方面，PR360 在 tokens/expert ≥256 时启用 N-major schedule，以减少大 M 下 weight tile 在 L2 中的 thrash；`DG_SM90_MOE_NMAJOR=-1|0|1` 可用于自动/关闭/强制。kernel 选择可用 `DG_SM90_MOE_KERNEL=auto|pingpong|cooperative` 覆盖，便于 A/B 测试。

---

## 6. API 与 Python 层差异

### 6.1 架构感知路由

| 维度 | pr-304 (SM100) | pr-323 | pr-352 | pr-360 |
|------|----------------|--------|--------|--------|
| **API 入口** | `csrc/apis/mega.hpp` 中仅注册 SM100 `fp8_fp4_mega_moe` | 独立的 `csrc/apis/sm90_mega.hpp` | 统一在 `csrc/apis/mega.hpp` 内 | `csrc/apis/mega.hpp` 内新增 SM90-only `fp8_mega_moe` 系列 |
| **Pybind11 注册** | `register_apis` 注册 `get_block_m_for_mega_moe`、`get_symm_buffer_size_for_mega_moe`、`fp8_fp4_mega_moe` | `register_sm90_apis` 单独注册 | `register_apis` 统一注册 `fp8_mega_moe` | `register_apis` 注册 auto、pingpong、cooperative 三个入口 |
| **Python 路由** | `deep_gemm/mega/__init__.py` 无 SM90 分支，直接调用 `_C.fp8_fp4_mega_moe` | `deep_gemm/mega/__init__.py` 通过 `_is_sm90()` 分发 | 同文件内通过 `arch_major == 9` 在 C++ 层路由 | Python 暴露三个 SM90 函数；auto 入口在 C++ 根据 token count 路由 |
| **架构检查** | C++ 中 `arch_major == 10` 否则 `Unsupported architecture` | Python 层区分 SM90/SM100 | C++ 统一入口内区分 SM90/SM100 | SM90-only FP8 path，独立于 SM100 FP8×FP4 path |
| **SM90 专用函数** | 无 | `get_token_alignment_for_sm90_mega_moe` | 统一为 `get_token_alignment_for_mega_moe` | `fp8_mega_moe_pingpong`、`fp8_mega_moe_cooperative`、`transform_weights_for_mega_moe_sm90` |
| **Buffer 分配** | symmetric buffer 包含 workspace、input views、L1/L2 pool、combine buffer | SM90 和 SM100 分别调用不同 C++ 函数 | 统一函数内部通过 `is_sm90` 分支处理 | 复用 MegaMoE symmetric buffer 语义，kernel 内通过 `sym_buffer.map` scatter/combine |

### 6.2 Weight Transform

- **pr-304**: `transform_weights_for_mega_moe` 是 SM100 FP4 path：L1 gate/up 按 8 粒度 interleave，然后 L1/L2 weight SF 都做 4×32 UTCCP transpose。
- **pr-323**: 有独立的 `transform_weights_for_mega_moe`（SM100 FP4 path，包含 UTCCP SF transpose）。
- **pr-352**: 有 `transform_weights_for_mega_moe`（SM100）和 `transform_weights_for_mega_moe_sm90`（SM90）。SM90 版本不需要 SF 的 4×32 UTCCP transpose，因为 SM90 没有 TMEM/UTCCP，SF 直接由 math warpgroup 从 global memory 读取。
- **pr-360**: `transform_weights_for_mega_moe_sm90` 只保留 L1 gate/up FP8 weight 的 gran-8 interleave。SM90 weight SF 保持自然 `(E, N/128, K/128)` MN-major layout，L1 SF 也保持 gate/up 沿 N 维拼接的原始布局；没有 SM100 的 4×32 UTCCP transpose。

### 6.3 测试架构

| 维度 | pr-304 (SM100) | pr-323 | pr-352 | pr-360 |
|------|----------------|--------|--------|--------|
| **测试文件** | `tests/test_mega_moe.py` (~264 lines) | `tests/test_mega_moe_hopper.py` (~1935 lines) | `tests/test_mega_moe_sm90.py` (~529 lines) | `tests/test_mega_moe_sm90.py` + `tests/bench_mega_moe_sm90.py` |
| **测试模式** | multi-process symmetric memory + bitwise baseline + benchmark | Benchmark 为主（含 fused vs baseline 对比） | 分层正确性测试为主（Smoke/Heuristic/Shape/Edge/Stress） | 分层 correctness + 强制 kernel 选择 + phase profile benchmark |
| **Baseline 对比** | DeepEP dispatch/combine + grouped FP8×FP4 GEMM + TileLang SwiGLU，要求 `torch.equal` | 完整（Triton SwiGLU + grouped GEMM + DeepEP） | 纯 PyTorch BF16/FP32 reference | PyTorch reference + DeepEP baseline；bench 支持 `--baseline-version v1/v2/both` |
| **Benchmark** | 内嵌在测试文件中，输出 TFLOPS/HBM/NVL/legacy speedup | 内嵌在测试文件中 | `tests/bench_mega_moe_sm90.py` (~259 lines) 独立 benchmark | 独立 benchmark，支持 `DG_SM90_MOE_PHASE_PROFILE` 与 `DG_SM90_MOE_KERNEL` |

### 6.4 PR304/PR357/PR360 关键 API

| API | 来源 | 用途 |
|---|---|---|
| `get_block_m_for_mega_moe` | pr-304 | 返回 SM100 MegaMoE token alignment block，原始实现固定 192 |
| `get_symm_buffer_size_for_mega_moe` | pr-304 | 计算 workspace/input/L1/L2/combine symmetric buffer 布局并返回 tensor views |
| `fp8_fp4_mega_moe` | pr-304 | 官方原始 SM100 FP8×FP4 fused MegaMoE 入口 |
| `transform_weights_for_mega_moe` | pr-304 | L1 gate/up interleave + UE8M0 SF UTCCP transpose |
| `SM100FP8FP4MegaMoESplitGraph` | pr-357 | CUDA graph 包装类，K1/K2/K3 split pipeline |
| `get_symm_buffer_for_mega_moe_split` | pr-357 | Split pipeline 专用 symmetric buffer |
| `fp8_mega_moe` | pr-360 | SM90 auto-routing 入口（token-count 选择 pingpong/cooperative） |
| `fp8_mega_moe_pingpong` | pr-360 | 强制 pingpong kernel (BLOCK_M=64) |
| `fp8_mega_moe_cooperative` | pr-360 | 强制 cooperative kernel (BLOCK_M=128) |
| `transform_weights_for_mega_moe_sm90` | pr-360 | SM90 weight transform（无 UTCCP transpose，仅 L1 gate/up interleave） |
| `DG_SM90_MOE_KERNEL` | pr-360 | 环境变量强制 kernel 选择 (auto/pingpong/cooperative) |
| `DG_SM90_MOE_COOPERATIVE_THRESHOLD` | pr-360 | cooperative 切换阈值（默认 256 tokens） |

---

## 7. 对 SM100 MegaMoE Kernel 的共享改动

pr-304 是 `sm100_fp8_fp4_mega_moe.cuh` 的官方原始实现，pr-347、pr-352 又对同一条 SM100 代码路径进行了不同程度的修改：

| 改动 | pr-304 | pr-347 | pr-352 |
|------|--------|--------|--------|
| 代码规模变化 | 原始 kernel ~1364 lines | 376 lines 重构 | 通过 API 间接影响 |
| Kernel 入口 | `sm100_fp8_fp4_mega_moe_impl` | 同源重构 | 同源保留，另加 SM90 路由 |
| Thread layout | 128 dispatch + 128 non-epilogue + 256 epilogue | 同口径 | 同口径 |
| BLOCK 配置 | `BLOCK_M=192`, `BLOCK_N=128`, `BLOCK_K=128` | 同口径 | 同口径 |
| TMEM alloc/dealloc sync fix | 原始写法，未包含后续修复 | ✓ | ✓ |
| UTCCP 对齐块尺寸 | 原始 128-aligned 写法 | 通过 `SM100ArchSpec::get_sf_uttcp_aligned_block_sizes` | 同左 |
| L2 arrival counter | `l2_arrival_mask` bit per L1 N block | 同左 | ✓ (可选/重构路径) |
| SF 量化粒度 | per-32 K packed UE8M0，UTCCP 128-token 对齐 | 同左 | 同左 |
| 调度器 cluster 支持 | 无（原始 2-CTA 硬编码） | 放宽了整除约束（`%==0` → `>0 && <=`） | 新增 `kClusterSize` 模板参数（但保留旧的严格整除约束，因分叉自 #316 早于 PR347） |

PR304 原版 SM100 kernel 的关键共享事实：
- `cute::TMEM::Allocator2Sm` 分配 2-SM tensor memory，`cute::cluster_sync()` 是 allocation/free 前后正确性的核心约束。
- FP8 input SF、L1/L2 weight SF、L2 output SF 都是 packed UE8M0 `torch.int`，并依赖 UTCCP 4×32 transpose 对齐。
- L1 output 和 L2 activation 复用同一个 `l2_acts` buffer；L1 epilogue 用 TMA store 写 FP8，L2 TMA 通过 `l2_arrival_mask` 等待对应 L1 N-block 完成。
- L2 epilogue 不先写本地最终输出，而是通过 symmetric memory map 写 remote `combine_token_buffer`，随后同一批 epilogue warps 做 top-k combine reduce。

---

---

## 8. SM90 MegaMoE 实现难点分析

基于六个 worktree 的代码 review 和实测数据，SM90/Hopper 上实现 MegaMoE 面临以下核心难点。PR304 作为 SM100 原始实现，是判断哪些问题来自架构差异的基线。

### 8.1 Register Budget 极度紧张（根本性约束）

H100/H200 每 SM 总共 **65,536 个 32-bit register**。DeepGEMM kernel 中常见的 guard 是 **64,512**，比物理 64K regs 少 1,024 regs 作为安全余量。MegaMoE 需要 Dispatch + TMA Loader + Math/Epilogue + Combine 全部塞进一个 CTA，因此 register budget 是 SM90 设计的第一约束。

| 实现 | Dispatch | TMA / Non-epi | Epilogue | 总 Reg | 占 65,536 | 占 64,512 guard |
|---|---|---|---|---|---|---|
| pr-304 SM100 fused | 128×48 | 128×40 | 256×208 | 64,512 | 98.4% | 100.0% |
| pr-323 (256-epi) | 64×48 | 64×40 | 256×168 | 48,640 | 74.2% | 75.4% |
| pr-352 compact | 64×48 | 64×48 | 256×208 | 59,392 | 90.6% | 92.1% |
| pr-360 pingpong/coop | 64×48 | 64×48 | 256×224 | 63,488 | 96.9% | 98.4% |

SM100 有 TMEM 将 accumulator 卸到 tensor memory，**SM90 没有 TMEM**——WGMMA 的 accumulator 全部留在 register 里。因此 SM90 不是简单把 PR304 的 UMMA issue warp 改成 WGMMA：math WG 和 epilogue 的 accumulator 生命周期、register reconfig、barrier 边界都要重写。PR304 自身也几乎吃满 64,512 guard，但它的压力主要来自 512-thread CTA + 256-thread epilogue，而不是 register-resident accumulator。

pr-352 的 `kWideNWarpgroups` 模式（注意：`kSerialNWarpgroups` 已被硬编码为 `false`，实际不可达）给 epilogue 配 256 reg/warp。Wide 模式要求 `kNumEpilogueWarpgroups==1`（128 epilogue threads），因此 256×128=32,768，在 budget 内安全。

### 8.2 Memory Bandwidth 是瓶颈，不是 Compute

g0033 (8×H200) 的 NCU profiling 数据一致显示：

| Kernel | DRAM % | L2 % | SM % | Occupancy |
|---|---|---|---|---|
| pr-323 fused | 89.71% | 93.68% | 18.95% | 18.57% |
| pr-352 L1 | 90.73% | 93.41% | 19.32% | — |
| pr-352 L2 | 91.42% | 93.67% | 19.71% | — |

**SM 利用率不到 20%**，但 DRAM 已接近极限。这意味着任何 kernel 优化（更好的 WGMMA 调度、更聪明的 register 分配）如果不减少 HBM 流量，就不会提升吞吐。

PR304 在 SM100 上通过 FP4 weights 直接把 weight bytes 降为 FP8 的一半，并用 2SM TMA/UMMA pipeline 把 NVLink dispatch、tensor core 计算和 epilogue 尽量重叠。SM90 版本失去 FP4/TMEM 后，pr-360 的 cooperative kernel 正是为了解决这个痛点：两个 math WG 共享 B/weight TMA load，**weight HBM 读取量减半**。这是目前唯一直接攻击 SM90 memory bottleneck 的设计。

### 8.3 Occupancy 被 Register + SMEM 双重限制

H100 上有 64 warps/SM 的理论上限，但 MegaMoE 的每个 CTA 就占 **12-16 个 warp**：

| 实现 | CTA warps | Max CTA/SM | Max warps/SM | Occupancy |
|---|---|---|---|---|
| pr-304 SM100 fused | 16 per CTA, 2-CTA cluster | 1 CTA/SM（cluster 占 2 SM） | 16/64 | 25%（参考，非 SM90） |
| pr-323 256-epi | 12 | 1 | 12/64 | 18.75% |
| pr-352 compact | 12 | 1 | 12/64 | 18.75% |
| pr-360 | 12 | 1 | 12/64 | 18.75% |
| pr-352 wide | 16 | 1 | 16/64 | 25% |

每个 SM 只能跑 **1 个 CTA**，意味着 latency hiding 极度依赖单 CTA 内部的 warp 级并行和 TMA async copy。18.75% theoretical occupancy 只能说明可驻留 warp 数少，不能直接推出具体等待时间比例；实际 stall breakdown 仍需要 scheduler/memory stall metrics 或 timeline 佐证。

### 8.4 Scale Factor 格式差异导致整条 Pipeline 不一致

| 维度 | PR304 SM100 | SM90 |
|---|---|---|
| Scale 类型 | UE8M0 int (8-bit) | float (32-bit) |
| Granularity | per-32 K，4 个 UE8M0 packed 到 `torch.int` | per-128 K |
| Layout | M-major / MN-major，并对 128 token 做 UTCCP 4×32 transpose | SM90 kernel/test 需要各自约定 float SF layout |
| SF 加载 | TMA + UTCCP copy 到 TMEM，UMMA block scaling 消费 | SMEM ld_shared，或 math WG 直接 ldg |
| SF 存储 | int packed，TMA-aligned stride | MN-major float，需保证 kernel 索引一致 |

pr-360 review 中发现的**中间 FP8 scale 语义不一致**（kernel 用 2 的幂次，test 用 exact float）正是这个问题的表现——SM90 的 float scale 使"应该用什么 scale 策略"变成了一个需要跨 kernel/test/baseline 协商的设计决策，而 SM100 的 UE8M0 天然是 2 的幂次，不存在歧义。

### 8.5 三种 Warpgroup 协作策略各有取舍

SM90 上无法像 SM100 那样用 2-CTA cluster UMMA 自然划分 work，必须在单 CTA 内用 warpgroup 拆分：

| 策略 | 代表 PR | 优点 | 缺点 |
|---|---|---|---|
| 2-CTA UMMA/TMEM fused | pr-304 | accumulator 不占普通寄存器，天然支持 FP8×FP4 + UE8M0 block scaling | SM100-only，依赖 TMEM/UMMA/2SM TMA |
| Fused single WG | pr-323 | 简单 | 无 pipeline overlap |
| Split phase (不同 CTA 执行 L1/L2) | pr-352 | L1/L2 独立配置 | kernel launch overhead，无法 overlap |
| Pingpong (WG 交替) | pr-360 | MMA/epilogue overlap | 仅对小 M 有效，BLOCK_M 固定 64 |
| Cooperative (WG 协作) | pr-360 | 共享 weight load，HBM 减半 | 需要 256-thread barrier 防 SMEM 别名冲突 |

没有一种策略在所有 token 量下最优——这也是 pr-360 引入 `DG_SM90_MOE_COOPERATIVE_THRESHOLD` auto-routing 的原因。

### 8.6 Dispatch+TMA 共享 Warpgroup 的设计约束

PR304/SM100 上 dispatch 独占 128 threads（完整 warpgroup），non-epilogue 也占 128 threads：warp4/5 做 A/B TMA，warp6 issue UMMA，warp7 仅做 register dealloc 后保留。SM90 上为了节省 register/SMEM，pr-352/360 把 dispatch 和 TMA 合并到同一个 WG0（各 64 threads，共 128）：

```
WG0: [dispatch warp 0] [dispatch warp 1] [TMA A loader] [TMA B loader]
```

但这带来两个约束：
- Dispatch 和 TMA 的 register dealloc 必须使用相同的 count（都是 48），不能各自独立调优
- Dispatch 的 register spilling 风险升高——dispatch warp 需要 rank round-robin 选择、SF copy、expert counting 等密集操作，48 reg/warp 偏紧

### 8.7 NCU Profiling 在 Distributed Symmetric-Memory 场景下不可靠

g0033 上多次尝试 NCU kernel replay / application replay / range replay，`.ncu-rep` 文件生成但 metrics 均为 `nan`，报 `LaunchFailed` / `unspecified launch failure`。这是因为 MegaMoE 的 `nvlink_barrier` + `sym_buffer.map` 跨 rank 同步与 NCU 的 kernel replay 机制冲突——replay 只重放单个 rank 的 kernel，跨 rank 的 symmetric memory 操作无法正确重放。

**这意味着 SM90 MegaMoE 的性能调试无法依赖标准的 NCU 单 kernel profiling 流程**，必须用 nsys timeline 或手动插入 NVTX marker 替代。

### 8.8 代码复杂度爆炸

| 实现 | Kernel 行数 | Heuristics 行数 | 配置参数数量 |
|---|---|---|---|
| pr-304 | 1,364 | 211 | ~10（固定 block + wave/stage 动态） |
| pr-323 | 1,935 | 180 | ~5 |
| pr-352 | 2,507 | 1,025 | ~15（多 bool flags + BLOCK_M/N 组合） |
| pr-360 | 1,418+1,419 | 320 | ~8 |

PR304 代码行数较少，是因为它把复杂性交给 SM100 TMEM/UMMA/UTCCP 专用硬件语义；SM90 无法复用这些语义，复杂度转移到 warpgroup 协作、SMEM staging 和 heuristics。pr-352 用 C++ 宏生成 4 种 phase 策略的 kernel 变体，每个变体有独立的 `__launch_bounds__` 和 register 分配，JIT 编译时间显著增长。pr-360 用两套独立 kernel（pingpong + cooperative）替代宏方案，降低了单 kernel 复杂度但增加了维护两个 ~1400 行 kernel 的负担。

### 8.9 难点优先级总览

| 优先级 | 难点 | 影响范围 | 是否有有效对策 |
|---|---|---|---|
| **P0** | Register budget 极度紧张 | 所有 SM90 实现；PR304 说明 SM100 也接近 guard | pr-360 的 compact dispatch+TMA WG 共享是局部缓解 |
| **P0** | HBM bandwidth bottleneck（~90% DRAM） | 所有 SM90 实现 | pr-360 cooperative 共享 B-load 是有意义的攻击方向 |
| **P1** | Occupancy ~18%，latency hiding 困难 | 所有 SM90 实现 | 单 CTA 内 TMA async + warp 级并行是仅有的手段 |
| **P1** | NCU profiling 不可靠 | 调试/优化流程 | nsys timeline + NVTX 替代 |
| **P2** | Float SF 语义不一致 | 测试/baseline 对比 | 需要统一策略决策 |
| **P2** | 多 warpgroup 策略选择 | pr-352/360 | auto-routing (pr-360) 是有效解决方案 |
| **P3** | 代码复杂度 | 维护成本 | pr-360 的双 kernel 比 pr-352 的宏方案更可维护 |

---

## 9. 技术亮点总结

### pr-304 原始 SM100 fused kernel 的基线价值

1. **完整端到端 fused pipeline**：单 kernel 覆盖 dispatch、L1 FP8×FP4、SwiGLU、L2 FP8×FP4、remote scatter 和 combine reduce。
2. **TMEM/UMMA accumulator pipeline**：`tcgen05` UMMA 写 TMEM，epilogue 从 TMEM load，避免 SM90 式 register accumulator 长生命周期。
3. **2-CTA cluster 语义**：通过 `Allocator2Sm`、2SM TMA load 和 cluster barrier 构建 Blackwell 原生的跨 CTA GEMM tile。
4. **UE8M0/UTCCP scale path**：activation/weight/intermediate scale 都按 packed UE8M0 和 4×32 UTCCP transpose 组织，baseline correctness 可以做到与 legacy path `torch.equal`。
5. **Symmetric-memory workspace 设计**：同一 workspace 管理 expert counts、arrival counters、source metadata、L1/L2 pools 和 combine buffer，是后续 SM90/SM100 split 路线共享的结构参照。

### pr-352 相对于 pr-323 的主要优势

1. **Split kernel 架构**：L1/L2 独立编译，各自最优配置（register/threads/cluster），减少不必要的 dispatch/combine 逻辑在不需要的 kernel 中。
2. **mma.sync 支持**：针对极小 token-per-expert（≤16）场景的专门优化路径。
3. **Profile 驱动的启发式**：针对 H20（≤80 SMs）和 H200（≥100 SMs）两种硬件配置的 per-bucket 最佳实践调优。
4. **Direct L2 scatter**：L2 输出直接通过 NVLink scatter 到目标 buffer，省去 SMEM staging 开销。
5. **Async L1 TMA store**：L1 epilogue 的 FP8 输出通过 double-buffered TMA store 实现异步写入。
6. **Cluster B-cast**：支持 cluster_size=2 的 TMA multicast，对 B 矩阵的相同 tile 进行广播。
7. **N-major scheduling**：在 token-per-expert 较大时，先跨 N block 再跨 M block 分配，提高 L2 cache 利用率。

### pr-360 的主要技术亮点

1. **双 fused kernel 策略**：保留 single-launch fused pipeline，但把小/中 M 与大 M 分别交给 pingpong 和 cooperative 两套 kernel，避免单个 SM90 kernel 同时承担互相冲突的优化目标。
2. **Pingpong warpgroup overlap**：`BLOCK_M=64` 下两个 math WG 处理不同 tile，通过 `OrderedSequenceBarrier<2,2>` 交替 MMA 和 epilogue，降低小 M 的 epilogue 暴露延迟。
3. **Cooperative weight sharing**：`BLOCK_M=128` 下两个 math WG 按 M 切分同一 tile，共享 B/weight TMA load，直接针对 SM90 MegaMoE 的 HBM bandwidth bottleneck。
4. **紧凑线程布局**：64 dispatch + 64 TMA/non-epilogue + 256 math/epilogue，384 threads 内完成 dispatch、TMA、WGMMA、L1 epilogue、L2 scatter 和 combine。
5. **SM90 专用 scale/layout 简化**：weight SF 保持自然 MN-major float layout，不做 SM100 UTCCP transpose；L1 gate/up 只 interleave FP8 weight 本体。
6. **可控 A/B 入口**：`DG_SM90_MOE_KERNEL`、`DG_SM90_MOE_COOPERATIVE_THRESHOLD`、`DG_SM90_MOE_NMAJOR` 让 kernel 选择和 L2 scheduling 可以在 benchmark 中显式覆盖。

### pr-323 相对于 pr-352 的潜在优势

1. **更简洁**：代码量少 \~30%（\~1935 vs \~2507 kernel lines），配置逻辑更直观。
2. **单 kernel launch**：减少了 kernel launch overhead，在 token-per-expert 较小的场景可能更优。
3. **支持 Split-MN warpgroup**：4-warpgroup 配置（BLOCK_M=128, BLOCK_N=256），在 prefill 场景可能有优势。
4. **shared-SF 机制**：更精确的 SF 量化粒度控制。

---

## 10. 当前状态评估

### 10.1 成熟度评估

| 维度 | pr-304 | pr-323 | pr-352 | pr-357 | pr-360 |
|------|--------|--------|--------|--------|--------|
| **目标架构** | SM100 | SM90 | SM90 | SM100 (CUDA 13.1+) | SM90 |
| **功能完整性** | 完整原始 fused kernel | 完整（decode + prefill） | 完整（decode + prefill） | 完整（3-kernel split pipeline） | 完整（pingpong + cooperative） |
| **硬件适配** | Blackwell TMEM/UMMA/FP4 | 通用 SM90 | H20/H200 专项调优 | SM100 green context | 通用 SM90 |
| **代码质量** | 原始实现紧凑，但硬件专用语义重 | 简洁，可读性好 | 功能丰富，复杂度高 | 结构清晰，graph 构造较复杂 | 两套 kernel 结构一致，设计精良 |
| **测试覆盖** | legacy baseline bitwise + benchmark | Benchmark 为主 | 分层正确性 + benchmark | `torch.equal` bitwise fused 对比 | 5 层分层测试 + benchmark |
| **与 SM100 兼容性** | 原生 SM100 | 独立 API | 统一 API | SM100 only | 独立 API（SM90 only） |
| **来源** | 官方 | 官方 | 官方 | RayWang96 (社区) | mpdfdfl (社区) |

### 10.2 结论

当前 DeepGEMM 社区中 MegaMoE 存在多条平行的实现路线：

**SM90 路线（三条）:**
- **pr-323** 代表 **fused kernel + 简化配置** 路线，代码简洁，适合快速验证。
- **pr-352** 代表 **split kernel + profile 驱动** 路线，功能最完整，对 H20/H200 有专项优化。
- **pr-360** 代表 **warpgroup 协作优化** 路线，通过 pingpong（MMA/epilogue 重叠）和 cooperative（共享 weight load）减少延迟和 HBM 压力。

**SM100 路线（两条）:**
- **pr-304** 代表 **官方原始 fused FP8×FP4 MegaMoE** 路线，依赖 Blackwell TMEM/UMMA/FP4/UE8M0，是所有移植工作的功能基线。
- **pr-357** 代表 **green-context split pipeline** 路线，利用 CUDA 13.1 green context 实现 K1/K2 真正并发，SM100 exclusive。

pr-347 提供了两条路线共享的基础设施重构（调度器 cluster 支持、barrier 优化等）。建议后续工作将两条路线的优势合并：
1. 以 pr-304 的 API/scale/workspace 语义为 correctness baseline，明确 SM90 FP8-only path 与 SM100 FP8×FP4 path 的差异边界
2. 吸收 pr-352 的 profile 驱动配置选择器到 pr-323/pr-360 的 fused kernel 中
3. 吸收 pr-323 的 Split-MN warpgroup 支持到 pr-352 中
4. 统一到 pr-352 的 unified API 设计，并保留 PR304/PR357 的 SM100-only fast path

---

## 11. 实测性能对比

测试环境：2× NVIDIA H100 80GB HBM3, NVCC 12.4（系统默认） + CUDA 12.8 toolkit, PyTorch 2.11.0+cu128, EP=2（2 ranks, 每 rank 192 experts）

PR304 未纳入本节实测表，因为它是 SM100-only kernel，依赖 `__CUDA_ARCH__ >= 1000`、TMEM/UMMA、FP4 weight path 和 2-CTA `Allocator2Sm`。在 H100/H200 上无法直接编译或运行该 kernel；这里的性能对比只覆盖 SM90 移植路线。未来若在 Blackwell 上补测，PR304 的官方测试入口是 `tests/test_mega_moe.py`，输出 `mega_moe` fused latency、TFLOPS、HBM GB/s、NVLink GB/s 以及 legacy baseline speedup。

### 11.1 Token Sweep 延迟对比

| Tokens | Recv Tokens | pr-323 (fused) | pr-352 (split) | Δ pr-352 vs pr-323 |
|--------|-------------|----------------|----------------|---------------------|
| 1      | 9           | 271 μs         | 264 μs         | -2.6%              |
| 2      | 13          | 350 μs         | 354 μs         | +1.1%              |
| 4      | 26          | 628 μs         | 621 μs         | -1.1%              |
| 8      | 46          | 992 μs         | 999 μs         | +0.7%              |
| 16     | 85          | 1773 μs        | 1769 μs        | -0.2%              |
| 32     | 194         | 2808 μs        | 2771 μs        | -1.3%              |
| 64     | 390         | 3939 μs        | 3877 μs        | -1.6%              |
| 128    | 755         | 4357 μs        | 4276 μs        | **-1.9%**          |
| 256    | 1531        | 4426 μs        | 4385 μs        | -0.9%              |
| 512    | 3079        | 4503 μs        | 4430 μs        | **-1.6%**          |

**结论**：pr-352（split kernel）在 4+ tokens 以上有 1-2% 的延迟优势，但在极端 decode（1-2 tokens）下两者基本持平。

### 11.2 Shape Sweep 延迟对比（128 tokens）

| Hidden × Intermediate | pr-323 (fused) | pr-352 (split) | Δ  |
|-----------------------|----------------|----------------|----|
| 7168 × 2048           | 2888 μs (25 TF) | 2911 μs (24 TF) | +0.8% |
| 7168 × 3072           | 4351 μs (23 TF) | 4266 μs (23 TF) | -2.0% |
| 4096 × 2048           | 1674 μs (22 TF) | 1694 μs (22 TF) | +1.2% |
| 5120 × 2048           | 2037 μs (24 TF) | 2064 μs (24 TF) | +1.3% |
| 7168 × 4096           | 5794 μs (24 TF) | 5696 μs (24 TF) | **-1.7%** |

**结论**：pr-352 在大 intermediate_hidden（3072, 4096）上有优势，在小模型上两者持平。

### 11.3 NCU Profiling 详细对比（32 tokens）

#### pr-323 (fused kernel: `sm90_fp8_mega_moe_impl`)
- **Duration**: 3.42 ms (132 blocks × 384 threads)
- **DRAM throughput**: 89.71%（极强 memory-bound）
- **L2 Cache throughput**: 93.68%（近饱和）
- **L1/TEX throughput**: 17.86%
- **Compute (SM) throughput**: 18.95%
- **Theoretical occupancy**: 18.75%（12/64 warps per SM）
- **Achieved occupancy**: 18.57%（11.89 warps）
- **Bottleneck**: **DRAM bandwidth + Register/Shared Memory**（Block Limit Registers=1, Shared Mem=1）

#### pr-352 L1 kernel (`sm90_fp8_mega_moe_l1_impl`)
- **Duration**: 2.23 ms
- **DRAM throughput**: 90.73%
- **L2 Cache throughput**: 93.41%
- **Compute (SM) throughput**: 19.32%

#### pr-352 L2 kernel (`sm90_fp8_mega_moe_l2_impl`)
- **Duration**: 1.10 ms
- **DRAM throughput**: 91.42%
- **L2 Cache throughput**: 93.67%
- **Compute (SM) throughput**: 19.71%

#### NCU 关键洞察

| 指标 | pr-323 fused | pr-352 L1+L2 | 解读 |
|------|-------------|-------------|------|
| Kernel 总时间 | 3.42 ms | 2.23 + 1.10 = 3.33 ms | pr-352 kernel 部分快 2.6% |
| Memory-bound 程度 | ~90% DRAM | ~91% DRAM | 两者都极度受限于 HBM 带宽 |
| L2 压力 | 93.68% | 93.41%/93.67% | L2 cache 都是瓶颈 |
| SM 利用率 | 18.95% | 19.32%/19.71% | 都极低，compute 远未饱和 |
| Occupancy 限制 | Registers + SMEM | Registers + SMEM | 两者 block 均被寄存器/共享内存限制到 1 block/SM |

**核心结论**：SM90 MegaMoE 在 H100 上的性能瓶颈是 **HBM 内存带宽**（~90% DRAM utilization），而非计算能力。register 和 shared memory 限制了 occupancy（仅 ~18.75%），进一步加剧了 memory latency hiding 的困难。

pr-352 的 split 架构优势在于：
1. L1/L2 独立 kernel 允许不同的线程配置和 register budget
2. L1 (2.23ms) 和 L2 (1.10ms) 的 work 量不同（L1=2× 算力），分拆后各阶段可以独立优化
3. 避免 fused kernel 中 L1→L2 切换时的 pipeline 重建开销
4. L2 kernel 不需要 dispatch 逻辑，减少了不必要的 SMEM 占用和 barrier

但 pr-323 的 fused 架构在极端 decode（1-2 tokens）场景表现相当，且单次 kernel launch 简化了上层调度。

### 11.4 总体评估

| 维度 | pr-323 fused | pr-352 split |
|------|-------------|-------------|
| Latency (4-512 tokens) | baseline | **-1~2%** |
| Latency (1-2 tokens) | **持平** | 持平 |
| Kernel 架构复杂度 | 低（single kernel） | 高（dual kernel） |
| Config 灵活性 | 低（2 block_m 选择） | 高（6 block_m + 多 bool flags） |
| H20/H200 适配 | 无 | 有（设备 profile 自动检测） |
| NCU 验证 | ✓ | ✓（L1+L2 分别剖面） |
| 代码量 | 1935 行 kernel | 2507 行 kernel + 1025 行 heuristics |
| API 设计 | 独立 SM90 入口 | 统一 SM100/SM90 入口 |

### 11.5 PR360 8×H200 历史补测摘要

PR360 另有一次单节点 8×H200 历史补测，目的是验证 PR 描述中的 fused kernel 性能是否能复现。该数据来自本地已记录的 code review 报告与远端日志路径；本次报告更新没有重新访问远端机器，也没有重新运行 benchmark。

测试条件摘要：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`。运行时设置 `NVSHMEM_IBGDA_ENABLE=0 NVSHMEM_DISABLE_IBGDA=1 EP_DISABLE_GIN=1 NCCL_NVLS_ENABLE=0`。

关键限制：`NCCL_NVLS_ENABLE=0` 是当时为了避开 Torch 2.12/NCCL 2.29 的 NVLS multicast bind 错误；PyTorch symmetric memory multicast 也走 fallback。因此这些数据能验证 PR360 fused kernel 的主延迟趋势，但不能证明 CUDA multicast/NVLS 路径可用。DeepEP v2 baseline 未验证，因为当前容器内 DeepEP 缺少 `deep_ep.ElasticBuffer`。

| Shape | Tokens 覆盖 | fused 相对 PR 描述偏差 | 结论 |
|-------|-------------|------------------------|------|
| DeepSeek-V4 Flash (`h=4096, i=2048, E=256, topk=6`) | 16–8192 | `-3.0% ~ +0.5%` | 复现通过 |
| DeepSeek-V4 Pro (`h=7168, i=3072, E=384, topk=6`) | 16–8192 | `-6.1% ~ -0.3%` | 趋势一致，但不能确认 v2 对比 |
| MiMo-V2.5 (`h=4096, i=2048, E=256, topk=8`) | 16–8192 | `-2.6% ~ +0.7%` | 复现度较高 |
| MiMo-V2.5-Pro (`h=6144, i=2048, E=384, topk=8`) | 16–8192 | `-3.7% ~ -1.1%` | 复现度较高 |

日志路径：
- `/mnt/beegfs/lxy/logs/pr360_h200_deepseek_flash_h4096_i2048_e256_topk6_baseline_v1_torch212.log`
- `/mnt/beegfs/lxy/logs/pr360_h200_deepseek_pro_h7168_i3072_e384_topk6_baseline_v1_torch212.log`
- `/mnt/beegfs/lxy/logs/pr360_h200_mimo_v25_h4096_i2048_e256_topk8_baseline_both_torch212.log`
- `/mnt/beegfs/lxy/logs/pr360_h200_mimo_v25_pro_h6144_i2048_e384_topk8_baseline_v1_torch212.log`

综合判断：PR360 fused kernel 的主性能结论基本成立，四个公开 shape 的 fused 时间均能在 8×H200 上以几个百分点内的偏差复现。但 baseline 数值对 DeepEP/TileLang/PyTorch 版本更敏感；如果要把“相对 DeepEP v2”的 speedup 作为合入依据，需要在包含 `deep_ep.ElasticBuffer` 的 DeepEP 版本上重跑。

---

## 12. PR357: SM100 Green-Context Split-Kernel MegaMoE

### 12.1 概述

pr-357 (`RayWang96/DeepGEMM:split_mega_moe`, commit `bb837421b`) 是一个 **SM100 only** 的 green-context split-kernel 实现。它以 PR304 的 SM100 FP8×FP4 fused kernel 为 correctness baseline，但不是 SM90 路径，也不是简单的 "把 fused kernel 拆成三个 kernel"。核心技术利用 CUDA Runtime 13.1 的 green context API，将 dispatch/L1/SwiGLU（K1）和 L2/combine（K2）放到互斥 SM partition 中**真正并发执行**。

来源: `git@github.com:RayWang96/DeepGEMM.git`, 分支 `split_mega_moe`

### 12.2 三段 Kernel 架构

| Kernel | 入口 | Grid | Threads | L1 Warp Layout | SM 分配 |
|---|---|---|---|---|---|
| K1: `dispatch_l1_swiglu` | `sm100_fp8_fp4_mega_moe_split_dispatch_l1_swiglu_impl` | `kernel1_sms` (默认 96) | 512 (dispatch=128, TMA=128, epilogue=256) | 4+4+8=16 warps, 4 WGs | Green Context 0 |
| K2: `l2_combine` | `sm100_fp8_fp4_mega_moe_split_l2_combine_impl` | `kernel2_sms` (默认 52) | 384 (TMA=128, epilogue=256) | 4+8=12 warps, 3 WGs (无 dispatch) | Green Context 1 |
| K3: `combine_reduce` | `sm100_fp8_fp4_mega_moe_split_combine_reduce_impl` | `num_tokens` | 512 (16 warps) | 1 CTA/token | Primary Context |

### 12.3 Green Context 并发模型

```
SM partition:
  ┌─────────────────────────────────────────────────────┐
  │ Green Context 0 (kernel1_sms=96 SMs)               │
  │   K1: dispatch → L1 GEMM → SwiGLU → write mask     │
  │        ↓ l2_arrival_mask (per pool block)           │
  │ Green Context 1 (kernel2_sms=52 SMs)               │
  │   K2: wait mask → L2 GEMM → NVLink scatter         │
  │        ↓ NVLink barrier                            │
  │ Primary Context (all SMs)                           │
  │   K3: combine reduce                                │
  └─────────────────────────────────────────────────────┘
```

关键设计点：
- K1/K2 仍保留 PR304 的 SM100 FP8×FP4、TMEM/UMMA、UE8M0 SF 和 512/384-thread thread layout，只是把 fused state machine 拆成 producer/consumer kernels
- K1/K2 通过 workspace `l2_arrival_mask` 实现 producer/consumer 流水 — K2 逐 block 消费 K1 的 L1 output
- K1 使用 route-based dispatch — `SplitWorkspace` 保存 route metadata，避免重复 remote token pull
- K3 依赖 CUDA graph node dependency（等 K1/K2 都完成）和 K2 的 NVLink barrier
- correctness 目标 bitwise identical 于 fused kernel（`torch.equal`）

### 12.4 与 SM90 路径的关系

| 维度 | pr-304 fused | pr-357 split | pr-323/352 pr-360 |
|---|---|---|---|
| 目标架构 | SM100 only | SM100 only (CUDA 13.1+) | SM90/Hopper |
| 精度 | FP8/FP4 (TMEM) | FP8/FP4 (TMEM) | FP8 only (WGMMA register accumulator) |
| 并发策略 | 单 kernel fused state machine | Green Context SM partition | Single launch / sequential launch |
| Kernel 数 | 1 | 3 (CUDA graph) | 1-2 |
| L1→L2 同步 | same-kernel `l2_arrival_mask` | cross-kernel `l2_arrival_mask` | CTA order 或 cross-kernel mask |
| 分支内旁路改动 | 无 | FP16 weights MQA logits | 无 |

---

## 13. PR360: SM90 Pingpong/Cooperative MegaMoE

### 13.1 概述

pr-360 (`mpdfdfl/DeepGEMM:sm90-mega-moe-pingpong-coop`, commit `f27fd5611`) 是 **第三套独立的 SM90 MegaMoE 实现**。它可以看作对 PR304 fused pipeline 的 Hopper 重构：保留 dispatch→L1→L2→combine 的单 kernel 目标，但用 WGMMA register accumulator、SMEM staging 和多 warpgroup 协作替代 SM100 的 TMEM/UMMA/2CTA 机制。它引入两种 math warpgroup 协作策略，在单一 fused kernel 内部实现 pipeline 重叠或 weight 共享。

来源: `git@github.com:mpdfdfl/DeepGEMM.git`, 分支 `sm90-mega-moe-pingpong-coop`

### 13.2 两种 Kernel 架构

#### Pingpong (BLOCK_M=64)

| Warp Index | 所属 WG | Role |
|---|---|---|
| 0–1 (2 warps) | WG0 | Dispatch |
| 2–3 (2 warps) | WG0 | TMA A+SFA / B loader |
| 4–7 (4 warps) | WG1 | Math WG0: 处理 tile 0,2,4... |
| 8–11 (4 warps) | WG2 | Math WG1: 处理 tile 1,3,5... |

**12 warps = 384 threads = 3 warpgroups. Register: 48+48+224+224 = 63,488 regs（占 65,536 的 96.9%，占 64,512 guard 的 98.4%）**

核心机制：`OrderedSequenceBarrier<2,2>` 实现两个 math WG 的交替执行：
```
WG0: [MMA tile 0] [epi tile 0] [MMA tile 2] [epi tile 2] ...
WG1:              [MMA tile 1] [epi tile 1] [MMA tile 3] ...
```
一个 WG 的 MMA 与另一个 WG 的 epilogue 重叠，隐藏 epilogue 延迟。

#### Cooperative (BLOCK_M=128)

| Warp Index | 所属 WG | Role |
|---|---|---|
| 0–3 (4 warps, 同 pingpong) | WG0 | Dispatch + TMA |
| 4–7 (4 warps) | WG1 | Math WG0: 处理 rows [0, 64) |
| 8–11 (4 warps) | WG2 | Math WG1: 处理 rows [64, 128) |

**12 warps = 384 threads = 3 warpgroups. 同 register 配置。**

核心机制：两个 math WG **协作处理同一 tile**，各负责 64 行 M-split，**共享一次 B/weight TMA load**（weight HBM 读取减半）。L2 epilogue 使用 256-thread barrier 防 smem_cd 别名冲突。

### 13.3 Token-Count Auto Routing

| 条件 | 选择 kernel | 适用场景 |
|---|---|---|
| `num_tokens < threshold` (默认 256) | pingpong | 小 M，epilogue 延迟主导 |
| `num_tokens >= threshold` | cooperative | 大 M，weight 重复读取主导 |

可通过 `DG_SM90_MOE_KERNEL=pingpong|cooperative` 强制选择。阈值可通过 `DG_SM90_MOE_COOPERATIVE_THRESHOLD` 覆盖。

### 13.4 与 pr-304/pr-323/pr-352 的架构对比

| 维度 | pr-304 SM100 fused | pr-323 fused | pr-352 split | pr-360 pingpong/coop |
|---|---|---|---|---|
| Kernel 数 | 1 | 1 | 2 (L1+L2) | 1 (auto-routing) |
| Accumulator | TMEM | registers | registers / SMEM staging | registers |
| Warpgroup 策略 | 2-CTA UMMA issue + TMEM epilogue | Split-N | Serial-N / Wide-N / Fused | Pingpong (交替) / Cooperative (协作) |
| Weight load 优化 | FP4 weight bytes + 2SM TMA path | 无 | cluster B/weight multicast in selected configs | Cooperative: 共享 B-load，HBM 减半 |
| Dispatch+TMA | 独立 dispatch WG + non-epilogue WG | 独立 WG 或共享 | 可选 compact | 固定 compact (共享 WG0) |
| Threads | 512 per CTA, 2-CTA cluster | 256/384/768 | 384/512 | 384 |
| N-major L2 schedule | scheduler 按 L1/L2 N blocks 遍历，要求偶数配对 | 无 | 支持 | 支持 (DG_SM90_MOE_NMAJOR) |

### 13.5 代码review主要发现

1. **波次启发式可能违反整除约束** — scheduler 要求 `kNumExpertsPerWave` 整除 `kNumExpertsPerRank`，但启发式不限制为除数
2. **调度器波次边界越界读取** — `get_num_tokens(wave_end)` 可能读取未初始化寄存器
3. **中间 FP8 scale 语义不一致** — kernel 使用 2 的幂次 scale，test/baseline 使用 exact float
4. **Phase-profile benchmark 路径不可用** — API 检查与 kernel 实现不匹配
5. **Weight scale layout 检查过宽** — 接受 kernel 无法正确索引的转置 stride

---

## 14. 总结与展望

### 14.1 SM90 MegaMoE 三条路线

| 路线 | 来源 | 策略 | 状态 |
|---|---|---|---|
| pr-323 | 官方 | Fused kernel, 简单可靠 | 早期实现 |
| pr-352 | 官方 | Split kernel, phase 分离, 多 config | 增强版 |
| pr-360 | mpdfdfl | Pingpong/cooperative fused, token routing, weight-load sharing | 社区 PR；8×H200 fused 延迟已历史复现 |

### 14.2 SM100 路线

| 路线 | 来源 | 策略 | 状态 |
|---|---|---|---|
| pr-304 | 官方 | 原始 fused FP8×FP4, TMEM/UMMA, 2-CTA cluster | 官方基线 |
| pr-357 | RayWang96 | Green-context split, K1/K2 真正并发 | 社区 PR |

pr-304 是功能和 layout 基线，pr-357 是在同一 SM100 能力集上的并发化探索。PR357 的技术方向（green-context concurrency）是 SM100 专属能力，无法直接移植到 SM90。但它证明了三段 split + CUDA graph + green context 的可行性，对于未来 SM100 上的大规模 MoE 部署有参考价值。SM90 后续工作应继续以 PR304 的 workspace/scale/correctness 语义为参照，同时承认 TMEM/UMMA/FP4 相关优化无法直接搬到 Hopper。

### 14.3 后续工作

1. **PR360 correctness/profile 补齐**：在当前本地可用环境中补跑 `tests/test_mega_moe_sm90.py` 和 `tests/bench_mega_moe_sm90.py`，分别覆盖 `DG_SM90_MOE_KERNEL=pingpong|cooperative|auto` 与 `DG_SM90_MOE_PHASE_PROFILE=1`。
2. **DeepEP v2 baseline 重测**：使用包含 `deep_ep.ElasticBuffer` 的 DeepEP 版本重跑 PR360 `--baseline-version v2` 或 `both`，确认相对 v2 的真实 speedup/slowdown。
3. **NVLS/multicast 路径验证**：在 PyTorch/NCCL/symmetric-memory 组合能正确启用 multicast 的环境下重测 PR352/PR360，区分 fallback 性能与 NVLS 性能。
4. **SM90 route 合并方向**：短期可把 pr-352 的 profile-driven device profile 与 PR360 的 pingpong/cooperative auto-routing 结合；长期需要统一 API、buffer、scale 语义，避免三套 SM90 实现继续分叉。
5. **SM100/SM90 边界固定**：PR304/PR357 的 FP8×FP4、TMEM、UE8M0/UTCCP 能力应保持 SM100-only fast path；SM90 路线应独立维护 FP8×FP8 float scale 语义，而不是尝试机械移植 SM100 layout。
