# MegaMoE SM90 Implementation Design (Final)

## 1. Executive Summary

本文档描述 DeepGEMM MegaMoE 在 SM90 (Hopper: H100/H200/H20) 架构上的最终实现设计。设计基于对 SM100 fused kernel (`sm100_fp8_fp4_mega_moe.cuh`, PR304) 的详细分析，以及对三条现有 SM90 探索路线 (PR323 fused, PR352 split, PR360 pingpong/cooperative) 的横向对比研究。

### 1.1 核心设计决策

| # | 决策项 | 最终选择 | 理由 |
|---|--------|---------|------|
| 1 | Kernel 架构 | **Fused single kernel** | 单 kernel launch 减少 overhead，代码简洁 (~1900 行) |
| 2 | Warpgroup 策略 | **Cooperative (BM=128) + single-WG (BM≤64)** | Cooperative 共享 weight load 主攻 HBM bottleneck；小 M 退化为 single-WG 避免 compute 浪费 |
| 3 | BLOCK_M | **动态 128/64/32** | 按 expected_tokens_per_expert 三档选择，覆盖 decode 到 prefill |
| 4 | BLOCK_N | **固定 128** | 64 个 accumulator regs/WG，寄存器压力可控 |
| 5 | Cluster | **固定 1** | Cooperative 已共享 B-load；cluster=2 的 amax 跨 CTA 同步过于复杂 |
| 6 | 线程布局 | **384 = 64 dispatch + 64 TMA + 256 math/epilogue** | 与 SM100 的 512 相比省出 WG 寄存器预算 |
| 7 | SF 策略 | **Exact float SF** | 与 PyTorch 标准 FP8 quantize 一致 |
| 8 | Weight SF 加载 | **Math WG `__ldg`** | 无额外 SMEM 占用，PR323/360 已验证 |
| 9 | L2 Act SF | **Per-64 dual-half SF** | 精度更高，PR323/360 采用 |
| 10 | L2 Scheduling | **Auto N-major** | 大 M 下减少 weight L2 cache thrash |
| 11 | Math WG Regs | **224 reg/warp** | 总占用 96.9%，PR360 已验证可行 |
| 12 | L1 TMA Store | **Double-buffered** | TMA store 与 epilogue 计算重叠 |
| 13 | API | **独立 SM90 入口** | Python 层显式区分 SM90/SM100 |
| 14 | Combine | **复用 SM100 逻辑** | Chunked TMA combine 与 TMEM/UMMA 无关，可直接移植 |
| 15 | 文件组织 | **单文件 `sm90_fp8_mega_moe.cuh`** | 模板参数区分 cooperative/single-WG |
| 16 | Clamp + FastMath | **两者都支持** | 模板参数，零开销 |

---

## 2. SM100 vs SM90 Architecture Gap

### 2.1 硬件能力对比

| 能力 | SM100 (Blackwell) | SM90 (Hopper) | 影响 |
|------|-------------------|---------------|------|
| Tensor Memory (TMEM) | 有 | **无** | accumulator 必须驻留在 registers |
| UMMA (Unified MMA) | 有（`tcgen05`） | **无** | 使用 WGMMA (`wgmma.mma_async`) |
| FP4 Tensor Core | 有（E2M1） | **无** | weight 必须使用 FP8，HBM 读取量翻倍 |
| UTCCP (SF transport) | 有 | **无** | SF 通过 TMA 或 `__ldg` 加载 |
| 2-CTA Cluster MMA | 有（`Allocator2Sm`） | 不适用 | cluster 仅用于 TMA multicast |
| Scale Factor Format | UE8M0 packed `int` | **float** | per-128-K granularity, 32-bit |
| SMEM Capacity | 232 KB | 232 KB | 相同 |
| Registers/SM | 65,536 | 65,536 | 相同，但 TMEM 缺失使 register 压力剧增 |
| WGMMA Shape | N/A (UMMA) | M=64, N=8..256, K=32 (FP8) | 固定 M=64 per warpgroup |

### 2.2 从 SM100 移植的关键改动点

1. **UMMA → WGMMA**：SM100 leader CTA 单 warp 发 UMMA，结果存 TMEM；SM90 每个 math WG（4 warps）独立发 WGMMA，结果驻留 WG register file
2. **TMEM epilogue → Register epilogue**：SM100 epilogue WG 从 TMEM load；SM90 math WG 自身持有 accumulator，math 和 epilogue 在同一 WG 中串行
3. **FP4 weights → FP8 weights**：weight HBM 读取量翻倍，是 SM90 bandwidth bottleneck 的根本原因
4. **UE8M0 SF → float SF**：activation/weight SF 使用 32-bit float，layout 简化（无需 4×32 UTCCP transpose）
5. **2-CTA cluster allocation → single CTA**：SM90 cluster 固定为 1，不做 2-SM MMA 协同

---

## 3. Kernel Architecture: Fused Single Kernel

### 3.1 设计选型理由

选择 **fused single kernel** 而非 dual-kernel 或 split-kernel：

| 方案 | 优势 | 劣势 | 结论 |
|------|------|------|------|
| **Fused single kernel** ✓ | 单次 launch，无 overhead；代码简洁 | 无法独立优化 L1/L2 的线程配置 | **选定** |
| Dual-kernel auto-routing | 各 kernel 独立优化 | 两套 ~1400 行 kernel 维护负担 | 不选 |
| Split L1/L2 | L1/L2 独立 register budget | Kernel launch overhead，无 pipeline overlap | 不选 |

