DeepGEMM/MEGAMOE_SM90_DESIGN.md

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：

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

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 三档选择

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

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

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

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 入口

// 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 层路由

# 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 + 单 tile L1 GEMM (Day 11-20)

目标：实现 dispatch + L1 WGMMA 的正确 TMA pipeline，不含 epilogue（accumulator 直接丢弃）。验证 dispatch pull、arrival count、TMA A/B 加载、WGMMA FP8×FP8 累积 的端到端正确性。

步骤	交付物	验证标准
2.1 Kernel 骨架	sm90_fp8_mega_moe.cuh 中的 warp role 分配：WG0 dispatch+TMA, WG1/WG2 math	编译通过，launch_bounds(384,1)
2.2 Barrier 初始化	所有 mbarrier (full/empty/dispatch/combine) 初始化 + __syncthreads	无 launch failure
2.3 Dispatch phase	64-thread dispatch: expert count、source metadata、grid sync、NVLink barrier、token pull	单 rank 模式下 l1_token_buffer 中数据正确
2.4 TMA A/SFA load	w2 TMA A warp：wait l1_arrival_count，TMA load L1 acts + float SFA to SMEM	dump SMEM A 区域，与 global 数据逐字节对比
2.5 TMA B load	w3 TMA B warp：TMA load L1 FP8 weights to SMEM	dump SMEM B 区域对比
2.6 WGMMA pipeline	WG1 math warp：warpgroup_arrive, 4× K-step WGMMA::wgmma, warpgroup_commit_batch, warpgroup_wait<0>	单 tile (1 expert, 1 block) 的 accumulator 与 PyTorch FP32 GEMM 结果对比（未 scale）
2.7 SF scaling	accum *= sfa * sfb (sfb via __ldg)	Scale 后的 accumulator 与 reference scaled GEMM 一致
2.8 Empty barrier	WGMMA 完成后 arrive empty_barrier → TMA 可重用 SMEM	多 K-block 的 pipeline 不死锁

验证里程碑：单 expert、BLOCK_M=128 的 L1 GEMM（无 epilogue）输出与 PyTorch (dequant(x) @ dequant(w).T) * sfa * sfb 在 FP32 accumulator 精度内 allclose(atol=1e-2, rtol=0.05)。

依赖：Phase 1 完成 关键挑战：

• 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 3: L1 Epilogue — SwiGLU + FP8 Quant (Day 21-30)

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

步骤	交付物	验证标准
3.1 Top-k weight load	从 l1_topk_weights_buffer load weight per row	值与 dispatch 写入一致
3.2 SwiGLU 计算	gate/up deinterleave + silu(gate) * up * topk_weight in registers	与 reference SwiGLU 一致（精度：atol=0.02）
3.3 Activation clamp	kActivationClamp 模板参数控制 clamp	clamp 后无超范围值
3.4 fast_math SwiGLU	__expf + fast_rcp 快速路径	与精确版相差 < 0.5%
3.5 Per-row amax	warp_reduce<4>(max(abs(...))) 跨 WG 内 4 warps reduce	amax 值与 PyTorch abs().max(dim=-1) 一致
3.6 Cooperative amax	两个 WG 独立 amax（各自 64 行），SMEM scratch 区域用于跨 warp reduce	每行 amax 正确
3.7 FP8 quantize	sf = amax / 448.0，fp8 = cast(result / sf)	反量化后与 reference 误差 < 1%
3.8 Float SF write	写入 l2_acts_sf (MN-major, per-64-K layout)	SF 值正确、layout 与 L2 TMA load 兼容
3.9 STSM + double-buf TMA store	STSM 到 smem_cd[stage]，TMA store 到 l2_acts	l2_acts 内容与 reference FP8 quant output 一致
3.10 l2_arrival_mask 设置	red_or_rel_gpu(mask, 1 << n_block_idx)	单 expert 全部 L1 N-blocks 完成后 mask == expected
3.11 Cooperative sync	sync_aligned(256) 等两个 WG 都 TMA store 完成后设置 mask	无 race condition

验证里程碑：单 expert 的 l2_acts (FP8) 和 l2_acts_sf (float) 与 reference SwiGLU → FP8_quant 一致。Multi-expert 场景下 l2_arrival_mask 正确触发。

