diff --git a/MEGAMOE_SM90_DESIGN.md b/MEGAMOE_SM90_DESIGN.md new file mode 100644 index 0000000..277b900 --- /dev/null +++ b/MEGAMOE_SM90_DESIGN.md @@ -0,0 +1,1021 @@ +# 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& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& 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& recipe, // (1, 1, 128) for FP8 + const std::string& activation, + const std::optional& 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 + 单 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 天）。 diff --git a/megamoe-research-reports/megamoe.jpg b/megamoe-research-reports/megamoe.jpg new file mode 100644 index 0000000..b5c6fcb Binary files /dev/null and b/megamoe-research-reports/megamoe.jpg differ diff --git a/megamoe-research-reports/megamoe_sm90_research_report.md b/megamoe-research-reports/megamoe_sm90_research_report.md new file mode 100644 index 0000000..abe8361 --- /dev/null +++ b/megamoe-research-reports/megamoe_sm90_research_report.md @@ -0,0 +1,1119 @@ +# 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()` 方法，为 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。 diff --git a/megamoe-research-reports/pr304_sm100_kernel_suite_code_review.md b/megamoe-research-reports/pr304_sm100_kernel_suite_code_review.md new file mode 100644 index 0000000..75840d5 --- /dev/null +++ b/megamoe-research-reports/pr304_sm100_kernel_suite_code_review.md @@ -0,0 +1,75 @@ +# PR304 SM100/FP8/FP4/BF16 GEMM Suite 代码review报告 + +## 范围 + +- Worktree: `pr-304` +- HEAD: `211d2678d` +- 审查方式: 代码review +- 这是 DeepGEMM 主分支的早期 snapshot，包含完整的 SM100 和部分 SM90 kernel 套件 + +## 实现概述 + +### SM100 FP8/FP4 MegaMoE Fused Kernel + +`deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (~1364 lines) + +**warp 分工与线程布局 (硬编码 dispatch=128, TMA=128, epilogue=256):** + +| Warp Index | 数量 | Threads | Role | 所属 Warpgroup | +|---:|---:|---:|---|---| +| 0–3 | 4 warps | 128 | Dispatch: 扫描 topk_idx，写 expert send/recv count，NVLink pull token+SF | WG0 | +| 4 | 1 warp | 32 | TMA A + SFA loader | WG1 | +| 5 | 1 warp | 32 | TMA B + SFB loader | WG1 | +| 6 | 1 warp | 32 | MMA issue warp (仅 leader CTA) | WG1 | +| 7 | 1 warp | 32 | Idle (warpgroup 占位) | WG1 | +| 8–15 | 8 warps | 256 | Math UMMA + SwiGLU epilogue + BF16 scatter + Combine reduce | WG2, WG3 | + +**精确 warp 统计: 16 warps = 512 threads = 4 warpgroups** | + +**寄存器分配 (典型配置: dispatch=128, TMA=128, epilogue=256):** + +| Role | 每 thread register 数 | Threads | 总 register 消耗 | +|---|---|---|---| +| Dispatch | 48 | 128 | 6,144 | +| TMA/non-epilogue | 40 | 128 | 5,120 | +| Math/epilogue | 208 | 256 | 53,248 | +| **总计** | | **512** | **64,512 (恰好 64K budget)** | + +**关键约束:** +- `kNumDispatchThreads % 128 == 0` → dispatch 独占整数个 warpgroup +- `kNumNonEpilogueThreads == 128` → TMA 恰好一个 warpgroup +- `kNumEpilogueThreads % 128 == 0` → math 独占整数个 warpgroup +- 2-CTA cluster MMA (SM100 UMMA)，cluster_size=2 + +### 其他 SM100 Kernel + +| Kernel | 线程布局 | 备注 | +|---|---|---| +| `sm100_fp8_gemm_1d1d` | dispatch≥128 + TMA=128 + math≥128 | FP8 1D-1D grouped GEMM，TMEM accumulator | +| `sm100_fp8_fp4_gemm_1d1d` | 同上 | FP8/FP4 混合精度，UE8M0 SF packing | +| `sm100_bf16_gemm` | 同上 | BF16 GEMM, TMEM, stage merging | +| `sm100_fp8_mqa_logits` | cluster=2, 128 specialized + 256 math | FP8 MQA attention logits | +| `sm100_tf32_hc_prenorm_gemm` | cluster=2, dispatch+math | HyperConnection prenorm GEMM | + +## 代码review发现 + +### 中: SM100 MegaMoE 调度器波次边界越界读取 + +`deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh:73-81` 中 `get_num_tokens()` 声明的 `valid_value` 未初始化，在波次最后一个 expert 完成时 `advance_expert_idx()` 会在无人拥有该 expert index 的 warp 中调用 `get_num_tokens(wave_end)`。 + +### 中: SM100 kernel 的 register budget 在默认配置下已达上限 + +dispatch=128*48 + TMA=128*40 + epilogue=256*208 = 64512，恰好等于 64K reg budget（SM 寄存器总数 65536，扣除 CUDA runtime 开销后约 64512 可用）。任何 register 增加都会溢出。 + +### 低: TMA TMA load warp 中 `warp_idx == kNumDispatchWarps + 3` 的 idle warp 仅执行 dealloc，无实际工作 + +## 正面评价 + +- 完整覆盖 SM100 GEMM、MegaMoE、MQA attention、paged attention 场景 +- 2-CTA cluster 充分利用 SM100 硬件能力 +- HEURISTICS 文件结构清晰，按架构分离 + +## 建议检查清单 + +- 修复 scheduler `valid_value` 未初始化 +- 所有 kernel 的 `warpgroup_reg_dealloc` 参考 SM90 已验证模式，当前无问题 diff --git a/megamoe-research-reports/pr323_sm90_fused_megamoe_code_review.md b/megamoe-research-reports/pr323_sm90_fused_megamoe_code_review.md new file mode 100644 index 0000000..4636bb8 --- /dev/null +++ b/megamoe-research-reports/pr323_sm90_fused_megamoe_code_review.md @@ -0,0 +1,93 @@ +# PR323 SM90 Fused MegaMoE 代码review报告 + +## 范围 + +- Worktree: `pr-323` +- HEAD: `23f46aa68` +- 审查方式: 代码review +- 主题: 首次引入 SM90/Hopper FP8 MegaMoE 单 kernel 实现 + +## 实现概述 + +### SM90 FP8 MegaMoE Fused Kernel + +`deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` (~1935 lines) + +**warp 分工与线程布局（3 种配置）:** + +**配置 1: 512-epilogue (dispatch+TMA 各自独占 warpgroup)** + +| Warp Index | 数量 | 所属 Warpgroup | Role | +|---:|---:|---|---| +| 0–3 | 4 warps (128 threads) | WG0 | Dispatch | +| 4–7 | 4 warps (128 threads) | WG1 | TMA A+SFA / B+SFB / MMA issue / idle | +| 8–23 | 16 warps (512 threads) | WG2–WG5 | Math WGMMA + epilogue + combine | + +**精确 warp 统计: 24 warps = 768 threads = 6 warpgroups** + +**配置 2: 256-epilogue (dispatch+TMA 共享 WG0, compact)** + +| Warp Index | 数量 | 所属 Warpgroup | Role | +|---:|---:|---|---| +| 0–1 | 2 warps (64 threads) | WG0 | Dispatch | +| 2–3 | 2 warps (64 threads) | WG0 | TMA A+SFA / B+SFB | +| 4–11 | 8 warps (256 threads) | WG1, WG2 | Math WGMMA + epilogue + combine | + +**精确 warp 统计: 12 warps = 384 threads = 3 warpgroups** + +**配置 3: ≤256-total (dispatch+TMA 共享 WG0, BLOCK_M=32)** + +| Warp Index | 数量 | 所属 Warpgroup | Role | +|---:|---:|---|---| +| 0–1 | 2 warps (64 threads) | WG0 | Dispatch | +| 2–3 | 2 warps (64 threads) | WG0 | TMA A+SFA / B+SFB | +| 4–7 | 4 warps (128 threads) | WG1 | Math WGMMA + epilogue + combine | + +**精确 warp 统计: 8 warps = 256 threads = 2 warpgroups** | + +**关键差异 vs SM100:** +- 无 TMEM，WGMMA accumulator 在 register 中 +- SF 为 per-128 float（非 UE8M0 int） +- cluster_size ≤ 2（无 2-CTA UMMA） +- dispatch 最小仅 64 threads（SM100 要求 ≥128），意味着 dispatch 可以和 TMA 共享 warpgroup + +**寄存器分配（多配置）:** + +| 配置 | Dispatch Threads | TMA Threads | Epilogue Threads | Total Threads | Disp Reg | TMA Reg | Epi Reg | Total Reg | +|---|---|---|---|---|---|---|---|---| +| 512-epilogue | 128 | 128 | 512 | 768 | 32 | 24 | 112 | 64,512 | +| 256-epilogue | 64 | 64 | 256 | 384 | 48 | 40 | 168 | 48,640 | +| ≤256-total | 64 | 64 | 128 | 256 | 48 | 40 | 256 | 38,400 | + +`kEpilogueRegisterBudget` 模板参数允许调用方精确控制 math warpgroup 的 register 分配。 + +**调度特点:** +- `sm90_fp8_mega_moe_for_each_block_split` 将 L1/L2 phase 拆分到不同 CTA，每个 CTA 只执行一种 phase +- 使用 `kNonEpilogueWarpsInWarpgroup` 确保 dispatch+TMA 恰好填满整数个 warpgroup + +## 代码review发现 + +### 高: 512-epilogue 配置下 dispatch register 仅 32，可能成为瓶颈 + +dispatch warp 需要大量 `smem_expert_count` 和 route metadata 操作。32 register/warp 对于复杂的 `read_topk_idx` lambda 和 rank round-robin 选择可能不够（register spilling 到 local memory）。 + +### 中: `kEpilogueRegisterBudget` 默认 0 的自动推导逻辑复杂 + +`sm90_fp8_mega_moe.cuh:427-432` 的自动 register budget 推导依赖 thread 数和总线程数，外部调用者可能不知情使用次优配置。 + +### 中: 与 SM100 kernel 共享 scheduler，但 SM90 无 cluster 概念 + +`mega_moe.cuh` scheduler 包含 `kClusterSize=2` 相关逻辑，SM90 cluster_size=1 时跳过但增加了编译期复杂度。 + +### 低: 256-total-thread 配置下 epilogue 获得 256 reg/warp，几乎占满 budget + +## 正面评价 + +- 首次将 MegaMoE 带到 SM90/Hopper，填补了架构空白 +- 灵活的 register budget 控制允许调用方按场景调优 +- L1/L2 phase 拆分减少了单一 CTA 的复杂度 + +## 建议检查清单 + +- 在 H100/H200 上实测 512-epilogue 配置的 dispatch spilling 情况 +- 简化或文档化 `kEpilogueRegisterBudget` 的默认行为 diff --git a/megamoe-research-reports/pr347_infra_refactor_code_review.md b/megamoe-research-reports/pr347_infra_refactor_code_review.md new file mode 100644 index 0000000..1415d1f --- /dev/null +++ b/megamoe-research-reports/pr347_infra_refactor_code_review.md @@ -0,0 +1,53 @@ +# PR347 基础设施/SM100 重构 代码review报告 + +## 范围 + +- Worktree: `pr-347` +- HEAD: `2b8dfd0e8` +- 审查方式: 代码review +- 主题: SM100 MegaMoE 重构，SM90 MegaMoE kernel **已被移除** + +## 实现概述 + +### 核心变化 + +pr-347 是一个**中间重构 PR**。它不包含 SM90 MegaMoE kernel（相比 pr-323 已移除 `sm90_fp8_mega_moe.cuh`），重点关注: + +1. **SM100 FP8/FP4 MegaMoE kernel 重构** — 代码格式、变量命名、注释统一 +2. **Heuristics 模块重组** — `mega_moe.hpp` 从 211 行扩展到 276 行，增加了更多 shape 分支 +3. **Scheduler 增强** — `mega_moe.cuh` 从 221 行扩展到 273 行，支持更多 block 分配策略 +4. **API 层清理** — 移除 SM90 专用的 `csrc/apis/sm90_mega.hpp` + +### SM100 FP8/FP4 MegaMoE Fused Kernel + +`deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (~1364 lines) + +线程和寄存器布局与 pr-304 一致: + +| Role | Threads | Warps | 所属 Warpgroup | Reg/thread | +|---|---|---|---|---| +| Dispatch | 128 | 4 | WG0 | 48 | +| TMA (A+SFA, B+SFB, MMA issue, idle) | 128 | 4 | WG1 | 40 | +| Math/epilogue + Combine | 256 | 8 (2 warpgroups) | WG2, WG3 | 208 | +| **总计** | **512** | **16** | **4 WGs** | **64,512/64K** | + +## 代码review发现 + +### 低: 这是一个重构 PR，kernel 逻辑变化小 + +主要改动是代码组织（命名、注释、文件结构调整），对运行时行为影响有限。风险评估较低。 + +### 低: SM90 MegaMoE 能力被移除 + +相比 pr-323，sm90_fp8_mega_moe.cuh 和相关 JIT host/heuristics 文件不再存在。如果这是故意的（为后续 PR 清理），则合理；如果是不小心丢失，需要注意。 + +## 正面评价 + +- 代码风格和注释质量有显著提升 +- Heuristics 模块增加了更多 shape 分支覆盖 +- 为后续 PR352 的分层 MegaMoE 打下基础 + +## 建议检查清单 + +- 确认 SM90 MegaMoE 移除是故意的，非 merge 错误 +- 验证重构后 SM100 MegaMoE 的 bitwise correctness 未退化 diff --git a/megamoe-research-reports/pr352_sm90_split_l1l2_megamoe_code_review.md b/megamoe-research-reports/pr352_sm90_split_l1l2_megamoe_code_review.md new file mode 100644 index 0000000..a6fcea2 --- /dev/null +++ b/megamoe-research-reports/pr352_sm90_split_l1l2_megamoe_code_review.md @@ -0,0 +1,130 @@ +# PR352 SM90 Split L1/L2 MegaMoE 代码review报告 + +## 范围 + +- Worktree: `pr-352` +- HEAD: `655075ef3` +- 审查方式: 代码review +- 主题: SM90 MegaMoE 增强 — phase 分离、compact frontend、多 epilogue 策略 + +## 实现概述 + +### SM90 FP8 MegaMoE Enhanced Kernel + +`deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` (~2507 lines) + +本 kernel 是 pr-323 的重度增强版。核心创新是使用 **C++ 宏** 将同一个 kernel body 参数化为多种配置变体。 + +**warp 分工与线程布局（按 Phase 策略枚举）:** + +**Compact 模式 (kCompactFrontendWarpgroup, topk=2, cluster=1, default):** + +| Warp Index | 所属 Warpgroup | Role | +|---:|---|---| +| 0–1 (2 warps, 64 threads) | WG0 | Dispatch | +| 2–3 (2 warps, 64 threads) | WG0 | TMA A+SFA / B+SFB | +| 4–11 (8 warps, 256 threads) | WG1, WG2 | Math WGMMA + epilogue + combine | + +**精确 warp 统计: 12 warps = 384 threads = 3 warpgroups** + +**Serial/Wide N 模式 (BLOCK_M=32, topk≥8, cluster=2):** + +| Warp Index | 所属 Warpgroup | Role | +|---:|---|---| +| 0–3 (4 warps, 128 threads) | WG0 | Dispatch | +| 4–7 (4 warps, 128 threads) | WG1 | TMA A+SFA / B+SFB / MMA issue / idle | +| 8–11 (4 warps, 128 threads) | WG2 | Math WGMMA + epilogue + combine | + +**精确 warp 统计: 12 warps = 384 threads = 3 warpgroups** + +**Serial/Wide N 模式 (BLOCK_M=64, topk≥8, cluster=2):** + +| Warp Index | 所属 Warpgroup | Role | +|---:|---|---| +| 0–3 (4 warps, 128 threads) | WG0 | Dispatch | +| 4–7 (4 warps, 128 threads) | WG1 | TMA A+SFA / B+SFB / MMA issue / idle | +| 8–15 (8 warps, 256 threads) | WG2, WG3 | Math WGMMA + epilogue + combine | + +**精确 warp 统计: 16 warps = 512 threads = 4 warpgroups** | + +**Compact Frontend 模式 (dispatch+TMA 共享 WG0):** + +当 `kCompactFrontendWarpgroup` 为 true 时，dispatch (2 warps) + TMA (2 warps) 共享一个 128-thread warpgroup: + +| Config | Dispatch | TMA | Epilogue | Total Threads | Reg Budget | +|---|---|---|---|---|---| +| Compact (default) | 64 | 64 | 256 | 384 | 48+48+208→59,392 | +| Wide (≥128) | 128 | 128 | 256 | 512 | 48+40+208=64,512 | + +**Phase 策略（通过宏模板参数选择）:** + +| 模式 | 说明 | +|---|---| +| `kSerialNWarpgroups` | math warpgroups 串行处理 L1 再 L2 | +| `kWideNWarpgroups` | math warpgroups 使用更多 N block 并行 | +| `kFusedL1L2Warpgroups` | 同一 WG 同时处理 L1 和 L2（与 pr-323 类似） | +| `kUseMMASync` | 使用 MMA sync 路径（BLOCK_M=32 时启用） | +| `kCompactFrontendWarpgroup` | dispatch+TMA 共享 warpgroup | + +**寄存器分配（宏推导）:** + +``` +kNumDispatchRegisters = 48 +kNumNonEpilogueRegisters = kCompactFrontendWarpgroup ? 48 : 40 +kNumEpilogueRegisters = (kSerialNWarpgroups or kWideNWarpgroups) ? 256 + : ((kUseMMASync and BLOCK_M==32) ? 240 : 208) +``` + +Compact mode 下 non-epilogue 必须与 dispatch 使用相同的 register 数 (48)，因为它们共享同一个 warpgroup（WG0）。 + +**megamoe_sm90 branch 具体配置（来自 `csrc/jit_kernels/heuristics/mega_moe.hpp` 的 1025 行扩展启发式）:** + +| Topk | Tokens | BLOCK_M | BLOCK_N | BLOCK_K | Epilogue Threads | Cluster | Compact | Phase | 备注 | +|---|---|---|---|---|---|---|---|---|---|---| +| 2 | 所有 | 64 | 128 | 128 | 256 | 1 | Yes | Fused | | +| 8 | ≤128 | 32 | 128 | 128 | 128 | 2 | No | Serial NW | `kSerialNWarpgroups` 硬编码为 false，实际不可达 | +| 8 | ≤576 | 64 | 128 | 128 | 256 | 2 | No | Serial NW | 同上 | +| 8 | >576 | 64 | 256 | 128 | 256 | 2 | No | Serial NW | 同上 | +| 9+ | ≤128 | 32 | 128 | 128 | 128 | 2 | No | Serial NW | 同上 | +| 9+ | ≤576 | 64 | 128 | 128 | 256 | 2 | No | Serial NW | 同上 | +| 9+ | >576 | 64 | 256 | 128 | 256 | 2 | No | Serial NW | 同上 | + +## 代码review发现 + +### 高: 宏驱动模板实例化导致编译膨胀 + +每次调用 `sm90_fp8_mega_moe` 会通过宏展开 4 个 kernel 实例（`INSTANTIATE_KERNEL_WITH_PHASE_POLICY` × 1 fuse + 1 serial_N + 1 wide_N + 1 mma_sync）。每个实例有独立的 `__launch_bounds__` 和 register 分配，JIT 编译时间较长。 + +### 中: Compact frontend 下 dispatch 使用 48 reg/warp — 仍然偏紧 + +dispatch warp 需要 rank round-robin 选择、SF copy 等复杂操作。48 reg/warp 可能在某些 shape 下触发 register spilling，但由于与 TMA warp 共享 WG0，无法单独增加而不破坏 budget。 + +### 中: Scheduler 保留严格整除约束 + +PR352 分叉自 #316（早于 PR347），因此其 scheduler (`mega_moe.cuh:38`) 仍使用 `kNumExpertsPerRank % kNumExpertsPerWave == 0` 的严格约束。PR347 的放宽修复（`> 0 && <=`）未被合入。对非 2 的幂 per-rank expert 数的 shape，可能触发编译期断言失败。 + +### 低: Compact 模式 register 余量尚充足 + +Compact 模式实际 register 消耗为 48×64+48×64+208×256=**59,392**（文档之前误写为 64,512），占预算约 90.6%，有约 5K register 余量。与 "tight" 的描述不同，实际还有一定 headroom。 + +### 中: 启发式文件增长到 1025 行，可维护性下降 + +`heuristics/mega_moe.hpp` 混合了 SM100 和 SM90 路径，且 SM90 部分包含大量硬编码的 topk/token 分支表。建议拆分为 `mega_moe_sm90.hpp` + `mega_moe_sm100.hpp`（类似 pr-360 的做法）。 + +### 低: Serial NW 模式下 `kNumEpilogueRegisters=256` 可能溢出 + +256 reg/warp × 256 epilogue threads = 65536，超过 64K reg budget。需要确认是否有其他约束（如减少 dispatch 线程数）来保证不溢出。 + +## 正面评价 + +- 宏驱动架构灵活，一个 kernel body 支持 4 种执行策略 +- Compact frontend 优化了资源利用率（H100 上 dispatch 不需要独占 warpgroup） +- 多种 phase 策略覆盖了不同 token 量的最优执行路径 +- Serial NW 模式的 epilogue 获得 256 reg/warp，适合计算密集场景 + +## 建议检查清单 + +- 验证 Serial NW 256 reg/warp 配置不超标 +- 实测 compact frontend 下 dispatch warp 的 spilling 情况 +- 考虑拆分 heuristics 文件降低维护成本 +- 确认所有 4 种 phase 策略的 correctness 测试覆盖 diff --git a/megamoe-research-reports/pr357_green_context_split_megamoe_notes.md b/megamoe-research-reports/pr357_green_context_split_megamoe_notes.md new file mode 100644 index 0000000..f72f8bb --- /dev/null +++ b/megamoe-research-reports/pr357_green_context_split_megamoe_notes.md @@ -0,0 +1,513 @@ +# PR357 Green-Context Split MegaMoE 代码review报告 + +## 1. 范围 + +- Remote: `git@github.com:RayWang96/DeepGEMM.git` +- 分支: `split_mega_moe` +- 本地 worktree: `pr-357` +- HEAD: `bb837421b` +- 审查方式: 代码review，不跑实验 + +`pr-357` 分支包含两段独立改动： + +| 范围 | Commit | 主题 | 规模 | +|---|---|---|---| +| 核心 PR357 | `bb837421b` | Add green-context split-kernel MegaMoE | 13 files, +3803/-4 | +| 旁路改动 | `41d89ee6c` | Add FP16-weights variant of FP8 MQA logits kernel | 7 files, +668/-9 | + +核心工作是 **SM100 FP8/FP4 MegaMoE split pipeline**，不是 SM90/Hopper 实现。它把原本 fused MegaMoE 拆成 `dispatch_l1_swiglu`、`l2_combine`、`combine_reduce` 三个 kernel，并用 CUDA Runtime 13.1 green context 将 K1/K2 放到互斥 SM 分区中并发执行。 + +--- + +## 实现概述 + +### K1: `dispatch_l1_swiglu` (dispatch + L1 GEMM + SwiGLU) + +`deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/dispatch_l1_swiglu.cuh` (1013 lines) + +K1 将 fused kernel 的 dispatch + L1 + SwiGLU 各阶段整合到一个 kernel 中。 + +**warp 分工:** + +| Warp Index | 数量 | Threads | Role | 所属 Warpgroup | +|---:|---:|---:|---|---| +| 0..`kNumDispatchWarps-1` | 4 (128 threads) | 128 | Route-based dispatch: 扫描 topk_idx，写 route count/entry，NVLink pull token+SF | WG0 | +| `kNumDispatchWarps` | 1 | 32 | TMA A (l1_acts) + SFA loader | WG1 | +| `kNumDispatchWarps+1` | 1 | 32 | TMA B (l1_weights) + SFB loader | WG1 | +| `kNumDispatchWarps+2` | 1 | 32 | MMA issue warp (leader CTA only) | WG1 | +| `kNumDispatchWarps+3` | 1 | 32 | Idle (warpgroup 占位) | WG1 | +| 剩余 | ≥4 (≥128 threads) | 128–256 | Math UMMA + SwiGLU epilogue + FP8 quant + TMA store + L2 arrival mask write | WG2+ | + +**寄存器分配 (与 SM100 fused kernel 一致, dispatch=128, TMA=128, epilogue=256):** + +| Role | Reg/thread | Threads | 总 Reg | +|---|---|---|---| +| Dispatch | 48 | 128 (4 warps) | 6,144 | +| TMA/non-epilogue | 40 | 128 (4 warps) | 5,120 | +| Math/epilogue | 208 | 256 (8 warps) | 53,248 | +| **总计** | | **512 (16 warps)** | **64,512/64K** | + +**关键约束:** +- `kNumDispatchThreads % 128 == 0` → dispatch 独占整数个 warpgroup +- `kNumNonEpilogueThreads == 128` → TMA 恰好一个 warpgroup +- Route-based dispatch: K1 不物化完整 routed token pool，而是在 dispatch warp 中通过 route metadata 边拉取边写 L1 arrival count +- K1 epilogue 写 `l2_arrival_mask`（release-or），通知 K2 可以开始消费 + +### K2: `l2_combine` (L2 GEMM + Cross-Rank Combine Scatter) + +`deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/l2_combine.cuh` (660 lines) + +K2 等待 K1 的 L1 output 就绪，执行 L2 GEMM 并将结果 scatter 到目标 rank 的 combine buffer。K2 无 dispatch warp，线程布局从 K1 config 继承 non_epilogue + epilogue。 + +**warp 分工 (non_epilogue=128, epilogue=256):** + +| Warp Index | 数量 | Threads | Role | 所属 Warpgroup | +|---:|---:|---:|---|---| +| 0 | 1 warp | 32 | TMA A (l2_acts) + SFA loader | WG0 | +| 1 | 1 warp | 32 | TMA B (l2_weights) + SFB loader | WG0 | +| 2 | 1 warp | 32 | MMA issue warp | WG0 | +| 3 | 1 warp | 32 | Idle (warpgroup 占位) | WG0 | +| 4–11 | 8 warps | 256 | Math UMMA + BF16 NVLink scatter epilogue | WG1, WG2 | + +**精确 warp 统计: 12 warps = 384 threads = 3 warpgroups** + +**寄存器分配:** + +| Role | Reg/thread | Threads | 总 Reg | +|---|---|---|---| +| TMA/non-epilogue | 40 | 128 (4 warps) | 5,120 | +| Math/epilogue | 208 | 256 (8 warps) | 53,248 | +| **总计** | | **384 (12 warps)** | **58,368/64K** | + +**关键设计:** +- 使用独立的 `Kernel2L2Scheduler`，等待 `expert_recv_count_sum` 高位达到 `kKernel1SMs * kNumRanks` 才开始调度 +- 每个 pool block 等待 `l2_arrival_mask == expected_mask`（由 K1 epilogue 写入） +- K2 结束前执行 NVLink barrier，保证 K3 reduce 时 combine buffer 跨 rank 可见 + +### K3: `combine_reduce` (最终 Top-K Reduce) + +`deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/combine_reduce.cuh` (142 lines) + +**warp 分工 (kNumThreads=512, JIT 模板实例化时设置，覆盖默认值 256):** + +| Warp Index | 数量 | Threads | Role | +|---:|---:|---:|---| +| 0–15 | 16 warps | 512 | Top-K reduce: 读取 top-k combine partials (BF16)，FP32 accumulate，BF16 store | + +**精确 warp 统计: 16 warps = 512 threads** + +**寄存器分配 (轻量 kernel, register 压力低):** + +| Role | Reg/thread | Threads | 总 Reg | +|---|---|---|---| +| Reduce | ~32 | 512 (16 warps) | ~16,384 | + +K3 不使用 green context，运行在 primary context。它依赖 graph node dependency（等 K1/K2 graph node 结束）和 K2 的 NVLink barrier 保证数据可见。 + +### Green Context 资源划分 + +| Context | Kernel | SM 分配 (默认) | Cluster Dim | +|---|---|---|---| +| Green Context 0 | K1 `dispatch_l1_swiglu` | `kernel1_sms` = 96 | 2 | +| Green Context 1 | K2 `l2_combine` | `kernel2_sms` = 52 | 2 | +| Primary Context | K3 `combine_reduce` | 无限制 | 1 | + +--- + +## 2. 代码差异总览 + +### 2.1 核心 split MegaMoE 文件 + +| 文件 | 作用 | +|---|---| +| `deep_gemm/mega/__init__.py` | Python API，新增 split buffer 和 `SM100FP8FP4MegaMoESplitGraph` 包装类 | +| `csrc/apis/mega.hpp` | C++ binding，新增 split buffer size API 和 pybind graph class | +| `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_split.hpp` | JIT runtime、配置推导、green context 创建、CUDA graph 构建 | +| `deep_gemm/include/deep_gemm/layout/mega_moe_split.cuh` | 新增 `SplitWorkspace`，扩展 route-based dispatch metadata | +| `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/dispatch_l1_swiglu.cuh` | K1，路由 dispatch、L1 GEMM、SwiGLU、FP8 quant、L2 arrival mask | +| `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/l2_combine.cuh` | K2，等待 K1 output、L2 GEMM、跨 rank combine scatter | +| `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/combine_reduce.cuh` | K3，按 top-k 对 combine partial 做最终 reduce | +| `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/common.cuh` | split pipeline 的 `state` tensor offset 定义 | +| `csrc/jit/handle.hpp` | lazy-load CUDA 13.1 green context 和 graph node API | +| `deep_gemm/include/deep_gemm/common/utils.cuh` | 新增 rank/token peel helper，用于 route-based pool slot 映射 | +| `deep_gemm/include/deep_gemm/ptx/ld_st.cuh` | 新增 CUDA 13 `longlong4_32a` 兼容和部分 load/store helper | +| `tests/test_mega_moe_split.py` | correctness 和 perf 测试，对 fused kernel 做 bitwise 对比 | + +### 2.2 旁路 MQA 改动 + +| 文件 | 作用 | +|---|---| +| `csrc/apis/attention.hpp` | `weights.dtype == fp16` 时选择新 SM100 FP8 MQA logits kernel | +| `csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp` | 新增 `SM100FP8MQALogitsF16WeightsRuntime` | +| `deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits_f16_weights.cuh` | FP8 input + FP16 weights 的 SM100 2-CTA logits kernel | +| `tests/test_attention.py` | 增加 FP16 weights 路径覆盖 | + +--- + +## 3. Public API 变化 + +### 3.1 Split symmetric buffer + +`deep_gemm/mega/__init__.py` 中 `SymmBuffer` 增加 `split: bool = False`。当 `split=True` 时，buffer sizing 使用 `_C.get_symm_buffer_size_for_mega_moe_split`。 + +新增 API： + +```python +deep_gemm.get_symm_buffer_for_mega_moe_split(...) +``` + +与 fused buffer 暴露相同的 tensor view： + +```python +x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf +``` + +区别在于 split buffer 的 workspace 头部更大，因为需要 route-based dispatch metadata。输入、pool、combine 区域的语义与 fused 对齐，方便 correctness 对比。 + +### 3.2 Weight transform + +`transform_weights_for_mega_moe` 负责把权重变成 SM100 MegaMoE 需要的 layout： + +| 权重 | 变换 | +|---|---| +| L1 weight | gate/up interleave | +| L1 scale factor | gate/up interleave，再做 UTCCP 需要的 transpose | +| L2 weight | 不 interleave | +| L2 scale factor | 做 UTCCP transpose | + +### 3.3 Split graph object + +新增 Python 包装类： + +```python +deep_gemm.SM100FP8FP4MegaMoESplitGraph(...) +graph.replay() +graph.get_green_context_ids() +``` + +构造函数接受 `states`、`ys`、`sym_buffers`、L1/L2 weights、stats，以及 `kernel1_sms`、`kernel2_sms`、`reduce_sms` 等参数。实际 C++ 构造函数会 JIT 编译三个 kernel，创建两个 green context，并构建一个 CUDA graph。 + +--- + +## 4. SplitWorkspace 布局 + +`layout::SplitWorkspace` 继承原 fused `Workspace`，保留 barrier、expert send/recv count、expert recv count sum、L1 arrival count、L2 arrival mask 等公共区域，同时新增 route-based dispatch 区域。 + +新增区域包括： + +| 区域 | 用途 | +|---|---| +| `src_token_topk_idx[expert][rank][token]` | 记录每个 local expert 收到的来自哪个 rank/token/topk slot | +| `src_route_count[rank][token]` | 记录一个 token 在目标 rank 上命中了几条 top-k route | +| `src_route_entry[rank][token][topk]` | 多 route 时记录 packed route entry | +| `token_src_metadata[pool_token]` | K2 scatter 时恢复目标 rank、原 token、topk slot | + +这个设计把 dispatch 从“先复制完整 routed token pool”改成“写路由元数据，按 expert pool 顺序拉取 token”，使 K1 可以边 dispatch、边 L1、边产生 L2 input。 + +--- + +## 5. CUDA Graph 与 Green Context + +### 5.1 CUDA 13.1 依赖 + +`sm100_fp8_fp4_mega_moe_split.hpp` 只在 `CUDART_VERSION >= 13010` 下启用真实实现。低于 CUDA Runtime 13.1 时，`SM100FP8FP4MegaMoESplitGraph` 是 stub，构造和 replay 都直接报错。 + +`csrc/jit/handle.hpp` lazy-load 了以下 runtime symbol： + +| API | 用途 | +|---|---| +| `cudaDeviceGetExecutionCtx` | 获取 primary execution context | +| `cudaDeviceGetDevResource` | 查询 SM/workqueue resource | +| `cudaDevSmResourceSplit` | 把 SM 资源切成两个 group | +| `cudaGreenCtxCreate` | 创建 green context | +| `cudaGraphAddNode` | 添加带 context 的 graph node | +| `cudaGraphKernelNodeSetAttribute` | 设置 cluster dim | +| `cudaGraphLaunch` | replay graph | + +### 5.2 SM 资源划分 + +`create_green_contexts()` 读取 device SM resource，按 `kernel1_sms` 和 `kernel2_sms` 切成两个 SM group： + +| Context | Kernel | 资源 | +|---|---|---| +| `green_contexts_[0]` | K1 `dispatch_l1_swiglu` | `kernel1_sms` | +| `green_contexts_[1]` | K2 `l2_combine` | `kernel2_sms` | +| `primary_context_` | K3 `combine_reduce` | 未 green-context 限制 | + +K1/K2 都用 `cluster_dim=2`，并设置 `coscheduledSmCount=2`。workqueue resource 使用 `cudaDevWorkqueueConfigScopeGreenCtxBalanced`。 + +### 5.3 Graph node dependency + +`build_graph()` 的 dependency 关系： + +| Node chain | Dependency | +|---|---| +| K1 nodes | 多 buffer 时 K1 串行依赖前一个 K1 | +| K2 nodes | 多 buffer 时 K2 串行依赖前一个 K2 | +| 第一个 K3 node | 依赖最后一个 K1 和最后一个 K2 | +| 后续 K3 nodes | 串行依赖前一个 K3 | + +单 buffer 下语义直接：K1 和 K2 同时启动，K2 在 kernel 内 busy-wait K1 的 arrival mask，K3 等 K1/K2 graph node 都结束后 reduce。 + +多 buffer 下 K1 chain 和 K2 chain 可以重叠，但 K3 等所有 K1/K2 chain 结束后再开始，不是 per-buffer K3 逐个接在各自 K1/K2 后面。 + +--- + +## 6. 三段 Kernel 实现 + +### 6.1 K1: `dispatch_l1_swiglu` + +入口： + +```cpp +sm100_fp8_fp4_mega_moe_split_dispatch_l1_swiglu_impl(...) +``` + +K1 同时承担三件事： + +| 阶段 | 行为 | +|---|---| +| Dispatch metadata | 扫描 `topk_idx`，写 expert send count、route count、route entry、remote `src_token_topk_idx` | +| L1 GEMM | 按 expert pool 顺序拉取输入 token 和 SF，做 L1 FP8/FP4 UMMA | +| SwiGLU/quant | epilogue 中做 clamp、SwiGLU、乘 top-k weight、量化成 FP8，写 `l2_acts` 和 `l2_acts_sf` | + +关键机制： + +| 机制 | 说明 | +|---|---| +| route-based pull | dispatch warp 通过 route metadata 拉取 remote input token，不提前物化全部 token pool | +| duplicate route 去重 | 如果一个 token 对同一目标 rank 有多个 top-k route，用 `route_count` 和 packed route entry 复用一次 remote pull | +| expert count 同步 | 本地写 send count，跨 rank 用 `atomic_add_sys` 汇总 recv count | +| NVLink barrier | dispatch metadata 完成后，通过 `comm::nvlink_barrier` 保证各 rank 可见 | +| L1 arrival count | 每个 pool block 的 token 写入后，`l1_arrival_count` release-add，L1 GEMM warp 等待完整 block | +| L2 arrival mask | K1 epilogue 完成一个 L1 output block 后，对 `l2_arrival_mask` 做 release-or，通知 K2 可消费 | + +K1 的 shared memory 与 fused kernel 不同。JIT 侧新增 `get_mega_moe_split_kernel1_pipeline()`，把 dispatch send buffer、expert count、CD staging、barrier、TMA stage 全部纳入 SMEM sizing，避免复用 fused heuristic 导致 SMEM 不足。 + +### 6.2 K2: `l2_combine` + +入口： + +```cpp +sm100_fp8_fp4_mega_moe_split_l2_combine_impl(...) +``` + +K2 只做 L2 和 combine scatter： + +| 阶段 | 行为 | +|---|---| +| Scheduling | `Kernel2L2Scheduler` 等待 `expert_recv_count_sum` 的高位达到 `kKernel1SMs * kNumRanks` | +| Readiness wait | 每个 pool block 等待 `l2_arrival_mask == expected_mask` | +| L2 GEMM | TMA load `l2_acts/l2_weights` 和 scale factors，执行 SM100 MXF8F6F4 UMMA | +| Combine scatter | epilogue 将 BF16 L2 output 通过 `sym_buffer.map` 写回目标 rank 的 combine buffer | +| Before reduce barrier | K2 epilogue 结束后执行 NVLink barrier，保证 K3 reduce 前 combine buffer 已跨 rank 可见 | + +K2 的 scheduler 使用 `block_idx += kKernel2SMs` 方式分配 block，因此 K2 的工作分区与其 green-context SM 数绑定。 + +### 6.3 K3: `combine_reduce` + +入口： + +```cpp +sm100_fp8_fp4_mega_moe_split_combine_reduce_impl(...) +``` + +K3 是最简单的一段： + +| 阶段 | 行为 | +|---|---| +| Grid | 一个 CTA 对应一个 original token | +| Load top-k | 读取该 token 的 `topk_idx` | +| Reduce | 对 `combine_token_buffer[topk_slot][token]` 中的 BF16 partial 做 FP32 accumulate | +| Store | cast 回 BF16，写 `y[token]` | + +K3 依赖 graph node dependency 和 K2 内部 NVLink barrier 来保证数据可见。它不使用 green context，也没有被 `reduce_sms` 限制。 + +--- + +## 7. 同步与状态设计 + +### 7.1 Workspace 同步 + +| 同步对象 | Producer | Consumer | 作用 | +|---|---|---|---| +| `expert_send_count` | K1 dispatch | K1/K2 scheduler | 每个 expert 的发送 token 数 | +| `expert_recv_count_sum` | K1 dispatch | K1/K2 scheduler | 每个 local expert 的总接收 token 数和到达计数 | +| `l1_arrival_count` | K1 dispatch pull | K1 L1 GEMM | pool block 输入 token 是否完整 | +| `l2_arrival_mask` | K1 epilogue | K2 L2 GEMM | L1 output 的 N-block 是否完整 | +| `token_src_metadata` | K1 dispatch pull | K2 combine scatter | scatter 回原 rank/token/topk slot | + +### 7.2 `state` tensor + +`common.cuh` 定义了一个至少 7 个 `int32` 的 state tensor： + +| Offset | 名称 | 当前用途 | +|---:|---|---| +| 0 | `K1ReadyTasks` | 已定义，当前实现中未见核心调度使用 | +| 1 | `K1DoneBlocks` | 已定义，当前实现中未见核心调度使用 | +| 2 | `K2ClaimCounter` | 已定义，当前实现中未见核心调度使用 | +| 3 | `K2DoneTasks` | K2 epilogue 完成 task 后 atomicAdd | +| 4 | `K2DoneBlocks` | K2 kernel block 完成后 atomicAdd | +| 5 | `K3DoneElements` | K3 每 token 完成后 atomicAdd | +| 6 | `K2Checksum` | 已定义，当前实现中未见核心调度使用 | + +当前实现主要依赖 workspace arrival/count 和 graph dependency，`state` 更像是调试/进度计数和未来扩展预留。 + +--- + +## 8. Correctness 与测试策略 + +`tests/test_mega_moe_split.py` 是主要验证入口，本文未运行。 + +测试逻辑： + +| 步骤 | 内容 | +|---|---| +| 构造输入 | FP8 activation、FP4 L1/L2 weights、相同 top-k routing | +| 分配 buffer | fused buffer 和 split buffer 分开分配 | +| 跑 fused | 调用 `deep_gemm.fp8_fp4_mega_moe` 得到 reference | +| 跑 split | 构造 `SM100FP8FP4MegaMoESplitGraph` 后 `replay()` | +| correctness | `torch.equal(y_split, y_fused)`，要求 bitwise identical | +| perf | 对 fused CUDA graph 和 split graph 做 best-of-N wall-clock，对比 TFLOPS/HBM/NVLink 粗略指标 | + +默认参数： + +| 参数 | 默认值 | +|---|---:| +| `num_processes` | 8 | +| `num_tokens` | 8192 | +| `hidden` | 7168 | +| `intermediate_hidden` | 3072 | +| `num_experts` | 384 | +| `num_topk` | 6 | +| `kernel1_sms` | 96 | +| `kernel2_sms` | 52 | +| `reduce_sms` | 148 | + +测试注释声称 split pipeline 通常略快于 fused，但这是代码注释，非本文实验结论。 + +--- + +## 9. 代码review发现与潜在问题 + +### 9.1 `reduce_sms` 和 `reduce_work_iters` 当前未实际生效 + +Python docstring 说 K3 在 `reduce_sms` 上运行，但 C++ `build_graph()` 中 K3 node 使用： + +```cpp +add_kernel_node(primary_context_, kernel3_graph_kernel_, num_tokens_, 512, 0, 1, ...) +``` + +具体情况： + +| 参数 | 当前情况 | +|---|---| +| `reduce_sms` | 构造时保存并做正数检查，但未用于 K3 grid 或 green context | +| `reduce_work_iters` | 构造时保存，但未传给 K3，也未影响调度 | +| K3 SM 限制 | 不使用 green context，运行在 primary context | + +这不一定是 bug，可能是 API 预留或尚未完成的功能。但文档和实现存在不一致，后续使用不要把 `reduce_sms` 当成有效的 SM partition 控制。 + +### 9.2 K1/K2 依赖 busy-wait arrival，异常路径可能死锁 + +K1/K2 之间的真实数据流依赖 workspace 中的 count/mask，而非 graph dependency。K2 会等待 `expert_recv_count_sum` 和 `l2_arrival_mask`。如果 K1 提前失败、route metadata 被破坏、或某个 arrival mask 没有写满，K2 可能长时间 spin 无法退出。 + +### 9.3 Split buffer 比 fused buffer 更大，且调用者必须显式 reset + +SplitWorkspace 增加了 route metadata 区域，buffer size 大于 fused。测试里每次 replay 前都会： + +```python +split_buffer.buffer.zero_() +state.zero_() +fill(split_buffer) +``` + +Public graph API 本身不自动 reset state/workspace。调用者如果 replay 多次，需要自行清零并重新填输入，否则 arrival/count/mask 会残留。 + +### 9.4 Graph construction 固化 shape 和指针 + +`SM100FP8FP4MegaMoESplitGraph` 在构造时固化： + +| 固化对象 | 后果 | +|---|---| +| `num_tokens`、hidden、intermediate、experts、topk | 不适合动态 shape 直接复用 | +| TMA descriptors | weights/buffer pointer 和 layout 固定 | +| green contexts | 生命周期绑定 graph object | +| graph node params | 多 buffer 列表长度固定 | + +动态 token 数或不同权重需要重新构造 graph。 + +### 9.5 SM100 only + +所有新增 split kernels 都要求 `__CUDA_ARCH__ >= 1000`，且 graph 实现要求 CUDA Runtime 13.1+。这不是 SM90/H100/H200 路径。 + +--- + +## 10. 与已有 MegaMoE split 方向的关系 + +pr357 的 split 不是简单的"L1 kernel 后接 L2 kernel"。其关键差异是 **green-context concurrency**： + +| 维度 | 常规 split L1/L2 | pr357 green-context split | +|---|---|---| +| K1/K2 launch | 两个 kernel 串行或外部并发 | 一个 CUDA graph 内两个 green context node | +| SM partition | 默认由调度器抢占/共享 | K1/K2 分别绑定固定 SM resource group | +| K2 等待方式 | 通常等 L1 kernel 完成 | K2 kernel 先启动，按 block 等 K1 arrival mask | +| 数据流 | L1 全量完成后 L2 消费 | K1 产生 pool block，K2 可逐 block 消费 | +| K3 | split 后单独 reduce | graph dependency 等 K1/K2 后 reduce | + +目标是重叠 dispatch/L1 和 L2/combine，并用固定 SM 切分减少 K1/K2 互相抢占。 + +--- + +## 11. FP16-weight MQA logits 旁路改动 + +这部分不是 split MegaMoE 主线，但在 `pr-357` 分支内。 + +### 11.1 API dispatch + +`fp8_fp4_mqa_logits` 中 `weights.scalar_type() == torch::kFloat16` 时选择新 kernel： + +```cpp +sm100_fp8_mqa_logits_f16_weights(...) +``` + +约束： + +| 约束 | 原因 | +|---|---| +| `arch_major == 10` | 只支持 SM100 | +| `seq_len % 4 == 0` | 2-CTA kernel 的 query tiling 没有 per-row bound check | +| 不支持 FP4 Q/KV | 代码显式 assert `not (is_fp4 and weights_is_f16)` | +| FP16 accumulation | QK score 和 per-head weighted sum 都用 FP16 accumulation，速度快但更容易溢出 | + +### 11.2 Kernel 结构 + +`sm100_fp8_mqa_logits_f16_weights.cuh` 是一个 2-CTA cluster kernel： + +| 组件 | 说明 | +|---|---| +| Q/KV | FP8 E4M3 | +| weights | FP16 | +| KV scale | FP32 | +| SMEM pipeline | Q、weights、KV、KV scale、KV offsets 多 stage | +| TMEM | 使用 SM100 tensor memory accumulator | +| Launch | `num_specialized_threads=128`，`num_math_threads=256`，`cluster_dim=2` | + +这是 attention/MQA 路径优化，不影响 split MegaMoE API。 + +--- + +## 12. 结论 + +pr357 的核心价值是引入了一个面向 SM100/CUDA 13.1 的 **green-context split-kernel MegaMoE execution model**。它不是单纯把 fused MegaMoE 拆成三个 kernel，而是让 K1 和 K2 在一个 CUDA graph 中用不同 green context 同时运行，借助 workspace arrival mask 实现 producer/consumer 流水线。 + +最重要的实现要点： + +| 结论 | 说明 | +|---|---| +| K1/K2 真正并发 | green context 将 `kernel1_sms` 和 `kernel2_sms` 切成独立 SM partition | +| 数据流按 pool block 流水化 | K1 写 `l2_arrival_mask`，K2 等 mask 后消费 | +| dispatch 变为 route-based pull | `SplitWorkspace` 保存 route metadata，减少重复 remote token pull | +| correctness 目标是 bitwise fused | 测试直接要求 `torch.equal(y_split, y_fused)` | +| API 仍有预留/不一致 | `reduce_sms` 和 `reduce_work_iters` 当前未控制 K3 | +| 分支含旁路 MQA 优化 | FP16-weight MQA logits 是独立 sidecar，不属于 MegaMoE split 主链路 | diff --git a/megamoe-research-reports/pr360_sm90_pingpong_coop_megamoe_code_review.md b/megamoe-research-reports/pr360_sm90_pingpong_coop_megamoe_code_review.md new file mode 100644 index 0000000..5850fe5 --- /dev/null +++ b/megamoe-research-reports/pr360_sm90_pingpong_coop_megamoe_code_review.md @@ -0,0 +1,375 @@ +# PR360 SM90 Pingpong/Cooperative MegaMoE 代码review报告 + +## 范围 + +- Worktree: `pr-360` +- Remote 分支: `pr-360/sm90-mega-moe-pingpong-coop` +- HEAD: `f27fd561164919925bf49aa2d015a1de5ec281db` +- 审查 diff 范围: `88965b078..f27fd5611` +- 变更规模: 13 个文件，+6200 / -7 行 +- 审查方式: 代码 review + 后续 8×H200 编译/benchmark 补测；未运行 correctness test 或 NCU profile + +## Commit 列表 + +- `f4f8473cd` sm90 megamoe: add pingpong and cooperative FP8 kernels +- `3f10f93a1` sm90 megamoe: token-count dispatch and Python bindings +- `89e3d351e` sm90 megamoe: correctness test and performance benchmark +- `7b65ae9b4` chore: ignore profiling artifacts (*.ncu-rep, *.nsys-rep, profile_logs/) +- `ba270a419` fix mul2 +- `0e6549b60` add sm90 megamoe bench +- `f27fd5611` Revert "chore: ignore profiling artifacts (*.ncu-rep, *.nsys-rep, profile_logs/)" + +## 架构概述 + +- 为 Hopper/SM90 架构新增 FP8 MegaMoE kernel，与现有 SM100 FP8/FP4 MegaMoE 路径独立。 +- `pingpong` 使用 `BLOCK_M=64`，两个 math warpgroup 分别处理不同 tile，通过 order barrier 交替执行 MMA 和 epilogue。 +- `cooperative` 使用 `BLOCK_M=128`，两个 math warpgroup 按 M 行协作切分同一 tile，共享一次 B/weight TMA load，将权重 HBM 读取量减半。 +- 两者均采用 persistent CTA-per-SM launch，共 64 dispatch threads + 64 TMA/non-epilogue threads + 256 math/epilogue threads。 +- 新增 public API: `fp8_mega_moe`、`fp8_mega_moe_pingpong`、`fp8_mega_moe_cooperative`、`transform_weights_for_mega_moe_sm90`。 + +## 8×H200 实验验证补充 + +本节为后续补测结果，目的是验证 PR 描述中的性能表是否能在单节点 8×H200 上复现。报告更新本身未重新访问远端机器。 + +**运行环境与准备:** + +- 机器: 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`。 +- PR360 `_C` 与 `/sgl-workspace/DeepEP/deep_ep_cpp` 均按 Torch 2.12/CUDA 13.0 临时重建；运行后默认 `.so` 已恢复为 base `torch 2.9.1+cu130` 可 import 状态，Torch 2.12 构建产物另存为 `.torch212cu130.so`。 +- 为适配 PyTorch 2.12 的 tensor metadata 限制，本地验证版本在 `csrc/apis/mega.hpp` 的 `slice_input_buffers` lambda 中加入 `torch::NoGradGuard no_grad;`。这是验证补丁，不属于原始 PR360 diff。 +- 环境变量: `NVSHMEM_IBGDA_ENABLE=0 NVSHMEM_DISABLE_IBGDA=1 EP_DISABLE_GIN=1 NCCL_NVLS_ENABLE=0`，并用 `PYTHONPATH=/sgl-workspace/DeepEP:$PYTHONPATH` 加载重建后的 DeepEP。 +- `NCCL_NVLS_ENABLE=0` 是必要项；否则 Torch 2.12/NCCL 2.29 会尝试 NVLS multicast，并报 `Failed to bind NVLink SHARP (NVLS) Multicast memory ... CUDA error 401`。 +- PyTorch symmetric memory multicast 仍报 `the operation cannot be performed in the present state`，但 Torch 2.12 会 `Gracefully skipping multicast initialization` 后走 fallback；因此本次验证不代表 CUDA multicast/NVLS 路径可用。 + +**Baseline 覆盖限制:** + +- `v1-contig` 和 `v1-ll` 已验证。 +- `v2 ElasticBuffer` 未能验证，因为当前容器内 DeepEP API 没有 `deep_ep.ElasticBuffer`；所有 PR 描述中关于 v2 的加速比本次未独立确认。 + +**DeepSeek-V4 Flash** (`hidden=4096, intermediate=2048, experts=256, topk=6`) + +| tokens | fused (us) | TFLOPS | HBM GB/s | v1-contig (us) | v1-ll (us) | +|---:|---:|---:|---:|---:|---:| +| 16 | 241.9 | 19.4 | 3232 | 1685.7 | 264.3 | +| 64 | 251.4 | 70.5 | 3226 | 1698.2 | 290.6 | +| 256 | 317.6 | 245.0 | 2616 | 2213.3 | 349.7 | +| 512 | 336.8 | 464.4 | 2542 | 1670.8 | 459.3 | +| 1024 | 586.8 | 527.0 | 1544 | 2125.0 | 711.5 | +| 4096 | 1799.1 | 684.6 | 670 | 4235.9 | skipped | +| 8192 | 3579.6 | 689.2 | 449 | 6951.4 | skipped | + +相对 PR 描述表中的 fused 时间，偏差范围约 `-3.0% ~ +0.5%`，可视为复现通过。 + +**DeepSeek-V4 Pro** (`hidden=7168, intermediate=3072, experts=384, topk=6`) + +| tokens | fused (us) | TFLOPS | HBM GB/s | v1-contig (us) | v1-ll (us) | +|---:|---:|---:|---:|---:|---:| +| 16 | 743.7 | 17.2 | 3912 | 2674.9 | 809.4 | +| 64 | 818.6 | 59.4 | 3886 | 2782.4 | 902.4 | +| 256 | 1015.4 | 197.5 | 3164 | 2688.5 | 1002.7 | +| 512 | 1063.0 | 387.1 | 3064 | 2750.7 | 1181.7 | +| 1024 | 1622.7 | 510.6 | 2061 | 3286.3 | 1604.4 | +| 4096 | 4684.3 | 689.4 | 821 | 7250.6 | skipped | +| 8192 | 10048.5 | 644.7 | 450 | 12512.8 | skipped | + +相对 PR 描述表中的 fused 时间，偏差范围约 `-6.1% ~ -0.3%`。趋势与 PR 描述一致：大 batch 下 fused 仍明显快于 v1-contig，但本次不能确认其相对 v2 的回退幅度。 + +**MiMo-V2.5** (`hidden=4096, intermediate=2048, experts=256, topk=8`) + +| tokens | fused (us) | TFLOPS | HBM GB/s | v1-contig (us) | v1-ll (us) | +|---:|---:|---:|---:|---:|---:| +| 16 | 246.7 | 25.1 | 3273 | 1684.0 | 271.4 | +| 64 | 248.7 | 97.5 | 3269 | 1637.8 | 299.6 | +| 256 | 323.7 | 326.4 | 2594 | 1597.7 | 386.6 | +| 512 | 469.9 | 443.2 | 1858 | 1900.1 | 541.6 | +| 1024 | 731.6 | 561.1 | 1283 | 2861.3 | 889.9 | +| 4096 | 2322.2 | 708.9 | 578 | 4933.5 | skipped | +| 8192 | 4677.9 | 703.4 | 401 | 8446.8 | skipped | + +相对 PR 描述表中的 fused 时间，偏差范围约 `-2.6% ~ +0.7%`，复现度较高。 + +**MiMo-V2.5-Pro** (`hidden=6144, intermediate=2048, experts=384, topk=8`) + +| tokens | fused (us) | TFLOPS | HBM GB/s | v1-contig (us) | v1-ll (us) | +|---:|---:|---:|---:|---:|---:| +| 16 | 471.1 | 20.5 | 3371 | 2210.8 | 518.8 | +| 64 | 502.5 | 73.9 | 3628 | 2375.6 | 578.2 | +| 256 | 642.6 | 242.4 | 2892 | 2331.6 | 677.7 | +| 512 | 675.7 | 454.0 | 2817 | 2232.6 | 857.7 | +| 1024 | 1226.9 | 502.9 | 1627 | 3110.1 | 1365.8 | +| 4096 | 3663.4 | 680.5 | 698 | 6787.8 | skipped | +| 8192 | 7482.0 | 663.3 | 440 | 11709.8 | skipped | + +相对 PR 描述表中的 fused 时间，偏差范围约 `-3.7% ~ -1.1%`，复现度较高。 + +**实验结论:** + +- PR360 fused kernel 的主性能结论基本成立：四个公开 shape 的 fused 时间均能在 8×H200 上复现，偏差通常在几个百分点内。 +- `v1-ll` 与 PR 描述基本吻合；`v1-contig` 波动较大，部分点偏差超过 10%，其中 DeepSeek-V4 Flash `tokens=256` 明显偏慢（`2213.3 us` vs PR 表 `1525.0 us`）。这不影响 fused kernel 复现结论，但说明 baseline 数值对 DeepEP/TileLang 版本和运行环境更敏感。 +- 由于 `ElasticBuffer` 缺失，本次不能验证 PR 描述中“相对 DeepEP V2”的结论；合入前若要以 v2 对比作为依据，需要在包含 `deep_ep.ElasticBuffer` 的 DeepEP 版本上重跑。 +- 本次所有 fused/baseline 结果均是在 symmetric-memory multicast fallback 条件下测得；如果目标环境能正确启用 NVLS/multicast，性能可能不同。 + +**实验日志:** + +- DeepSeek-V4 Flash: `/mnt/beegfs/lxy/logs/pr360_h200_deepseek_flash_h4096_i2048_e256_topk6_baseline_v1_torch212.log` +- DeepSeek-V4 Pro: `/mnt/beegfs/lxy/logs/pr360_h200_deepseek_pro_h7168_i3072_e384_topk6_baseline_v1_torch212.log` +- MiMo-V2.5: `/mnt/beegfs/lxy/logs/pr360_h200_mimo_v25_h4096_i2048_e256_topk8_baseline_both_torch212.log` +- MiMo-V2.5-Pro: `/mnt/beegfs/lxy/logs/pr360_h200_mimo_v25_pro_h6144_i2048_e384_topk8_baseline_v1_torch212.log` + +## 实现概述 + +### Pingpong Kernel (`BLOCK_M=64`) + +`deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh` (1418 lines) + +**warp 分工与 WGMMA 流水线:** + +| Warp Index | 数量 | Threads | Role | 所属 Warpgroup | +|---:|---:|---:|---|---| +| 0, 1 | 2 (64 threads) | 64 | Dispatch: NVLink pull token+SF, topk weight copy, TMA store to L1 pool | WG0 | +| 2 | 1 | 32 | TMA A (acts) + SFA loader | WG0 | +| 3 | 1 | 32 | TMA B (weights) loader | WG0 | +| 4–7 | 4 (128 threads) | 128 | Math WG0: m64 WGMMA → pingpong epilogue → TMA store L1 / BF16 scatter | WG1 | +| 8–11 | 4 (128 threads) | 128 | Math WG1: m64 WGMMA → pingpong epilogue → TMA store L1 / BF16 scatter | WG2 | + +**OrderedSequenceBarrier pingpong 时序:** + +``` + WG1: [MMA tile 0] [epilogue tile 0] [MMA tile 2] [epilogue tile 2] ... + WG2: [MMA tile 1] [epilogue tile 1] [MMA tile 3] [epilogue tile 3] ... +``` + +同一个 tile 的 MMA 和 epilogue 在不同 WG 上交替执行，通过 `order_arrive`/`order_wait` 门控 stage 0 (MMA) 和 stage 1 (epilogue) 的访问权。 + +**寄存器分配:** + +| Role | Reg/thread | Threads | 总 Reg | +|---|---|---|---| +| Dispatch | 48 | 64 | 3,072 | +| TMA (A loader) | 48 | 32 | 1,536 | +| TMA (B loader) | 48 | 32 | 1,536 | +| Math WG0 | 224 | 128 | 28,672 | +| Math WG1 | 224 | 128 | 28,672 | +| **总计** | | **384** | **63,488/64K** | + +**关键约束:** +- `kNumDispatchThreads + kNumNonEpilogueThreads == 128` → dispatch (64) + TMA (64) 形成一个完整的 WG0 +- `kNumEpilogueWarpgroups == 2` → 恰好两个 math warpgroup 参与 pingpong +- WG0 内所有 warp 使用相同的 dealloc count (48)，与 SM100 生产代码一致的模式 +- BLOCK_M=64，每个 math WG 拥有**整个 tile**（一次 m64 WGMMA 完成） + +### Cooperative Kernel (`BLOCK_M=128`) + +`deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh` (1419 lines) + +**warp 分工:** + +| Warp Index | 数量 | Threads | Role | 所属 Warpgroup | +|---:|---:|---:|---|---| +| 0, 1 | 2 (64 threads) | 64 | Dispatch (同 pingpong) | WG0 | +| 2 | 1 | 32 | TMA A (acts) + SFA loader | WG0 | +| 3 | 1 | 32 | TMA B (weights) loader | WG0 | +| 4–7 | 4 (128 threads) | 128 | Math WG0: 处理 tile 的 rows [0, 64) | WG1 | +| 8–11 | 4 (128 threads) | 128 | Math WG1: 处理 tile 的 rows [64, 128) | WG2 | + +**WG 协作模式:** + +``` + 同一 tile (BLOCK_M=128): + WG1: owns rows [0, 64) → m64 WGMMA on smem_a[0:64] + WG2: owns rows [64, 128) → m64 WGMMA on smem_a[64:128] + + B/weight: 共享一次 TMA load → 两个 WG 各自读取同一 smem_b +``` + +两个 WG 各自执行自己的 epilogue，通过 256-thread barrier (`kEpilogueFullBarrierIdx`) 在 L2 phase 同步，防止一个 WG 的下一 tile 写入覆盖另一个 WG 仍在 scatter 的本 tile 数据。 + +**寄存器分配 (与 pingpong 完全一致):** + +| Role | Reg/thread | Threads | 总 Reg | +|---|---|---|---| +| Dispatch | 48 | 64 | 3,072 | +| TMA (A+B loader) | 48 | 64 | 3,072 | +| Math WG0 | 224 | 128 | 28,672 | +| Math WG1 | 224 | 128 | 28,672 | +| **总计** | | **384** | **63,488/64K** | + +**关键设计:** +- WG 通过 `row_block_offset = epilogue_wg_idx * WG_BLOCK_M` 偏移 smem_a/smem_sfa/smem_cd 的读写位置 +- L1 epilogue: 两个 WG 各自 TMA store 自己 64-row 带，使用 `block_m / 2 = 64` 的 descriptor box +- L2 epilogue: 跨 WG 256-thread barrier 防止 smem_cd 别名冲突 +- Cooperative 适合大 M：B/weight 仅 load 一次，weight HBM 读数减半 + +### Token-Count Auto Routing + +`fp8_mega_moe` 按 token 数自动选择 kernel: + +| num_tokens | 选择 kernel | 适用场景 | +|---:---|---| +| < `DG_SM90_MOE_COOPERATIVE_THRESHOLD` (默认 256) | pingpong | 小/中等 M，epilogue 延迟主导 | +| ≥ threshold | cooperative | 大 M，weight 重复读取主导 | + +也可通过 `DG_SM90_MOE_KERNEL=pingpong|cooperative` 强制选择。 + +## 代码review发现 + +### ~~严重: `setmaxnreg` 在 warpgroup 分歧分支中执行~~ [代码形态与 SM100 同构，风险降低] + +初步审阅时怀疑 SM90 kernel 将 WG0 的 4 个 warp 拆分在不同 `if/else if` 分支中调用 `warpgroup_reg_dealloc`，可能违反 PTX 对 warpgroup 一致性的要求。经与主仓库 SM100 fused kernel（`sm100_fp8_fp4_mega_moe.cuh`）交叉验证，确认该代码形态与生产代码同构，实际风险低于最初判断；但不应将其概括为“PC 地址无关”。 + +关键事实： + +| | SM100 fused（主仓库，生产验证） | SM90 PR360 | +|---|---|---| +| WG 划分 | TMA warps (4个) 独占 WG1 | dispatch (2) + TMA (2) 共享 WG0 | +| dealloc 调用方式 | 4 个 TMA warp 分别从 4 个不同的 `else if` 分支调用 `dealloc<40>` | 4 个 warp 分别从不同的 `if/else if` 分支调用 `dealloc<48>` | +| 同一 WG 内目标指令与 operand 是否一致？ | 是，均为 `setmaxnreg.dec.sync.aligned.u32 40` | 是，均为 `setmaxnreg.dec.sync.aligned.u32 48` | + +PTX 规范要求 warpgroup 内所有 warp 在收敛的 warpgroup 上执行相同的 `setmaxnreg` 指令（相同 opcode + operand，且满足 `.sync.aligned` 的同步语义）。这里能降低风险的依据是：分支条件按 warp role 决定，所有相关分支最终发出相同 operand 的 `setmaxnreg.dec.sync.aligned`，并且 SM100 生产路径使用同构模式。严格来说，是否满足“同一动态 warpgroup 指令”的约束仍应以生成 SASS/实机测试为准，不能把“源码中不同分支但同 operand”泛化成任意 PC 都安全。当前 8×H200 benchmark 已覆盖该路径的运行稳定性，但未替代专门的 SASS 审核或 sanitizer 验证。 + +### 高: 波次大小启发式可能违反调度器的整除约束 + +SM90 调度器在编译期要求 `kNumExpertsPerWave` 精确整除 `kNumExpertsPerRank`。复用的 host 启发式按尾部比例搜索最佳值，但不限制候选值为除数。 + +证据: + +- `deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:37` 包含 `DG_STATIC_ASSERT(kNumExpertsPerRank % kNumExpertsPerWave == 0, "Invalid wave config")`。 +- `csrc/jit_kernels/heuristics/sm90_mega_moe.hpp:116-124` 将 SM90 波次选择委托给 `get_num_experts_per_wave_for_mega_moe`。 +- `csrc/jit_kernels/heuristics/mega_moe.hpp:129-144` 在 `[min_num_experts_per_wave, min_num_experts_per_wave * 2]` 内按尾部比例搜索并返回 `best_num_experts_per_wave`，不筛选除数。 +- `csrc/jit_kernels/impls/sm90_fp8_mega_moe_pingpong.hpp:101-105` 和 `csrc/jit_kernels/impls/sm90_fp8_mega_moe_cooperative.hpp:101-105` 将启发式结果传入 kernel 模板。 + +影响: + +- 合法的 public API shape 可能因 `Invalid wave config` 在 JIT 编译时失败。 +- 现有测试多使用 2 的幂次 per-rank expert 数，可能掩盖问题。API 本身仅检查 `num_experts == num_experts_per_rank * num_ranks`，不限制 per-rank expert 为 2 的幂或其他与启发式兼容的值。 + +建议修复: + +- 约束启发式只返回 `num_experts_per_rank` 的除数，或放松调度器不变式以支持尾部非完整波次。 +- 增加非 2 的幂 `num_experts_per_rank` 的针对性测试。 + +### 中: 调度器在波次边界读取越界 expert count + +`advance_expert_idx()` 递增当前 expert 并立即调用 `get_num_tokens(current_local_expert_idx)`。当 consume 完一个波次的最后一个 expert 后，会调用 `get_num_tokens(wave_end)`。如果 `wave_end == kNumExpertsPerRank` 且 per-rank expert 数是 32 的倍数，没有 lane 拥有该 expert index，`ptx::exchange` 的源值为未初始化。 + +证据: + +- `deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:73-81` 声明 `uint32_t valid_value` 且未初始化，仅在 lane 的缓存 expert index 匹配 `expert_idx` 时才赋值。 +- `deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:94-98` 递增 `current_local_expert_idx` 并无条件调用 `get_num_tokens`。 +- `deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:126-137` 和 `141-163` 在完成一个 expert 后调用 `advance_expert_idx()`，包括波次最后一个 expert。 +- `deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh:61-63` 为每个 `(i * 32 + lane_idx)` expert 缓存一个值。当 `kNumExpertsPerRank == 32` 时 `kNumExpertsPerLane == 1`，expert index 32 无人拥有。 + +影响: + +- 在 32 个 local expert 等常见 shape 下，设备调度器使用未初始化寄存器值。 +- 该值通常在波次切换后被覆盖，不一定会造成可见的错误输出，但仍是未定义行为，在 sanitizer 或编译器变化时可能暴露。 + +建议修复: + +- `get_num_tokens` 中初始化 `valid_value = 0`。 +- `advance_expert_idx` 中仅当 `current_local_expert_idx < kNumExpertsPerRank` 时调用 `get_num_tokens`，否则将 `current_num_tokens` 设为 0。 + +### 中: 中间 FP8 量化 scale 语义在 kernel、测试、baseline 间不一致 + +fused kernel 对 L1/SwiGLU 输出使用 2 的幂次 scale 进行量化，但 correctness reference 和 benchmark baseline 使用精确的 `amax / 448.0` float scale。导致 correctness 检查与性能 baseline 比较的是不同的数值管线。 + +证据: + +- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh:1125-1135` 通过 `math::get_e4m3_sf_and_sf_inv` 计算 L1 输出 scale。 +- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh:1115-1125` 同样如此。 +- `deep_gemm/include/deep_gemm/common/math.cuh:98-106` 将 `get_e4m3_sf_and_sf_inv` 实现为 `2 ** ceil(log2(amax / 448))`，并静态断言 UE8M0 风格行为。 +- `tests/test_mega_moe_sm90.py:231-238` 在 reference L1 quant/dequant 路径中使用精确的 `sf2 = amax / 448.0`。 +- `tests/bench_mega_moe_sm90.py:323-325` 在 V1 contiguous baseline 中使用 `use_ue8m0_scale=False`。 +- `tests/bench_mega_moe_sm90.py:592-599` 在 V1 low-latency baseline 中使用 `use_ue8m0_scale=False`。 + +影响: + +- `calc_diff` 测量的并非 fused kernel 与数值等价 reference 的差异。 +- 如果非融合 baseline 使用不同的 L1 量化 scale 策略，benchmark 加速比不是严格同类对比。 +- 较宽松的 `diff_tol=0.01` 可能掩盖此问题，但会削弱测试信号并隐藏 scale 敏感的回归。 + +建议修复: + +- 明确 SM90 中间 activation scale 应使用精确 float 还是 2 的幂次 float。 +- 若采用精确 float，将 L1 epilogue 中的 `get_e4m3_sf_and_sf_inv` 替换为精确的 `amax / 448.0` 及其倒数。 +- 若采用 2 的幂次，更新 PyTorch reference 和所有 benchmark baseline 使用相同的幂次 scale，并文档说明 float tensor 中存储 UE8M0 等效值。 + +### 中: Phase-profiling benchmark 路径与 API 不兼容且 kernel 未实现 + +benchmark 中 `DG_SM90_MOE_PHASE_PROFILE` 路径在 per-expert stats 后追加 64 个 int，但所有 public SM90 wrapper 要求 stats tensor 长度精确等于 `num_experts_per_rank`。kernel 中也仅写入 per-expert recv count，不包含 phase metric 写入。 + +证据: + +- `tests/bench_mega_moe_sm90.py:712-714` 在 `DG_SM90_MOE_PHASE_PROFILE` 被设置时分配 `num_experts_per_rank + 64` 个 int。 +- `tests/bench_mega_moe_sm90.py:810-828` 从追加的 64-int 区域读取 phase metrics。 +- `csrc/apis/mega.hpp:321-324`、`450-453`、`555-558` 断言 `cumulative_local_expert_recv_stats->numel() == num_experts_per_rank`。 +- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh:602-603` 和 `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh:610-611` 仅累加 per-expert recv count。 +- 在整个 PR 中搜索，未发现 `DG_SM90_MOE_PHASE_PROFILE` 或追加 phase slot 的 kernel 侧写入。 + +影响: + +- 启用文档中描述/环境变量驱动的 phase profile 路径会在 host API 断言处失败。 +- 即使放宽断言，打印的 phase metrics 也会是零或过期数据，因为 kernel 从未写入这些值。 + +建议修复: + +- 删除 benchmark phase-profile 路径，或实现独立的 stats tensor/API 合约及 kernel 侧 instrumentation。 + +### 中: SM90 weight scale layout 检查允许 kernel 无法正确索引的 layout + +API 接受两种 weight scale 布局: 连续的 `(E, N/128, K/128)` 或最后两维转置后连续。SM90 kernel 使用原始指针算术，假定 K 是连续的最内层维度。 + +证据: + +- `csrc/apis/mega.hpp:315-318`、`444-447`、`549-552` 使用 `check_sf_layout(..., sm90_sfb_check=true, torch::kFloat)` 检查 weight scale。 +- `csrc/utils/layout.hpp:109-114` 同时接受 `(stride(-1) == 1 and stride(-2) == size(-1))` 和 `(stride(-1) == size(-2) and stride(-2) == 1)`。 +- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_pingpong.cuh:858-867`、`913-918`、`1006-1010` 将 weight scale 视为以 K-block 为最快维度的连续原始数组进行索引。 +- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe_cooperative.cuh:847-857`、`903-908`、`996-1000` 使用相同索引方式。 +- `deep_gemm/mega/__init__.py:135-158` 未对 SM90 scale tensor 做任何变换，测试仅覆盖自然连续布局。 + +影响: + +- 转置 stride 的 scale tensor 可以合法通过校验，但在 kernel 中被按错误线性顺序解释。 +- 将产生错误的 dequantization scale，且不会触发 host 断言失败。 + +建议修复: + +- 对 SM90 原始指针 weight scale，要求严格自然连续性: group stride 为 `size(-2) * size(-1)`，`stride(-2) == size(-1)`，`stride(-1) == 1`。 +- 或者传入并在 kernel 中使用显式 stride，但对此路径可能得不偿失。 + +### 低: Benchmark HBM expert 计数在单 rank 或无过滤场景下可能少计 + +benchmark 将 unique expert 数减 1 以剔除插入的 `-1` 哨兵值。如果 tensor 中不存在 `-1`，则会少计一个 touched expert。 + +证据: + +- `tests/bench_mega_moe_sm90.py:784-792` 将非本地 expert 过滤为 `-1`，然后计算 `torch.unique(...).numel() - 1`。 +- 在单 rank 运行或无 route 被过滤/屏蔽的场景下，`-1` 可能不存在。 + +影响: + +- 这些场景下报告的 HBM GB/s 偏低。 +- 不影响 kernel 正确性。 + +建议修复: + +- 改为统计 `torch.unique(gathered_topk_idx[gathered_topk_idx >= 0]).numel()`。 + +## 正面评价 + +- PR 将 SM90 代码与现有 SM100 MegaMoE 路径分离，降低了跨架构回归风险。 +- Public API 同时提供自动路由和强制 `pingpong`/`cooperative` 变体，便于 A/B 测试。 +- 测试按分层结构组织，覆盖 smoke、启发式、shape、边界条件和随机压力场景。 +- cooperative kernel 在共享 `smem_cd` 重用处显式加入了跨 warpgroup barrier，正确处理了一类行带别名的数据竞争风险。 + +## 建议合入前检查清单 + +- 使 `num_experts_per_wave` 与调度器合约一致，并增加非 2 的幂 expert 数覆盖。 +- 修复调度器波次边界越界 token count 读取。 +- 统一 L1 中间 FP8 scale 语义（kernel / 测试 / baseline）。 +- 删除或实现 `DG_SM90_MOE_PHASE_PROFILE`。 +- 收紧 SM90 weight scale layout 检查。 +- 若 PR 描述继续引用 DeepEP V2 加速比，需在包含 `deep_ep.ElasticBuffer` 的 DeepEP 版本上重跑 `--baseline-version v2/both`。 +- 在 Fabric Manager / NVSwitch / NCCL NVLS multicast 正常的环境上复测一次，确认非 fallback symmetric-memory 路径下的性能变化。 +- 修复后，在 SM90/H100 或 H200 上运行 `tests/test_mega_moe_sm90.py` 的 layers 1-5，分别测试 `DG_SM90_MOE_KERNEL=pingpong` 和 `DG_SM90_MOE_KERNEL=cooperative`。 diff --git a/scripts/run_nsys_mega_moe_sm90.sh b/scripts/run_nsys_mega_moe_sm90.sh new file mode 100755 index 0000000..12018f6 --- /dev/null +++ b/scripts/run_nsys_mega_moe_sm90.sh @@ -0,0 +1,52 @@ +#!/bin/bash +set -e + +NUM_RANKS=${NUM_RANKS:-2} +NUM_TOKENS=${NUM_TOKENS:-128} +HIDDEN=${HIDDEN:-4096} +INTERMEDIATE=${INTERMEDIATE:-2048} +NUM_EXPERTS=${NUM_EXPERTS:-16} +NUM_TOPK=${NUM_TOPK:-6} +OUTPUT_DIR=${OUTPUT_DIR:-./nsys-results} + +mkdir -p "$OUTPUT_DIR" + +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +REPORT_NAME="nsys_sm90_megamoe_t${NUM_TOKENS}_r${NUM_RANKS}_${TIMESTAMP}" + +echo "=== SM90 MegaMoE nsys profiling ===" +echo " Ranks: $NUM_RANKS" +echo " Tokens: $NUM_TOKENS" +echo " Hidden: $HIDDEN" +echo " Intermediate: $INTERMEDIATE" +echo " Experts: $NUM_EXPERTS" +echo " TopK: $NUM_TOPK" +echo " Output: $OUTPUT_DIR/$REPORT_NAME" +echo "" + +NSYS_FLAGS=( + --trace=cuda,nvtx + --cuda-memory-usage=true + --force-overwrite=true + --output="$OUTPUT_DIR/$REPORT_NAME" + --stats=true +) + +TORCH_FLAGS=( + --nproc_per_node="$NUM_RANKS" +) + +TEST_FLAGS=( + --num-tokens="$NUM_TOKENS" + --hidden="$HIDDEN" + --intermediate-hidden="$INTERMEDIATE" + --num-experts="$NUM_EXPERTS" + --num-topk="$NUM_TOPK" +) + +nsys profile "${NSYS_FLAGS[@]}" \ + torchrun "${TORCH_FLAGS[@]}" \ + tests/test_mega_moe_sm90.py "${TEST_FLAGS[@]}" + +echo "" +echo "=== Done. Report: $OUTPUT_DIR/$REPORT_NAME.nsys-rep ===" diff --git a/tests/test_mega_moe_sm90.py b/tests/test_mega_moe_sm90.py new file mode 100644 index 0000000..3a0c04d --- /dev/null +++ b/tests/test_mega_moe_sm90.py @@ -0,0 +1,214 @@ +import argparse +import os +import random +import sys +import torch +import torch.distributed as dist +from typing import Tuple, Optional + +from deep_gemm.utils import per_token_cast_to_fp8 +from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather +from deep_gemm.testing import bench_kineto + + +def reference_swiglu(gate: torch.Tensor, up: torch.Tensor, + activation_clamp: Optional[float] = None) -> torch.Tensor: + if activation_clamp is not None: + gate = gate.clamp(-activation_clamp, activation_clamp) + up = up.clamp(-activation_clamp, activation_clamp) + return torch.nn.functional.silu(gate) * up + + +def reference_fp8_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + fp8_max = 448.0 + amax = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + sf = amax / fp8_max + x_scaled = x / sf + x_fp8 = x_scaled.to(torch.float8_e4m3fn) + return x_fp8, sf.squeeze(-1) + + +def reference_fp8_dequantize(x_fp8: torch.Tensor, sf: torch.Tensor, + sf_weights: torch.Tensor, + per_k_gran: int = 128) -> torch.Tensor: + x_f32 = x_fp8.to(torch.float32) + k = x_f32.shape[-1] + num_groups = k // per_k_gran + x_f32 = x_f32.reshape(*x_f32.shape[:-1], num_groups, per_k_gran) + sf_a = sf.unsqueeze(-1) + sf_b = sf_weights.unsqueeze(-2) + x_f32 = x_f32 * sf_a * sf_b + return x_f32.reshape(*x_f32.shape[:-2], k) + + +def reference_mega_moe( + x_bf16: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + l1_weights_bf16: torch.Tensor, + l2_weights_bf16: torch.Tensor, + num_experts_per_rank: int, + rank_idx: int, + num_ranks: int, + activation_clamp: Optional[float] = None, +) -> torch.Tensor: + num_tokens, hidden = x_bf16.shape + num_topk = topk_idx.shape[1] + intermediate_hidden = l2_weights_bf16.shape[2] + y = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device=x_bf16.device) + + local_expert_start = rank_idx * num_experts_per_rank + local_expert_end = local_expert_start + num_experts_per_rank + + for token_idx in range(num_tokens): + for topk_slot in range(num_topk): + expert_idx = topk_idx[token_idx, topk_slot].item() + if expert_idx < 0: + continue + + if expert_idx < local_expert_start or expert_idx >= local_expert_end: + continue + + local_expert = expert_idx - local_expert_start + weight = topk_weights[token_idx, topk_slot].item() + + token = x_bf16[token_idx].float() + + w1 = l1_weights_bf16[local_expert].float() + gate_up = token @ w1.t() + gate = gate_up[:intermediate_hidden] + up = gate_up[intermediate_hidden:] + + h = reference_swiglu(gate, up, activation_clamp) * weight + + w2 = l2_weights_bf16[local_expert].float() + out = h @ w2.t() + + y[token_idx] += out.to(torch.bfloat16) + + return y + + +def reference_mega_moe_batched( + x_bf16: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + l1_weights_bf16: torch.Tensor, + l2_weights_bf16: torch.Tensor, + num_experts_per_rank: int, + rank_idx: int, + num_ranks: int, + activation_clamp: Optional[float] = None, +) -> torch.Tensor: + num_tokens, hidden = x_bf16.shape + num_topk = topk_idx.shape[1] + intermediate_hidden = l2_weights_bf16.shape[2] + y = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device=x_bf16.device) + + local_expert_start = rank_idx * num_experts_per_rank + local_expert_end = local_expert_start + num_experts_per_rank + + for local_expert in range(num_experts_per_rank): + expert_idx = local_expert_start + local_expert + mask = (topk_idx == expert_idx) + token_indices, topk_slots = mask.nonzero(as_tuple=True) + + if token_indices.numel() == 0: + continue + + tokens = x_bf16[token_indices].float() + weights = topk_weights[token_indices, topk_slots].unsqueeze(-1) + + w1 = l1_weights_bf16[local_expert].float() + gate_up = tokens @ w1.t() + gate = gate_up[:, :intermediate_hidden] + up = gate_up[:, intermediate_hidden:] + + h = reference_swiglu(gate, up, activation_clamp) * weights + + w2 = l2_weights_bf16[local_expert].float() + out = h @ w2.t() + + y.index_add_(0, token_indices, out.to(torch.bfloat16)) + + return y + + +def test_correctness(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + torch.manual_seed(rank_idx + 42) + random.seed(rank_idx + 42) + + num_tokens = args.num_tokens if args.num_tokens > 0 else 32 + hidden = args.hidden + intermediate_hidden = args.intermediate_hidden + num_experts = args.num_experts + num_topk = args.num_topk + num_experts_per_rank = num_experts // num_ranks + activation_clamp = args.activation_clamp + + dist_print(f'[SM90 MegaMoE Test] ranks={num_ranks}, tokens={num_tokens}, ' + f'hidden={hidden}, intermediate={intermediate_hidden}, ' + f'experts={num_experts}, topk={num_topk}') + + x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1_weights = torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, device='cuda') * 0.01 + l2_weights = torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.bfloat16, device='cuda') * 0.01 + + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + + if args.masked_ratio > 0: + rand_mask = torch.rand_like(topk_idx, dtype=torch.float) + topk_idx.masked_fill_(rand_mask < args.masked_ratio, -1) + topk_weights.masked_fill_(topk_idx < 0, 0) + + all_x = uneven_all_gather(x_bf16, group) + all_topk_idx = uneven_all_gather(topk_idx, group) + all_topk_weights = uneven_all_gather(topk_weights, group) + + ref_y = reference_mega_moe_batched( + all_x, all_topk_idx, all_topk_weights, + l1_weights, l2_weights, + num_experts_per_rank, rank_idx, num_ranks, + activation_clamp + ) + + ref_y_per_rank = ref_y[:num_tokens] + + dist_print(f'[Reference] y norm: {ref_y_per_rank.float().norm().item():.4f}, ' + f'y abs max: {ref_y_per_rank.float().abs().max().item():.6f}') + + # TODO: Phase 5+ will add SM90 kernel call here and compare with ref_y_per_rank + dist_print('[SM90 MegaMoE] Reference baseline computed successfully. ' + 'Kernel comparison will be added in Phase 5.') + + group.barrier() + dist_print('[PASSED] Phase 0 reference baseline test') + + +def main(): + parser = argparse.ArgumentParser(description='SM90 MegaMoE Test') + parser.add_argument('--num-tokens', type=int, default=32) + parser.add_argument('--num-max-tokens-per-rank', type=int, default=512) + parser.add_argument('--hidden', type=int, default=4096) + parser.add_argument('--intermediate-hidden', type=int, default=2048) + parser.add_argument('--num-experts', type=int, default=16) + parser.add_argument('--num-topk', type=int, default=6) + parser.add_argument('--masked-ratio', type=float, default=0.0) + parser.add_argument('--activation-clamp', type=float, default=None) + parser.add_argument('--local-rank', type=int, default=None) + args = parser.parse_args() + + num_local_ranks = int(os.environ.get('LOCAL_WORLD_SIZE', '1')) + local_rank = args.local_rank if args.local_rank is not None else int(os.environ.get('LOCAL_RANK', '0')) + + test_correctness(local_rank, num_local_ranks, args) + + +if __name__ == '__main__': + main()