### 3.2 模板参数驱动的 Warpgroup 模式

通过模板参数 `kCooperativeMode` 在编译期选择 cooperative 或 single-WG：

```cpp
template <..., bool kCooperativeMode, ...>
void sm90_fp8_mega_moe_impl(...) {
    // kCooperativeMode=true:  BLOCK_M=128, 两个 math WG M-split
    // kCooperativeMode=false: BLOCK_M=64/32, 只有 WG1 工作, WG2 空闲
}
```

| BLOCK_M | kCooperativeMode | Math WGs | 每 WG 行数 | 场景 |
|---------|-----------------|----------|-----------|------|
| 128 | true | 2 (cooperative) | 64 | 大 M prefill |
| 64 | false | 1 (single) | 64 | 中 M decode |
| 32 | false | 1 (single) | 32 (WGMMA M=64, 用 32 行) | 极小 M decode |

### 3.3 Fused Kernel Pipeline

单 kernel 覆盖 dispatch → L1 GEMM → SwiGLU → L2 GEMM → NVLink scatter → combine reduce：

```
Phase 1: Dispatch
  WG0 dispatch warps → count experts, write source metadata, grid sync
  → NVLink barrier → pull remote tokens/SF/weights into local L1 pool
  → l1_arrival_count release

Phase 2: Linear1 (L1 GEMM + SwiGLU)
  WG0 TMA warps → wait l1_arrival_count, TMA load L1 acts/weights
  WG1 (+ WG2 if cooperative) → WGMMA FP8×FP8, float SF scaling
  → SwiGLU epilogue in registers, top-k weight multiply
  → per-row amax, FP8 quantize, double-buffered TMA store l2_acts
  → write float l2_acts_sf, set l2_arrival_mask

Phase 3: Linear2 (L2 GEMM + NVLink scatter)
  WG0 TMA warps → wait l2_arrival_mask, TMA load L2 acts/weights
  WG1 (+ WG2 if cooperative) → WGMMA FP8×FP8, per-64 dual-half SF
  → BF16 convert, SMEM staging, NVLink scatter to remote combine buffer

Phase 4: Combine Reduce
  NVLink barrier → math/epilogue warps TMA-load top-k BF16 chunks
  → FP32 accumulate → BF16 cast → TMA store output y
```

---

## 4. Thread Layout

### 4.1 384 Threads = 12 Warps = 3 Warpgroups

```
WG0 (128 threads = 4 warps):
  w0-w1 (64 threads):  Dispatch + NVLink pull + workspace cleanup
  w2    (32 threads):   TMA load A (activations + SFA)
  w3    (32 threads):   TMA load B (weights); weight SF 由 math WG ldg

WG1 (128 threads = 4 warps):
  w4-w7: Math WG0 - WGMMA + epilogue
         Cooperative: rows [0, 64)
         Single-WG:  rows [0, BLOCK_M)

WG2 (128 threads = 4 warps):
  w8-w11: Math WG1
          Cooperative: rows [64, 128)
          Single-WG:  idle (register-deallocated)
```

### 4.2 Register Allocation

| Role | Warps | Reg/Warp | Total |
|------|-------|----------|-------|
| WG0 (dispatch+TMA) | 4 | 48 | 6,144 |
| WG1 (math+epilogue) | 4 | 224 | 28,672 |
| WG2 (math+epilogue) | 4 | 224 | 28,672 |
| **Total** | 12 | — | **63,488 / 65,536 (96.9%)** |

Single-WG mode 下 WG2 同样分配 224 但不执行计算（寄存器仍被占用以保持 `__launch_bounds__` 一致性）。

---

## 5. WGMMA Pipeline

### 5.1 MMA Instruction

```
WGMMA FP8: MMA_64x128x32_F32E4M3E4M3_SS_TN
  M = 64 (per warpgroup)
  N = 128 (= BLOCK_N)
  K = 32 (sub-step, 4 steps per BLOCK_K=128)
  Accumulator: float[64] per WG (64×128/128 = 64 registers)
```

### 5.2 K-Loop Inner Pipeline

```
for k_block_idx in range(num_k_blocks):
    // Wait TMA full barrier (SMEM A/B ready)
    full_barriers[stage].wait(phase)

    // WGMMA fence + commit
    warpgroup_arrive()
    for k_step in range(BLOCK_K / 32):  // 4 steps
        desc_a = advance_gmma_desc(smem_a[stage], wg_m_offset, k_step * 32)
        desc_b = advance_gmma_desc(smem_b[stage], 0, k_step * 32)
        WGMMA::wgmma(desc_a, desc_b, accum, k_step > 0 or k_block_idx > 0)
    warpgroup_commit_batch()
    warpgroup_wait<0>()

    // Signal TMA empty barrier (SMEM consumed)
    empty_barriers[stage].arrive()
    advance_pipeline(stage, phase)
```

`wg_m_offset`：cooperative mode 下 WG1=0, WG2=64；single-WG mode 下 WG1=0。

### 5.3 Scale Factor Application

WGMMA 完成后在 registers 中软件 apply SF：

