chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
062cb160cff9e9316a41ac9365650e0ed16b4b68
DeepGEMM/megamoe-research-reports/megamoe_sm90_research_report.md
Xinyi Liu 062cb160cf Phase 0: SM90 MegaMoE design doc, reference baseline, nsys script 
- MEGAMOE_SM90_DESIGN.md: complete design document with finalized decisions
  (fused single kernel, cooperative + single-WG, dynamic BLOCK_M, etc.)
- tests/test_mega_moe_sm90.py: PyTorch FP32/BF16 reference implementation
  for dispatch → L1 GEMM → SwiGLU → L2 GEMM → combine pipeline
- scripts/run_nsys_mega_moe_sm90.sh: nsys profiling wrapper script
- megamoe-research-reports/: research analysis of PR304/323/347/352/357/360
2026-06-16 18:01:12 +08:00

76 KiB
Raw Blame History

MegaMoE on SM100/SM90 研究现状报告

1. 概述

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

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

2. Git 历史与分支关系

2.1 关键基线 commit

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

2.2 各 worktree 概览

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

2.3 血缘关系图

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

关键发现：

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

3. MegaMoE Kernel 技术对比

3.1 核心架构差异

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

设计理念差异：

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

3.2 BLOCK 配置灵活性

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

pr-352 的 BLOCK_M=16/32 支持 是关键差异化 feature：

  • 使用 mma.sync 指令替代 WGMMA，专门针对 token-per-expert 极小的 decode 场景（如 RL long-tail rollout）。
  • mma.sync 的 FP32 accumulator 通过 SMEM staging 实现 epilogue logical row/column 映射。
  • 此时 swizzle mode 设为 0（禁用 TMA swizzle，使用 32B 粒度加载）。

3.3 Warpgroup 拆分策略

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

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

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

3.4 调度器 (Scheduler) 扩展

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

pr-304 原始 SM100 scheduler：

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

pr-347 引入的变更（被 pr-323 继承）：

  • 修复了 kNumExpertsPerWave 的整除约束变为更宽松的 > 0 && <= kNumExpertsPerRank
  • get_wave_expert_end_idx 增加 clamp 保护部分 wave

pr-352 引入的额外变更：

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

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

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

3.5 SMEM 与 Pipeline 管理

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

3.6 L1 Epilogue 差异

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

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

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

通用符号说明

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

4.1 pr-304: SM100 Fused MegaMoE

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

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

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

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

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

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

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

PR304 warp role map:

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

PR304 同步关键点:

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

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

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

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

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

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

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

数据流关键路径:

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

Warp 间同步关键点:

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

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

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

K1-->K2 跨 kernel 同步:

K1 dispatch 写完 pool block  --> l1_arrival_count release-add
K1 TMA load 等待 l1_arrival_count == num_tokens_per_block
K1 math epilogue 写完 L1 output block -->
    l2_arrival_mask <-- red_or_rel(bit n_block_idx)

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

4.4 pr-360: Pingpong

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

Pingpong 时序展开:

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

4.5 pr-360: Cooperative

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

4.6 四实现横向对比

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

4.7 同步原语对比

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

5. Heuristics 策略对比

5.1 配置选择器架构

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

5.2 pr-304 的 SM100 配置逻辑

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

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

关键规则：

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

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

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

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

5.4 pr-323 的配置逻辑

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

5.5 pr-360 的 token-count 路由

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

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

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

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

6. API 与 Python 层差异

6.1 架构感知路由

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

6.2 Weight Transform

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

6.3 测试架构

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

6.4 PR304/PR357/PR360 关键 API

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

7. 对 SM100 MegaMoE Kernel 的共享改动

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

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

PR304 原版 SM100 kernel 的关键共享事实：

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

8. SM90 MegaMoE 实现难点分析

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

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

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

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

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

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

8.2 Memory Bandwidth 是瓶颈，不是 Compute

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

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

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

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

8.3 Occupancy 被 Register + SMEM 双重限制

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

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

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

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

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

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

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

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

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

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

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

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

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

但这带来两个约束：

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

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

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

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

8.8 代码复杂度爆炸

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

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

8.9 难点优先级总览

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