依赖：Phase 2 完成 关键挑战：

• 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 4: L2 GEMM + NVLink Scatter (Day 31-40)

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

步骤	交付物	验证标准
4.1 L2 TMA wait	TMA A warp wait l2_arrival_mask == expected	正确等待 L1 所有 N-blocks 完成
4.2 L2 TMA load	TMA load l2_acts + 两个 per-64 SFA halves	SMEM 数据正确
4.3 L2 WGMMA dual-half	K-step 0-1 累积 half0，K-step 2-3 累积 half1	两组 partial accum 正确
4.4 Dual-half SF apply	final = half0 * sfa0 * sfb + half1 * sfa1 * sfb	与 reference L2 dequant GEMM 一致
4.5 Weight SF __ldg	Math WG prefetch L2 weight SF from global	值正确
4.6 BF16 cast + STSM	FP32 → BF16 cast，STSM 到 smem_cd_l2 (BF16 swizzle)	SMEM 内容正确
4.7 Source metadata read	token_src_metadata[pool_token_idx + row] → {rank, token, topk}	值与 dispatch 写入一致
4.8 NVLink scatter	ld_shared + sym_buffer.map + remote store	单 rank 下本地 combine buffer 数据正确
4.9 Cooperative smem_cd guard	sync_aligned(256) 防止下一 tile 覆盖 smem_cd	多 tile 序列无 data corruption
4.10 Multi-expert persistent loop	Scheduler for_each_block 遍历多 expert 的 L1→L2 cycle	全量 expert 输出正确

验证里程碑：多 expert、单 rank 下 combine_token_buffer 内容与 reference L2 GEMM BF16 输出一致。l2_arrival_mask 正确控制 L1→L2 依赖。

依赖：Phase 3 完成 关键挑战：

• 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 5: Combine Reduce + 端到端 Fused (Day 41-48)

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

步骤	交付物	验证标准
5.1 NVLink barrier	Epilogue grid sync + cross-rank signal	多 rank 下所有 rank 到达 barrier
5.2 Dispatch↔Epilogue handshake	kDispatchWithEpilogueBarrierIdx sync	Dispatch 和 epilogue 阶段正确交替
5.3 Combine TMA load	Double-buffered TMA 从 combine buffer load BF16 chunks	Load 数据与 scatter 写入一致
5.4 FP32 accumulate	逐 top-k slot 累加到 float registers	累加结果与 reference sum(topk_weights * expert_outputs) 一致
5.5 BF16 cast + TMA store	Cast + TMA store 到 output y	y 与 reference 在 BF16 精度内一致
5.6 Workspace cleanup	Dispatch warps 清理 expert counts、arrival counters、metadata	下次 kernel 调用前 workspace 已归零
5.7 NVLink cleanup barrier	跨 rank 确认 workspace 清理完成	连续两次 kernel 调用无 data corruption
5.8 端到端 correctness	完整 fp8_mega_moe(...) 调用	test_mega_moe_sm90.py Smoke test 通过 (2 ranks, 4096×2048, topk=6)

验证里程碑：2-rank Smoke test (4 tokens, 8 experts) allclose(atol=0.05, rtol=0.1) 通过。

依赖：Phase 4 完成 关键挑战：

• Combine 逻辑直接移植自 SM100 (与 TMEM/UMMA 无关)，但需要确认 SM90 下 TMA 1D load/store 的行为
• NVLink barrier 的 phase 计数在多次 kernel 调用间的正确性
• Workspace 清理时序：dispatch warps 与 epilogue warps 并行执行不同的清理任务

---

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

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