```
// After WGMMA accumulation for a tile:
// 1. Load SFA (per-128 activation scale) from SMEM via ld_shared
// 2. Load SFB (per-128 weight scale) from global via __ldg (software prefetch)
// 3. Apply: accum[i] *= sfa[row_i] * sfb[col_i]
```

### 5.4 L2 Per-64 Dual-Half SF Scaling

L1 epilogue 对每 64 列 output 生成独立 float SF。L2 GEMM K=128 需要拆成两个 per-64 half：

```
// L2 WGMMA K-loop produces raw accum over K=128
// After all K iterations:

// Load two per-64 SFA halves
sfa_half0 = l2_acts_sf[k_group_0 * stride + pool_token_idx]  // k=[0,64)
sfa_half1 = l2_acts_sf[k_group_1 * stride + pool_token_idx]  // k=[64,128)

// Apply: need to track partial sums per half during K-loop
// Option A: Two separate WGMMA passes (K=64 each), scale, add
// Option B: Single K=128 pass, then decompose using SF ratio
//
// Chosen: Option A (two K=64 passes with intermediate scaling)
// K-step 0..1 → accum_half0, scale by sfa_half0 * sfb
// K-step 2..3 → accum_half1, scale by sfa_half1 * sfb
// final = accum_half0 + accum_half1
```

---

## 6. Cooperative Mode (BLOCK_M=128)

### 6.1 M-Split Layout

```
smem_A [128 rows × BLOCK_K]:
  WG1 → desc points to rows [0, 64)
  WG2 → desc points to rows [64, 128)

smem_B [BLOCK_N × BLOCK_K]:
  SHARED - TMA loads ONCE, both WGs consume

smem_cd [128 rows × output_cols]:
  WG1 → writes rows [0, 64)
  WG2 → writes rows [64, 128)
```

### 6.2 TMA Load

```
TMA A warp (w2): load full 128-row A tile
  → full_barriers[stage].arrive_and_expect_tx(SMEM_A_SIZE)
TMA B warp (w3): load BLOCK_N-row B tile (shared)
  → full_barriers[stage].arrive_and_expect_tx(SMEM_B_SIZE)
```

Weight HBM 读取量：相比 single-WG 的两次独立 B load，cooperative 只 load 一次 → **weight HBM 减半**。

### 6.3 L1 Epilogue Sync

两个 WG 各自完成自己 64 行的 SwiGLU + FP8 quant + TMA store 后：

```
// 等两个 WG 的 TMA store 都完成
tma_store_wait<0>()
sync_aligned(256, kEpilogueFullBarrierIdx)

// 一个 WG 代表设置 l2_arrival_mask
if (epilogue_warp_idx == 0 and elect_one_sync())
    red_or_rel_gpu(l2_arrival_mask[pool_block_idx], 1 << n_block_idx)
```

### 6.4 smem_cd Aliasing Guard

L1 (FP8) 和 L2 (BF16) 的 smem_cd 复用同一 SMEM 区域。L2 epilogue scatter 完成前不得进入下一 tile：

```
// L2 epilogue 末尾：
sync_aligned(256, kEpilogueFullBarrierIdx)  // 等所有 WG scatter 完成
```

---

## 7. Single-WG Mode (BLOCK_M=64/32)

### 7.1 退化行为

当 BLOCK_M ≤ 64 时，`kCooperativeMode=false`：

- **WG1** 执行全部 WGMMA + epilogue 工作
- **WG2** 在 kernel 开始后立即 `warpgroup_reg_dealloc` 并 idle
- TMA A warp 只加载 BLOCK_M 行（而非 128 行）
- TMA B warp 不变

### 7.2 BLOCK_M=32 的 WGMMA 适配

WGMMA M 固定为 64，但 BLOCK_M=32 时只有 32 个有效行。处理方式：

- WGMMA 仍计算 64×128 的 tile，但只有 accumulator 中前 32 行的结果被 epilogue 使用
- epilogue 中 `valid_m = min(num_tokens_in_block, BLOCK_M)` 控制实际写入行数
- 计算利用率 50%，但在 token-per-expert ≤ 16 的极端 decode 场景下 tile 数量少，总 compute 量不大

---

## 8. Shared Memory Layout

### 8.1 SMEM Regions

```
smem_buffer (232 KB total, 1024-byte aligned):

┌──────────────────────────────────────────────────────┐
│ Expert Count (kNumExperts × 4B, align to 1024B)      │  Dispatch region
├──────────────────────────────────────────────────────┤
│ Send Buffers (num_dispatch_warps × hidden, aligned)  │
├──────────────────────────────────────────────────────┤
│ smem_cd: max(L1 FP8 double-buf, L2 BF16 staging)    │
│   L1: store_block_m × (BN/2) × sizeof(FP8) × 2     │
│   L2: BLOCK_M × BN × sizeof(BF16)                   │
├──────────────────────────────────────────────────────┤
│ smem_a[0..num_stages-1]:                             │  Pipeline stages
│   BLOCK_M × BLOCK_K × sizeof(FP8) per stage         │
├──────────────────────────────────────────────────────┤
│ smem_b[0..num_stages-1]:                             │
│   BLOCK_N × BLOCK_K × sizeof(FP8) per stage         │
├──────────────────────────────────────────────────────┤
│ smem_sfa[0..num_stages-1]:                           │  Float SF
│   align(2 × BLOCK_M × sizeof(float), 128) per stage │
├──────────────────────────────────────────────────────┤
│ Amax reduction scratch                               │
├──────────────────────────────────────────────────────┤
│ Barriers + combine barriers                          │
└──────────────────────────────────────────────────────┘
```