9. 技术亮点总结

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

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

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

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

pr-360 的主要技术亮点

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

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

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

10. 当前状态评估

10.1 成熟度评估

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

10.2 结论

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

SM90 路线（三条）:

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

SM100 路线（两条）:

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

pr-347 提供了两条路线共享的基础设施重构（调度器 cluster 支持、barrier 优化等）。建议后续工作将两条路线的优势合并：

  1. 以 pr-304 的 API/scale/workspace 语义为 correctness baseline，明确 SM90 FP8-only path 与 SM100 FP8×FP4 path 的差异边界
  2. 吸收 pr-352 的 profile 驱动配置选择器到 pr-323/pr-360 的 fused kernel 中
  3. 吸收 pr-323 的 Split-MN warpgroup 支持到 pr-352 中
  4. 统一到 pr-352 的 unified API 设计，并保留 PR304/PR357 的 SM100-only fast path

11. 实测性能对比

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

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

11.1 Token Sweep 延迟对比

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

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

11.2 Shape Sweep 延迟对比（128 tokens）

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

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

11.3 NCU Profiling 详细对比（32 tokens）

pr-323 (fused kernel: sm90_fp8_mega_moe_impl)

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

pr-352 L1 kernel (sm90_fp8_mega_moe_l1_impl)

  • Duration: 2.23 ms
  • DRAM throughput: 90.73%
  • L2 Cache throughput: 93.41%
  • Compute (SM) throughput: 19.32%

pr-352 L2 kernel (sm90_fp8_mega_moe_l2_impl)

  • Duration: 1.10 ms
  • DRAM throughput: 91.42%
  • L2 Cache throughput: 93.67%
  • Compute (SM) throughput: 19.71%

NCU 关键洞察

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

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

pr-352 的 split 架构优势在于：

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

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

11.4 总体评估

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

11.5 PR360 8×H200 历史补测摘要

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

测试条件摘要：g0033 单节点 8×H200，8 ranks，每 rank 1 GPU；代码路径 /mnt/beegfs/lxy/sm90-bench/pr-360；Python 环境 /mnt/beegfs/lxy/venvs/pr360-torch212-cu130；torch 2.12.0+cu130，tilelang 0.1.9。运行时设置 NVSHMEM_IBGDA_ENABLE=0 NVSHMEM_DISABLE_IBGDA=1 EP_DISABLE_GIN=1 NCCL_NVLS_ENABLE=0。

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

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

日志路径：

  • /mnt/beegfs/lxy/logs/pr360_h200_deepseek_flash_h4096_i2048_e256_topk6_baseline_v1_torch212.log
  • /mnt/beegfs/lxy/logs/pr360_h200_deepseek_pro_h7168_i3072_e384_topk6_baseline_v1_torch212.log
  • /mnt/beegfs/lxy/logs/pr360_h200_mimo_v25_h4096_i2048_e256_topk8_baseline_both_torch212.log
  • /mnt/beegfs/lxy/logs/pr360_h200_mimo_v25_pro_h6144_i2048_e384_topk8_baseline_v1_torch212.log

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

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

12.1 概述

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

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

12.2 三段 Kernel 架构

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

12.3 Green Context 并发模型

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

关键设计点：

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

12.4 与 SM90 路径的关系

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

13. PR360: SM90 Pingpong/Cooperative MegaMoE

13.1 概述

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

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

13.2 两种 Kernel 架构

Pingpong (BLOCK_M=64)

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

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

核心机制：OrderedSequenceBarrier<2,2> 实现两个 math WG 的交替执行：

WG0: [MMA tile 0] [epi tile 0] [MMA tile 2] [epi tile 2] ...
WG1:              [MMA tile 1] [epi tile 1] [MMA tile 3] ...

一个 WG 的 MMA 与另一个 WG 的 epilogue 重叠，隐藏 epilogue 延迟。

Cooperative (BLOCK_M=128)

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

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

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

13.3 Token-Count Auto Routing

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

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

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

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

13.5 代码review主要发现

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

14. 总结与展望

14.1 SM90 MegaMoE 三条路线

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

14.2 SM100 路线

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

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

14.3 后续工作

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