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
This commit is contained in:
Xinyi Liu
1021
MEGAMOE_SM90_DESIGN.md
Normal file
1021
MEGAMOE_SM90_DESIGN.md
Normal file
File diff suppressed because it is too large
Load Diff
BIN
megamoe-research-reports/megamoe.jpg
Normal file
BIN
megamoe-research-reports/megamoe.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 9.8 MiB |
1119
megamoe-research-reports/megamoe_sm90_research_report.md
Normal file
1119
megamoe-research-reports/megamoe_sm90_research_report.md
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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 已验证模式，当前无问题
|
||||||
@@ -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` 的默认行为
|
||||||
53
megamoe-research-reports/pr347_infra_refactor_code_review.md
Normal file
53
megamoe-research-reports/pr347_infra_refactor_code_review.md
Normal file
@@ -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 未退化
|
||||||
@@ -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 测试覆盖
|
||||||
@@ -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 主链路 |
|
||||||
@@ -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`。
|
||||||
52
scripts/run_nsys_mega_moe_sm90.sh
Executable file
52
scripts/run_nsys_mega_moe_sm90.sh
Executable file
@@ -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 ==="
|
||||||
214
tests/test_mega_moe_sm90.py
Normal file
214
tests/test_mega_moe_sm90.py
Normal file
@@ -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()
|
||||||
Reference in New Issue
Block a user