### 8.2 Pipeline Stage Budget

```
Cooperative (BLOCK_M=128, BLOCK_N=128, BLOCK_K=128):
  smem_a = 128 × 128 × 1B = 16 KB
  smem_b = 128 × 128 × 1B = 16 KB
  smem_sfa = align(256 × 4B, 128) = 1 KB
  Total per stage ≈ 33 KB

  Fixed ≈ 42 KB (expert_count + send_buffers + smem_cd + amax + barriers)
  Available = 232 - 42 = 190 KB
  Num stages = 190 / 33 ≈ **5 stages**

Single-WG (BLOCK_M=64):
  smem_a = 64 × 128 × 1B = 8 KB
  smem_b = 128 × 128 × 1B = 16 KB
  smem_sfa = align(128 × 4B, 128) = 512 B
  Total per stage ≈ 24.5 KB

  Num stages = 190 / 24.5 ≈ **7 stages**

Single-WG (BLOCK_M=32):
  smem_a = 32 × 128 × 1B = 4 KB
  Total per stage ≈ 20.5 KB
  Num stages = 190 / 20.5 ≈ **9 stages**
```

---

## 9. Dispatch Phase

### 9.1 与 SM100 的差异

| 项目 | SM100 | SM90 |
|------|-------|------|
| Dispatch threads | 128 (WG0 全部) | **64** (WG0 w0-w1) |
| Register budget | 48 reg/warp | 48 reg/warp (相同) |
| SF copy | UTCCP 4×32 transpose 写入 l1_sf_buffer | **直接按自然 layout 写入** (无 transpose) |
| Token pull | TMA load 1D + TMA store 1D | 相同 |

### 9.2 Dispatch Pull Path

```
for each assigned token (round-robin across global warps):
    // Round-robin rank selection via iterative min-peeling (same as SM100)
    select source rank

    // TMA load token from remote rank into SMEM send buffer
    tma_load_1d(pull_buffer, remote_token_ptr, pull_mbarrier, kHidden)

    // Load float SF from remote rank (direct copy, no UTCCP transpose)
    for each sf_element:
        local_sf[k_group * stride + pool_token_idx] = remote_sf[...]

    // Load top-k weight, wait TMA, store token via TMA store 1D
    ...

    // Write source metadata + signal block ready
    token_src_metadata[pool_token_idx] = {rank, token_idx, topk_idx}
    red_add_rel(l1_arrival_count[pool_block_idx], 1)
```

---

## 10. L1 Epilogue: SwiGLU + FP8 Quantization

### 10.1 Register-Based Flow

SM90 L1 epilogue 在 math WG 自身的 registers 中执行：

```
// After WGMMA completes for a L1 tile:
// accum[64] holds FP32 values

// 1. Apply SF: accum[i] *= sfa[row_i] * sfb[col_i]

// 2. Load top-k weight for each row
topk_weight = l1_topk_weights[pool_token_idx + row_offset]

// 3. SwiGLU: silu(gate) * up * topk_weight
//    Gate/up pairs interleaved in N dimension (granularity 8)
for each gate-up pair:
    if kActivationClamp != inf:
        gate = clamp(gate, -clamp, clamp)
        up = clamp(up, -clamp, clamp)
    silu_gate = gate / (1 + (kFastMath ? __expf(-gate) : expf(-gate)))
    result = silu_gate * up * topk_weight

// 4. Per-row amax reduction across 4 warps in WG
amax = warp_reduce<4>(max(abs(result_values)))

// 5. Compute exact float SF, quantize to FP8 E4M3
sf = amax / 448.0
fp8_result = cast_to_fp8(result / sf)

// 6. Write float SF to l2_acts_sf (MN-major, per-64-K layout)
l2_acts_sf[k_idx * stride + pool_token_idx] = sf

// 7. STSM to smem_cd, double-buffered TMA store to l2_acts
stsm(smem_cd[tma_stage], fp8_result)
tma_store_2d(tensor_map_l1_output, smem_cd[tma_stage], n_idx, m_idx)

// 8. Signal L2 ready
red_or_rel_gpu(l2_arrival_mask[pool_block_idx], 1 << n_block_idx)
```

### 10.2 Cooperative L1 Epilogue

- WG1 处理 rows [0, 64), WG2 处理 rows [64, 128)
- 各自独立计算 amax、SF、FP8 quant、TMA store
- `sync_aligned(256, kEpilogueFullBarrierIdx)` 等两个 band 都完成后设置 `l2_arrival_mask`

---

## 11. L2 Epilogue: BF16 Scatter + Combine

### 11.1 L2 BF16 Epilogue