步骤	交付物	验证标准
6.1 模板分支	if constexpr (kCooperativeMode) 控制 WG2 行为	两种模式均编译通过
6.2 WG2 idle	WG2 在 kCooperativeMode=false 时 warpgroup_reg_dealloc 后 idle	无 deadlock / launch failure
6.3 TMA A 调整	TMA A 只加载 BLOCK_M 行（64 或 32）而非 128	SMEM A 数据正确
6.4 WGMMA 有效行	valid_m = min(tokens_in_block, BLOCK_M) 控制 epilogue 写入范围	无 OOB 写入
6.5 Scheduler 适配	动态 BLOCK_M 下 scheduler 的 block 分配和 pool_block_offset 正确	全量 expert 覆盖，无漏 tile
6.6 Heuristics JIT 分发	根据 expected_tokens_per_expert 选择 BLOCK_M，JIT 编译对应模板	正确选择 config 并编译
6.7 Small M correctness	BLOCK_M=64: 单 WG 处理完整 64 行	Smoke test 通过
6.8 Extreme decode	BLOCK_M=32: token-per-expert=1-2	极端 decode 场景输出正确

验证里程碑：BLOCK_M=64 和 BLOCK_M=32 的 Smoke test 通过。动态 BLOCK_M 选择逻辑根据 token count 正确切换。

依赖：Phase 5 完成（cooperative 端到端已验证） 关键挑战：

• BLOCK_M=32 时 WGMMA 计算 64 行但只用 32 行的 register 浪费
• Single-WG mode 下 WG2 的 named barrier arrive 需要调整（WG2 不参与 epilogue barrier）

---

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

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

步骤	交付物	验证标准
7.1 Multi-rank launch	8-rank torchrun 启动 MegaMoE kernel	无 hang / crash
7.2 Correctness sweep	Token sweep (1-8192) × Shape sweep × BLOCK_M 三档	全部 allclose 通过
7.3 Edge cases	0 tokens, unbalanced routing, single-expert-per-token	无 crash，输出正确
7.4 Benchmark framework	bench_mega_moe_sm90.py：latency, TFLOPS, HBM GB/s	可生成 performance table
7.5 nsys timeline	Phase profile: dispatch / L1 / L2 / combine 各阶段时间	识别主要 stall 来源
7.6 DeepEP v1 baseline	对比 DeepEP dispatch + grouped GEMM + combine	输出 speedup ratio
7.7 Regression test	CI 集成到 tests/test_mega_moe_sm90.py	pytest 通过

验证里程碑：8-rank × DSV4 Flash shape (4096×2048, E256, topk6)：correctness 通过 + fused latency < 3000us (128 tokens)。

依赖：Phase 6 完成

---

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

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

步骤	优化方向	预期收益	依据
8.1 N-major scheduling tuning	调整 tokens/expert >= 256 阈值	大 M 下 weight L2 cache 命中率提升	NCU L2 sector miss 数据
8.2 SFA 软件 prefetch	L2 per-64 SFA 的 __ldg prefetch 提前一个 K-block	减少 math stall waiting on SF	nsys timeline 中 SF load 占比
8.3 Combine chunk 大小调优	根据 hidden dim 选择 1 vs 2 chunks	减少 combine 阶段 TMA round-trip	combine 占总时间比例
8.4 Dispatch pull 优化	增加 dispatch warp 内的 token parallelism	减少 dispatch phase 时间	nsys dispatch 耗时
8.5 Register spill 消除	检查 SASS 中的 local memory access，调整变量生命周期	消除 spill 带来的 DRAM 流量	cuobjdump --dump-sass
8.6 wave_count 调优	根据实测 per-shape 最优 experts_per_wave	减少 tail effect	benchmark sweep
8.7 Optional: cluster=2 探索	TMA multicast B tile（需解决 amax 跨 CTA sync）	weight HBM 进一步减半	若 cooperative 的 HBM 仍是瓶颈

验证里程碑：相对 Phase 7 baseline，主要 shape 上 latency 降低 10-20%。

---

Phase 依赖关系图

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

时间线总结

Phase	天数	累计	核心交付
0	3	3	环境 + baseline
1	7	10	基础设施 (可 JIT 编译)
2	10	20	Dispatch + L1 WGMMA 正确
3	10	30	L1 SwiGLU + FP8 quant 正确
4	10	40	L2 GEMM + NVLink scatter 正确
5	8	48	端到端 fused kernel 正确 (2-rank)
6	7	55	动态 BLOCK_M 全变体正确
7	7	62	8-rank 集成 + benchmark
8	13	75	性能调优完成

总工期：约 75 工作日（10-11 周），含充分的验证和调试时间。关键路径是 Phase 2-5 的串行 kernel 开发（40 天）。