```
// After WGMMA with dual-half SF scaling:

// 1. Cast FP32 → BF16, STSM to smem_cd_l2
bf16_values = float_to_bf16(final_accum)
stsm(smem_cd_l2, bf16_values)

// 2. sync WG-internal
sync_aligned(128, kEpilogueWGBarrierStartIdx + wg_idx)

// 3. Read source metadata, NVLink scatter
for each valid row:
    metadata = token_src_metadata[m_idx + row_offset]
    packed = ld_shared(smem_cd_l2 + row * BLOCK_N * sizeof(BF16))
    *sym_buffer.map(combine_dst_ptr, metadata.rank_idx) = packed

// 4. Cooperative: sync_aligned(256) guard smem_cd alias
sync_aligned(256, kEpilogueFullBarrierIdx)
```

### 11.2 Combine Reduce (复用 SM100 逻辑)

NVLink barrier 后，epilogue warps 执行 chunked TMA-based top-k reduce：

```
// Per warp processes one token at a time
for token_idx in stride(sm_idx * num_epi_warps + warp_idx, num_tokens,
                        num_sms * num_epi_warps):

    // Double-buffered TMA load all active top-k BF16 chunks
    // Accumulate in float registers
    float2 reduced[chunk_elems] = {0}
    for each active slot:
        tma_load_1d(load_buffer[stage], combine_chunk_ptr, barrier, chunk_bytes)
        barrier.wait()
        for each element:
            reduced[i] += bf16_to_float2(load_buffer[stage][i])
        flip stage

    // Cast to BF16, TMA store to output y
    tma_store_1d(y + token_idx * hidden + chunk_offset, store_buffer, chunk_bytes)
```

---

## 12. Scheduler

### 12.1 Wave-Based Expert Scheduling

复用 SM100 `MegaMoEScheduler` 核心设计，移除 2-CTA cluster 约束：

- `block_idx = blockIdx.x`（SM100: `blockIdx.x` 隐含 2-CTA 配对）
- L1/L2 block N counts 不再要求偶数
- `kClusterSize = 1` 固定

Wave 结构不变：每 wave 处理 `kNumExpertsPerWave` 个 expert，先遍历 L1 blocks 再回到 wave 起点遍历 L2 blocks。

### 12.2 Auto N-Major

```cpp
bool use_n_major_l2 = (expected_tokens_per_expert >= 256);
// 可通过 DG_SM90_MOE_NMAJOR 覆盖: -1=auto, 0=off, 1=on

// M-major: block_idx → (m_block, n_block) = (idx / num_n_blocks, idx % num_n_blocks)
// N-major: block_idx → (m_block, n_block) = (idx % num_m_blocks, idx / num_m_blocks)
```

---

## 13. Heuristics

### 13.1 BLOCK_M 三档选择

```cpp
static auto get_block_config_for_sm90_mega_moe(
    int num_ranks, int num_experts, int num_max_tokens_per_rank,
    int num_topk, int num_tokens) {

    float expected = float(num_tokens) * num_ranks * num_topk / num_experts;

    if (expected <= 16.5) {
        // 极小 decode: RL long-tail, 极大 EP
        return {.block_m = 32, .cooperative = false, .store_block_m = 16};
    } else if (expected <= 64.5) {
        // 中等 decode
        return {.block_m = 64, .cooperative = false, .store_block_m = 32};
    } else {
        // 大 M prefill / 大 EP decode
        return {.block_m = 128, .cooperative = true, .store_block_m = 32};
    }
}
```

### 13.2 Wave Count

```cpp
static int get_num_experts_per_wave(...) {
    // Same logic as SM100:
    // 1. Estimate L1 blocks per expert
    // 2. Find smallest value whose total blocks keep all SMs busy
    // 3. Round up to nearest divisor of num_experts_per_rank
}
```

### 13.3 Pipeline Stages

```cpp
static auto get_pipeline_config(...) {
    constexpr int smem_capacity = 232448;  // SM90

    // Fixed regions: expert_count + send_buffers + smem_cd + amax + barriers
    // Per-stage: smem_a + smem_b + smem_sfa + 2 barriers

    int num_stages = (smem_capacity - smem_fixed) / smem_per_stage;
    assert(num_stages >= 3);
    return {num_stages, total_smem};
}
```

---

## 14. Synchronization Map

### 14.1 Barrier Inventory

| Barrier | Type | Count | Usage |
|---------|------|-------|-------|
| `dispatch_barriers` | `ClusterTransactionBarrier` | 2 (dispatch warps) | Dispatch pull TMA mbarrier |
| `full_barriers` | `ClusterTransactionBarrier` | num_stages | TMA→math: SMEM data ready |
| `empty_barriers` | `ClusterTransactionBarrier` | num_stages | Math→TMA: SMEM consumed |
| `combine_barriers` | `ClusterTransactionBarrier` | num_epi_warps × 2 | Combine TMA double-buffer |

### 14.2 Intra-SM Named Barriers

| ID | Name | Scope | Usage |
|----|------|-------|-------|
| 0 | `kDispatchBarrierIdx` | 64 dispatch threads | Dispatch internal phases |
| 1 | `kDispatchWithEpilogueBarrierIdx` | 64 + 256 threads | Dispatch↔epilogue handshake |
| 2 | `kEpilogueFullBarrierIdx` | 256 epilogue threads | L1/L2 cooperative sync + smem_cd alias guard |
| 3+ | `kEpilogueWGBarrierStartIdx` | 128 threads (per WG) | WG-local smem_cd write sync |

### 14.3 Cross-Phase Sync

| Sync Point | Mechanism | Producer | Consumer |
|------------|-----------|----------|----------|
| Dispatch → L1 TMA | `l1_arrival_count` release-add | Dispatch warp | TMA A warp |
| L1 Epilogue → L2 TMA | `l2_arrival_mask` bit-OR release | Math WG | TMA A warp |
| L2 Scatter → Combine | NVLink barrier | All epilogue threads | All epilogue threads |
| Combine → Cleanup | `kDispatchWithEpilogueBarrierIdx` | Epilogue | Dispatch |

---

## 15. Weight Transform

### 15.1 SM90 Weight Transform

```python
def transform_weights_for_mega_moe_sm90(l1_weights, l2_weights):
    """
    SM90 weight transform:
    - L1: gate/up FP8 weight interleave along N (granularity 8)
    - L2: no transform
    - Weight SF: natural MN-major float layout, NO UTCCP transpose
    """
    E, N, K = l1_weights.shape  # N = 2 * intermediate_hidden
    half_N = N // 2
    interleaved = torch.empty_like(l1_weights)
    for i in range(0, half_N, 8):
        interleaved[:, 2*i:2*i+8, :] = l1_weights[:, i:i+8, :]          # gate
        interleaved[:, 2*i+8:2*i+16, :] = l1_weights[:, half_N+i:half_N+i+8, :]  # up
    return interleaved, l2_weights
```

### 15.2 Weight SF Layout

```
SM100: UE8M0 packed int, UTCCP 4×32 transposed, per-32-K
SM90:  float, natural MN-major, per-128-K

Shape: (num_experts, N/128, K/128) float32
Kernel indexing: sf[expert * (N/128 * K/128) + n_block * (K/128) + k_block]
```

---

## 16. API Design

### 16.1 独立 SM90 入口

```cpp
// csrc/apis/sm90_mega.hpp (新文件)

static void fp8_mega_moe(
    const torch::Tensor& y,
    const std::tuple<torch::Tensor, torch::Tensor>& l1_weights_tuple,
    const std::tuple<torch::Tensor, torch::Tensor>& l2_weights_tuple,
    const std::optional<torch::Tensor>& cumulative_local_expert_recv_stats,
    const torch::Tensor& sym_buffer,
    const std::vector<int64_t>& sym_buffer_ptrs, const int& rank_idx,
    const int& num_max_tokens_per_rank,
    const int& num_experts, const int& num_topk,
    const std::tuple<int, int, int>& recipe,  // (1, 1, 128) for FP8
    const std::string& activation,
    const std::optional<float>& activation_clamp_opt,
    const bool& fast_math
);

static void register_sm90_apis(pybind11::module_& m) {
    m.def("get_token_alignment_for_sm90_mega_moe", ...);
    m.def("get_symm_buffer_size_for_sm90_mega_moe", ...);
    m.def("fp8_mega_moe", &fp8_mega_moe);
    m.def("transform_weights_for_mega_moe_sm90", ...);
}
```

### 16.2 Python 层路由

```python
# deep_gemm/mega/__init__.py

def fp8_mega_moe(y, l1_weights, l2_weights, ...):
    if _is_sm90():
        _C.fp8_mega_moe(...)
    elif _is_sm100():
        _C.fp8_fp4_mega_moe(...)
    else:
        raise RuntimeError("Unsupported architecture")

def transform_weights_for_mega_moe(l1_weights, l2_weights):
    if _is_sm90():
        return _C.transform_weights_for_mega_moe_sm90(...)
    else:
        return _C.transform_weights_for_mega_moe(...)  # SM100 FP4+UTCCP
```

---

## 17. File Layout

```
deep_gemm/include/deep_gemm/
  impls/
    sm90_fp8_mega_moe.cuh              # Fused kernel (cooperative + single-WG via template)
  scheduler/
    sm90_mega_moe.cuh                  # SM90 scheduler (no 2-CTA, N-major support)

csrc/
  jit_kernels/
    heuristics/
      sm90_mega_moe.hpp                # SM90 config: BLOCK_M selection, wave, stages
    impls/
      sm90_fp8_mega_moe.hpp            # JIT runtime: template instantiation, TMA desc, launch
  apis/
    sm90_mega.hpp                      # SM90-only API: fp8_mega_moe, buffer alloc, pybind11

deep_gemm/
  mega/__init__.py                     # Extended: SM90 routing + transform

tests/
  test_mega_moe_sm90.py               # SM90 correctness + benchmark
```

---

## 18. Environment Variables

| Variable | Values | Default | Description |
|----------|--------|---------|-------------|
| `DG_SM90_MOE_NMAJOR` | `-1\|0\|1` | `-1` (auto) | L2 N-major scheduling |
| `DG_SM90_MOE_PHASE_PROFILE` | `0\|1` | `0` | 启用 per-phase nsys profiling |

---

## 19. Known Risks & Mitigations

### 19.1 Register Budget (P0)

63,488 / 65,536 = 96.9% 利用率。余量 2,048 registers。

- 严格 `__launch_bounds__(384, 1)` + `warpgroup_reg_alloc<224>` / `warpgroup_reg_dealloc<48>`
- Per-64 dual-half SF 在 L2 中需要维护两组 partial accumulator —— 需要仔细编写避免 spill
- 若 224 不够可降至 208 (总 60,416, 92.2%)

### 19.2 HBM Bandwidth Bottleneck (P0)

SM90 FP8 weights 比 SM100 FP4 大 2×，NCU 显示 ~90% DRAM utilization。

- Cooperative kernel 共享 B-load，weight 读取减半
- Auto N-major L2 scheduling 改善 weight tile L2 cache 命中率
- 未来可考虑 cluster=2 TMA multicast（需解决 amax 跨 CTA 同步）

### 19.3 Single-WG BLOCK_M=32 Compute 浪费 (P2)

WGMMA M=64 但只用 32 行，50% compute 浪费。在极端 decode (token-per-expert ≤ 16) 场景下可接受，因为总 compute 量极小。若需进一步优化，未来可考虑 `mma.sync` 路径（PR352 已探索）。

### 19.4 NCU Profiling Limitations (P1)

Symmetric memory + NVLink barrier 与 NCU kernel replay 冲突。

- 使用 `nsys profile` + NVTX markers 做 timeline 分析
- 单 rank 模式 (num_ranks=1) 下可用 NCU
- `DG_SM90_MOE_PHASE_PROFILE` 分段计时

---

## 20. Correctness Verification

### 20.1 Baseline

PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 quant → L2 dequant GEMM → combine。

### 20.2 Test Matrix

| Test | Tokens | Experts | Hidden × Inter | TopK | Ranks |
|------|--------|---------|----------------|------|-------|
| Smoke | 1-4 | 8 | 4096×2048 | 6 | 2 |
| Decode | 16-128 | 48 | 7168×2048 | 6 | 8 |
| Prefill | 256-8192 | 48 | 7168×3072 | 6 | 8 |
| DSV4 Flash | 16-8192 | 256 | 4096×2048 | 6 | 8 |
| DSV4 Pro | 16-8192 | 384 | 7168×3072 | 6 | 8 |
| MiMo-V2.5 | 16-8192 | 256 | 4096×2048 | 8 | 8 |
| Edge: 0 tokens | 0 | any | any | any | 2 |

### 20.3 Tolerance

- BF16 output: `allclose(atol=0.05, rtol=0.1)`

---

## 21. 工程实现 Phases

### Phase 0: 环境搭建与 Baseline 验证 (Day 1-3)

**目标**：建立开发环境、确认 SM100 fused kernel 可正常运行、准备 SM90 PyTorch reference baseline。

| 步骤 | 交付物 | 验证标准 |
|------|--------|----------|
| 0.1 在 H200 集群搭建开发环境 | 可编译运行的 DeepGEMM 环境 | `tests/test_mega_moe.py` SM100 path 通过 |
| 0.2 编写 PyTorch FP32/BF16 reference | `tests/test_mega_moe_sm90.py` baseline 函数 | 与 SM100 kernel output 在 BF16 精度内一致 |
| 0.3 搭建 SM90 单 rank 测试框架 | 可在 H200 上以 `__CUDA_ARCH__=900` 编译的 stub kernel | JIT 编译通过，kernel launch 不崩溃 |
| 0.4 建立 nsys profiling 脚本 | `scripts/run_nsys_mega_moe_sm90.sh` | 能生成 timeline report |

**依赖**：无
**风险**：PyTorch symmetric memory API 在特定 NCCL/CUDA 版本下 multicast 不可用 → 使用 `NCCL_NVLS_ENABLE=0` fallback

---

### Phase 1: 基础设施层 (Day 4-10)

**目标**：完成 kernel 外围的所有 C++/Python 基础设施，使 JIT 编译框架可以实例化空 kernel 并正确分配 buffer。

| 步骤 | 文件 | 交付物 | 验证标准 |
|------|------|--------|----------|
| 1.1 SM90 Scheduler | `deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh` | 移除 2-CTA 约束的 scheduler，支持 N-major | 单元测试：给定 expert token counts 和 SM 数，验证 block 分配覆盖所有 expert tiles |
| 1.2 SM90 Heuristics | `csrc/jit_kernels/heuristics/sm90_mega_moe.hpp` | BLOCK_M 三档 (32/64/128)，wave count，pipeline stages 计算 | 给定 shape 参数验证输出 config 合理（stages ≥ 3，smem ≤ 232KB） |
| 1.3 JIT Runtime | `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp` | Template instantiation 代码生成、TMA descriptor 创建、kernel launch | 空 kernel 可 JIT 编译 + launch（立即 return） |
| 1.4 API & Buffer | `csrc/apis/sm90_mega.hpp` | `fp8_mega_moe`、`get_symm_buffer_size_for_sm90_mega_moe`、pybind11 注册 | Python 可调用 `_C.get_symm_buffer_size_for_sm90_mega_moe(...)` |
| 1.5 Weight Transform | Python `mega/__init__.py` 扩展 | `transform_weights_for_mega_moe_sm90` (L1 gate/up interleave) | 对比 SM100 的 interleave 结果在 FP8 数据维度一致 |
| 1.6 Python Routing | `deep_gemm/mega/__init__.py` | `_is_sm90()` 检测 + 路由到 SM90 API | 在 H200 上 import deep_gemm 并调用 SM90 路径不报错 |

**验证里程碑**：`python -c "import deep_gemm; deep_gemm.get_symm_buffer_for_mega_moe(...)"`  在 H200 上正确返回 buffer views。

**依赖**：Phase 0 完成
**关键参考**：
- SM100 对应文件：`csrc/jit_kernels/heuristics/mega_moe.hpp` (heuristics)、`csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (JIT)
- SM90 GEMM 参考：`csrc/jit_kernels/heuristics/sm90.hpp` (arch spec)

---

### Phase 2: Dispatch + L1 Pool 填充 (Day 11-15)

**目标**：只实现 dispatch 数据路径，不进入 TMA/WGMMA。验证 expert count、token routing、source metadata、跨 rank pull、`l1_arrival_count` 等 dispatch 产物正确，为后续 L1 GEMM 提供稳定输入。

| 步骤 | 交付物 | 验证标准 |
|------|--------|----------|
| 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 2 完成
**关键挑战**：
- WGMMA descriptor 构造（参考 `deep_gemm/mma/sm90.cuh` 中的 `make_gmma_desc`）
- Cooperative mode 下两个 WG 的 desc 指向 smem_a 的不同 64-row 偏移
- `warpgroup_arrive()` / `warpgroup_commit_batch()` / `warpgroup_wait<0>()` 的 fence 语义

---

### Phase 4: L1 Epilogue — SwiGLU + FP8 Quant (Day 21-30)

**目标**：在 Phase 3 的 WGMMA 结果上完成 L1 epilogue 全流程，输出正确的 FP8 `l2_acts` 和 float `l2_acts_sf`。

| 步骤 | 交付物 | 验证标准 |
|------|--------|----------|
| 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 3 完成
**关键挑战**：
- WGMMA register layout 到 gate/up 列的映射关系（参考 SM100 的 `SM100_TMEM_LOAD_16dp256b1x` 读取模式，SM90 的 register layout 不同）
- Double-buffered TMA store 的 `tma_store_wait<1>` 与 smem_cd 复用
- 寄存器压力：SwiGLU 中间变量 + amax + SF + FP8 cast 共享 224 reg/warp 的 math WG budget

---

### Phase 5: L2 GEMM + NVLink Scatter (Day 31-40)

**目标**：L2 WGMMA pipeline + per-64 dual-half SF + BF16 epilogue + NVLink scatter 到 remote combine buffer。

| 步骤 | 交付物 | 验证标准 |
|------|--------|----------|
| 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 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 一致）
- smem_cd_l2 (BF16) 和 smem_cd_l1 (FP8) 的 aliasing — 同一 SMEM 地址

---

### Phase 6: Combine Reduce + 端到端 Fused (Day 41-48)

**目标**：完成 combine reduce、NVLink barrier、workspace 清理。实现完整端到端 fused kernel。

| 步骤 | 交付物 | 验证标准 |
|------|--------|----------|
| 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 5 完成
**关键挑战**：
- Combine 逻辑直接移植自 SM100 (与 TMEM/UMMA 无关)，但需要确认 SM90 下 TMA 1D load/store 的行为
- NVLink barrier 的 phase 计数在多次 kernel 调用间的正确性
- Workspace 清理时序：dispatch warps 与 epilogue warps 并行执行不同的清理任务

---

### Phase 7: Single-WG 变体 (BLOCK_M=64/32) (Day 49-55)

**目标**：实现 `kCooperativeMode=false` 路径，支持小 BLOCK_M 场景。

| 步骤 | 交付物 | 验证标准 |
|------|--------|----------|
| 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 6 完成（cooperative 端到端已验证）
**关键挑战**：
- BLOCK_M=32 时 WGMMA 计算 64 行但只用 32 行的 register 浪费
- Single-WG mode 下 WG2 的 named barrier arrive 需要调整（WG2 不参与 epilogue barrier）

---

### Phase 8: 多 Rank 集成测试 (Day 56-62)

**目标**：在 8×H200 集群上验证完整 multi-rank 正确性和性能。

| 步骤 | 交付物 | 验证标准 |
|------|--------|----------|
| 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 7 完成

---

### Phase 9: 性能优化 (Day 63-75)

**目标**：基于 nsys timeline 分析，针对瓶颈进行优化迭代。

| 步骤 | 优化方向 | 预期收益 | 依据 |
|------|---------|---------|------|
| 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 8 baseline，主要 shape 上 latency 降低 10-20%。

---

### Phase 依赖关系图

```
Phase 0 (环境)
    │
    v
Phase 1 (基础设施)
    │
    v
Phase 2 (Dispatch)
    │
    v
Phase 3 (L1 TMA + WGMMA)
    │
    v
Phase 4 (L1 Epilogue)
    │
    v
Phase 5 (L2 GEMM + Scatter)
    │
    v
Phase 6 (Combine + 端到端)
    │
    ├──────────────────┐
    v                  v
Phase 7 (Single-WG)   Phase 8 (多 Rank 集成)
    │                  │
    └──────┬───────────┘
           v
    Phase 9 (性能优化)
```

### 时间线总结

| Phase | 天数 | 累计 | 核心交付 |
|-------|------|------|---------|
| 0 | 3 | 3 | 环境 + baseline |
| 1 | 7 | 10 | 基础设施 (可 JIT 编译) |
| 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-6 的串行 kernel 开发。