chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add sm90 megamoe phase1 interfaces

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-17 23:54:49 +08:00
parent 062cb160cf
commit 1dafd0f173
12 changed files with 1520 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files

254
AGENTS.md Normal file
View File

@@ -0,0 +1,254 @@
# DeepGEMM 算子开发工作流 (AGENTS.md)
## 项目概述
DeepGEMM 是一个高性能 CUDA kernel 库（C++20/CUDA，JIT 编译），支持 SM90 (H100/H200) 与 SM100 (B200) 架构。
核心开发文件：
- Kernel 实现: `deep_gemm/include/deep_gemm/impls/*.cuh`
- JIT / host 端代码: `csrc/**/*.hpp`, `csrc/**/*.cpp`, `csrc/**/*.cu`
- Python API / 测试: `deep_gemm/**/*.py`, `tests/*.py`
---
## SSH 远程开发工作流（核心）
### 原则：单向同步 Local → Remote，禁止反向
```
本地编辑 → scp 传输 → ssh 远程测试 → 循环
↑ |
└──────────── 禁止 scp remote→local ─────┘
```
### 标准流程
1. **本地编辑** — 在本地修改 kernel/host 代码
2. **scp 传输** — 将变更文件推送到远程服务器
```bash
# 示例：传输单个文件
scp deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh <REMOTE_HOST>:<REMOTE_PATH>/deep_gemm/include/deep_gemm/impls/
# 示例：传输整个目录
scp -r deep_gemm/include/deep_gemm/impls/ <REMOTE_HOST>:<REMOTE_PATH>/deep_gemm/include/deep_gemm/impls/
```
3. **ssh 远程测试** — 在远程服务器上编译并运行测试，长耗时或可能 hang 的命令必须设置合理 timeout
```bash
ssh <REMOTE_HOST> "cd <REMOTE_PATH> && timeout 60s python tests/test_fp8_fp4.py"
```
4. **迭代** — 根据远程测试结果，在本地继续修改，重复步骤 2-3
### 连接配置
- 远程主机: `<REMOTE_HOST>`（从 `~/.ssh/config` 或环境变量读取）
- 远程路径: `<REMOTE_PATH>`（项目在远程服务器的根目录）
- 传输前确认远程路径存在: `ssh <REMOTE_HOST> "ls <REMOTE_PATH>"`
### Timeout 规则
- 运行远程 build / pytest / torchrun / benchmark 命令时，必须根据预期耗时设置 timeout，避免 kernel hang、NCCL hang 或 barrier hang 长时间占用会话。
- 远端 Linux/container 命令优先使用 shell `timeout` 包裹实际测试命令，例如 `timeout 60s python ...` 或 `timeout 60s torchrun ...`。
- 如果调用工具本身也支持 timeout 参数，也必须设置外层 timeout；外层 timeout 应略大于远端 `timeout`，便于收集远端退出信息。
- 经验值：常规 smoke / build / `torchrun` 测试优先设置 60 秒左右 timeout；性能 benchmark 按 case 数量单独估算。
---
## Debug 代码规则（关键）
### 必须加 `DEBUG` 注释
**每行 debug 代码末尾必须加 `// DEBUG` 注释（C++/CUDA）或 `# DEBUG`（Python）。**
```cpp
// 正确示例
printf("m=%d n=%d k=%d\n", m, n, k); // DEBUG
float* debug_buf = new float[M * N]; // DEBUG
// 错误示例（禁止）
printf("m=%d n=%d k=%d\n", m, n, k); // 缺少 DEBUG 标记
```
```python
# 正确示例
print(f"shape: {x.shape}, dtype: {x.dtype}") # DEBUG
torch.save(result, "/tmp/debug_result.pt") # DEBUG
```
### Debug 代码是临时的，禁止提交
Debug 代码仅用于定位问题，**严禁进入 git commit**。
### Debug → 正式修复的标准流程
```
1. 用 debug 代码定位到问题根因
↓
2. git stash -m "DEBUG: <描述本次调试目的>"
（debug 代码被暂存，工作区恢复干净）
↓
3. git stash/git restore
↓
4. 编写正式修复（干净代码，不含 DEBUG 注释）
↓
5. git add <修复的文件>
git commit -m "<修复描述>"
```
**stash 消息必须有意义**，清晰描述调试目标：
```bash
# 好
git stash -m "DEBUG: 排查 sm90_fp8_gemm 的 block scheduling index 越界问题"
# 差
git stash -m "test"
git stash -m "debug"
```
---
## 测试脚本规则
### 存放位置：`./megamoe_dev_test_scripts/phase<phase id>/`
MegaMoE / SM90 开发测试脚本必须按 phase 存放到仓库内，并随对应工作提交到 git：
```bash
mkdir -p megamoe_dev_test_scripts/phase2
```
示例路径：
```text
megamoe_dev_test_scripts/phase0/megamoe_phase0_smoke.py
megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py
```
### 生命周期：随代码提交，禁止 stash-only
```bash
# 1. 编写或更新测试脚本
# megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py
# 2. 与对应实现一起暂存
git add deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
git add megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py
# 3. 随 clean 工作提交进入 git history
git commit -m "feat: 实现 sm90 megamoe dispatch-only 路径"
```
测试脚本是开发过程的一部分，必须可追溯、可复跑。不要再把 MegaMoE 开发测试脚本放到 `.tmp/test_scripts/`，也不要通过 `git stash` 作为唯一保存方式。
---
## 禁止操作清单
| 操作 | 说明 |
|------|------|
| `sed` 编辑代码 | **严格禁止**。使用 Edit 工具修改文件 |
| `scp remote→local` | 禁止从远程拉回变更，所有修改在本地进行 |
| Debug 代码 commit | 禁止将含 `// DEBUG` / `# DEBUG` 的代码提交到 git |
| MegaMoE 测试脚本 stash-only | 测试脚本必须放入 `megamoe_dev_test_scripts/phase<phase id>/` 并随 commit 提交 |
| MegaMoE 测试脚本放 `.tmp/test_scripts/` | 测试脚本不再放 `.tmp`，`.tmp` 只用于非提交的临时文件 |
| 无 timeout 的远程长命令 | build / pytest / torchrun / benchmark 必须设置合理 timeout，避免 hang |
| git stash pop | 避免git conflict |
| ssh xxxx bash -c "大段测试脚本" | 避免直接运行测试，必须持久化成临时脚本 |
---
## Git 管理规范
### 每次 clean 工作提交后的开发日志
每次完成用户请求并创建一次 clean 的工作 commit 后，必须把本次工作内容和详细流程追加到 `MEGAMOE_SM90_DEV.md`，并把该开发日志文档提交到 git。
由于 git commit hash 只有在 commit 创建后才确定，标准流程是两步提交：
```bash
# 1. 提交干净的代码/测试/文档变更
git add <本次工作文件>
git commit -m "<本次工作描述>"
# 2. 获取刚完成的工作 commit hash
git rev-parse HEAD
# 3. 追加开发日志到 MEGAMOE_SM90_DEV.md
# 日志必须记录上一步输出的 HEAD hash
# 4. 单独提交开发日志
git add MEGAMOE_SM90_DEV.md
git commit -m "docs: 记录 <本次工作描述> 开发日志"
```
`MEGAMOE_SM90_DEV.md` 每条日志至少包含：
- 日期和时间
- 对应 clean 工作 commit 的 git HEAD hash
- 用户请求摘要
- 本次提交的核心改动
- 关键文件列表
- 详细开发流程（本地修改、远程同步、编译/测试命令）
- 测试结果和已知问题
- 后续待办
日志 commit 记录的是刚完成的 clean 工作 commit hash，不要求记录日志 commit 自身 hash。
### Debug 周期完整示例
```bash
# === 阶段 1: 添加 debug 代码 ===
# 编辑 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh
# 加入 printf + DEBUG 标记
# === 阶段 2: scp 传输到远程 ===
scp deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh gpu01:/workspace/DeepGEMM/deep_gemm/include/deep_gemm/impls/
# === 阶段 3: ssh 远程测试 ===
ssh gpu01 "cd /workspace/DeepGEMM && python tests/test_fp8_fp4.py"
# === 阶段 4: 定位到问题后，stash debug 代码 ===
git stash -m "DEBUG: 排查 fp8_gemm_1d1d 的 warp scheduling 偏移错误"
# === 阶段 5: 编写正式修复 ===
# 编辑文件，写入干净的修复代码（无 DEBUG 注释）
# === 阶段 6: 提交修复 ===
git add deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh
git commit -m "fix: 修正 sm90_fp8_gemm 1d1d warp scheduling 偏移计算"
```
### 测试脚本管理示例
```bash
# 编写 Phase 2 开发测试
mkdir -p megamoe_dev_test_scripts/phase2
# megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py
# 与对应实现一起提交
git add megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py
git add deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
git commit -m "feat: 实现 sm90 megamoe phase2 dispatch-only"
```
---
## 远程常用命令参考
```bash
# 传输单个 kernel 文件
scp deep_gemm/include/deep_gemm/impls/<kernel>.cuh <HOST>:<PATH>/deep_gemm/include/deep_gemm/impls/
# 传输整个 impls 目录
scp -r deep_gemm/include/deep_gemm/impls/ <HOST>:<PATH>/deep_gemm/include/deep_gemm/impls/
# 传输 host 端代码
scp csrc/**/*.hpp <HOST>:<PATH>/csrc/
# 远程运行测试
ssh <HOST> "cd <PATH> && timeout 60s python tests/test_fp8_fp4.py"
# 远程运行单个测试函数
ssh <HOST> "cd <PATH> && timeout 60s python -c 'from tests.test_fp8_fp4 import test_fp8_gemm_nt; test_fp8_gemm_nt()'"
# 远程 build（如需重新编译 _C.so）
ssh <HOST> "cd <PATH> && timeout 60s bash develop.sh"
# 远程多 rank 测试
ssh <HOST> "cd <PATH> && timeout 60s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py"
```

206
MEGAMOE_SM90_DESIGN.md
View File

@@ -808,24 +808,48 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
--- ---
### Phase 2: Dispatch + 单 tile L1 GEMM (Day 11-20) ### Phase 2: Dispatch + L1 Pool 填充 (Day 11-15)
**目标**：实现 dispatch + L1 WGMMA 的正确 TMA pipeline，**不含 epilogue**（accumulator 直接丢弃）。验证 dispatch pull、arrival count、TMA A/B 加载、WGMMA FP8×FP8 累积 的端到端正确性。 **目标**：只实现 dispatch 数据路径，不进入 TMA/WGMMA。验证 expert count、token routing、source metadata、跨 rank pull、`l1_arrival_count` 等 dispatch 产物正确，为后续 L1 GEMM 提供稳定输入。
| 步骤 | 交付物 | 验证标准 | | 步骤 | 交付物 | 验证标准 |
|------|--------|----------| |------|--------|----------|
| 2.1 Kernel 骨架 | `sm90_fp8_mega_moe.cuh` 中的 warp role 分配：WG0 dispatch+TMA, WG1/WG2 math | 编译通过，launch_bounds(384,1) | | 2.1 Dispatch-only kernel 骨架 | `sm90_fp8_mega_moe.cuh` 中 WG0 的 64-thread dispatch 角色分配；TMA/math WG 可 idle | 编译通过，launch_bounds(384,1)，launch 不崩溃 |
| 2.2 Barrier 初始化 | 所有 mbarrier (full/empty/dispatch/combine) 初始化 + `__syncthreads` | 无 launch failure | | 2.2 Dispatch barrier / workspace 初始化 | dispatch mbarrier、grid sync counter、expert count、arrival counter 初始化 | 连续两次调用无残留 workspace 污染 |
| 2.3 Dispatch phase | 64-thread dispatch: expert count、source metadata、grid sync、NVLink barrier、token pull | 单 rank 模式下 `l1_token_buffer` 中数据正确 | | 2.3 Expert count | 统计本 rank token/top-k 到 local experts 的 recv count，并写入 workspace | expert token counts 与 PyTorch routing reference 一致 |
| 2.4 TMA A/SFA load | w2 TMA A warp：wait `l1_arrival_count`，TMA load L1 acts + float SFA to SMEM | dump SMEM A 区域，与 global 数据逐字节对比 | | 2.4 Source metadata | 写入 `token_src_metadata[pool_token_idx] = {rank, token_idx, topk_idx}` | metadata 与 routing reference 逐项一致 |
| 2.5 TMA B load | w3 TMA B warp：TMA load L1 FP8 weights to SMEM | dump SMEM B 区域对比 | | 2.5 Token/SF/top-k pull | 将输入 token、float SFA、top-k weight 拉入 local L1 pool | 单 rank 下 `l1_acts` / `l1_acts_sf` / `l1_topk_weights` 与输入逐字节或逐值一致 |
| 2.6 WGMMA pipeline | WG1 math warp：`warpgroup_arrive`, 4× K-step `WGMMA::wgmma`, `warpgroup_commit_batch`, `warpgroup_wait<0>` | 单 tile (1 expert, 1 block) 的 accumulator 与 PyTorch FP32 GEMM 结果对比（未 scale） | | 2.6 Cross-rank dispatch smoke | NVLink barrier + remote pull 路径打通 | 2-rank / 8-rank 下无 hang，pool 内容与 reference 一致 |
| 2.7 SF scaling | accum *= sfa * sfb (sfb via `__ldg`) | Scale 后的 accumulator 与 reference scaled GEMM 一致 | | 2.7 `l1_arrival_count` release | 每个 pool block dispatch 完成后 release-add arrival count | arrival count 达到 block 内 token 数，TMA consumer 可安全等待 |
| 2.8 Empty barrier | WGMMA 完成后 arrive empty_barrier → TMA 可重用 SMEM | 多 K-block 的 pipeline 不死锁 |
**验证里程碑**：单 rank 和多 rank 下 dispatch-only 输出 `l1_acts`、`l1_acts_sf`、`l1_topk_weights`、`token_src_metadata`、`l1_arrival_count` 与 PyTorch routing reference 一致；不要求任何 GEMM 计算。
**依赖**：Phase 1 完成
**关键挑战**：
- pool token index / expert block offset 计算必须与 scheduler 后续消费保持一致
- symmetric memory `sym_buffer.map` 与 NVLink barrier 的 phase 计数必须可连续调用
- debug dump 只能临时使用，必须按 AGENTS.md 加 `// DEBUG` / `# DEBUG` 并 stash
---
### Phase 3: L1 TMA + 单 tile WGMMA (Day 16-20)
**目标**：在 Phase 2 dispatch 产物上实现 L1 TMA/WGMMA pipeline，**不含 epilogue**（accumulator 可直接丢弃或写 debug buffer）。验证 TMA A/SFA/B 加载、WGMMA FP8×FP8 累积、SF scaling、empty barrier 的正确性。
| 步骤 | 交付物 | 验证标准 |
|------|--------|----------|
| 3.1 TMA/math kernel 骨架 | WG0 w2/w3 作为 TMA A/B warp，WG1/WG2 作为 math warpgroup | 编译通过，launch_bounds(384,1) |
| 3.2 Full/empty barrier 初始化 | full/empty mbarrier per pipeline stage 初始化 + `__syncthreads` | 无 launch failure / deadlock |
| 3.3 TMA A/SFA load | w2 TMA A warp：wait `l1_arrival_count`，TMA load L1 acts + float SFA to SMEM | dump SMEM A/SFA，与 global pool 数据逐字节/逐值对比 |
| 3.4 TMA B load | w3 TMA B warp：TMA load L1 FP8 weights to SMEM | dump SMEM B 区域对比 |
| 3.5 WGMMA descriptor | 构造 A/B GMMA descriptor；cooperative 下 WG1/WG2 指向不同 M offset | descriptor 地址和 swizzle 与 SM90 GEMM reference 一致 |
| 3.6 WGMMA pipeline | math WG：`warpgroup_arrive`, 4× K-step `WGMMA::wgmma`, `warpgroup_commit_batch`, `warpgroup_wait<0>` | 单 tile accumulator 与 PyTorch FP32 GEMM 结果对比（未 scale） |
| 3.7 SF scaling | accum *= sfa * sfb (sfb via `__ldg`) | Scale 后 accumulator 与 reference scaled GEMM 一致 |
| 3.8 Empty barrier | WGMMA 完成后 arrive empty_barrier → TMA 可重用 SMEM | 多 K-block pipeline 不死锁 |
**验证里程碑**：单 expert、BLOCK_M=128 的 L1 GEMM（无 epilogue）输出与 PyTorch `(dequant(x) @ dequant(w).T) * sfa * sfb` 在 FP32 accumulator 精度内 `allclose(atol=1e-2, rtol=0.05)`。 **验证里程碑**：单 expert、BLOCK_M=128 的 L1 GEMM（无 epilogue）输出与 PyTorch `(dequant(x) @ dequant(w).T) * sfa * sfb` 在 FP32 accumulator 精度内 `allclose(atol=1e-2, rtol=0.05)`。
**依赖**：Phase 1 完成 **依赖**：Phase 2 完成
**关键挑战**： **关键挑战**：
- WGMMA descriptor 构造（参考 `deep_gemm/mma/sm90.cuh` 中的 `make_gmma_desc`） - WGMMA descriptor 构造（参考 `deep_gemm/mma/sm90.cuh` 中的 `make_gmma_desc`）
- Cooperative mode 下两个 WG 的 desc 指向 smem_a 的不同 64-row 偏移 - Cooperative mode 下两个 WG 的 desc 指向 smem_a 的不同 64-row 偏移
@@ -833,27 +857,27 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
--- ---
### Phase 3: L1 Epilogue — SwiGLU + FP8 Quant (Day 21-30) ### Phase 4: L1 Epilogue — SwiGLU + FP8 Quant (Day 21-30)
**目标**：在 Phase 2 的 WGMMA 结果上完成 L1 epilogue 全流程，输出正确的 FP8 `l2_acts` 和 float `l2_acts_sf`。 **目标**：在 Phase 3 的 WGMMA 结果上完成 L1 epilogue 全流程，输出正确的 FP8 `l2_acts` 和 float `l2_acts_sf`。
| 步骤 | 交付物 | 验证标准 | | 步骤 | 交付物 | 验证标准 |
|------|--------|----------| |------|--------|----------|
| 3.1 Top-k weight load | 从 `l1_topk_weights_buffer` load weight per row | 值与 dispatch 写入一致 | | 4.1 Top-k weight load | 从 `l1_topk_weights_buffer` load weight per row | 值与 dispatch 写入一致 |
| 3.2 SwiGLU 计算 | gate/up deinterleave + `silu(gate) * up * topk_weight` in registers | 与 reference SwiGLU 一致（精度：atol=0.02） | | 4.2 SwiGLU 计算 | gate/up deinterleave + `silu(gate) * up * topk_weight` in registers | 与 reference SwiGLU 一致（精度：atol=0.02） |
| 3.3 Activation clamp | `kActivationClamp` 模板参数控制 clamp | clamp 后无超范围值 | | 4.3 Activation clamp | `kActivationClamp` 模板参数控制 clamp | clamp 后无超范围值 |
| 3.4 fast_math SwiGLU | `__expf` + `fast_rcp` 快速路径 | 与精确版相差 < 0.5% | | 4.4 fast_math SwiGLU | `__expf` + `fast_rcp` 快速路径 | 与精确版相差 < 0.5% |
| 3.5 Per-row amax | `warp_reduce<4>(max(abs(...)))` 跨 WG 内 4 warps reduce | amax 值与 PyTorch `abs().max(dim=-1)` 一致 | | 4.5 Per-row amax | `warp_reduce<4>(max(abs(...)))` 跨 WG 内 4 warps reduce | amax 值与 PyTorch `abs().max(dim=-1)` 一致 |
| 3.6 Cooperative amax | 两个 WG 独立 amax（各自 64 行），SMEM scratch 区域用于跨 warp reduce | 每行 amax 正确 | | 4.6 Cooperative amax | 两个 WG 独立 amax（各自 64 行），SMEM scratch 区域用于跨 warp reduce | 每行 amax 正确 |
| 3.7 FP8 quantize | `sf = amax / 448.0`，`fp8 = cast(result / sf)` | 反量化后与 reference 误差 < 1% | | 4.7 FP8 quantize | `sf = amax / 448.0`，`fp8 = cast(result / sf)` | 反量化后与 reference 误差 < 1% |
| 3.8 Float SF write | 写入 `l2_acts_sf` (MN-major, per-64-K layout) | SF 值正确、layout 与 L2 TMA load 兼容 | | 4.8 Float SF write | 写入 `l2_acts_sf` (MN-major, per-64-K layout) | SF 值正确、layout 与 L2 TMA load 兼容 |
| 3.9 STSM + double-buf TMA store | STSM 到 `smem_cd[stage]`，TMA store 到 `l2_acts` | `l2_acts` 内容与 reference FP8 quant output 一致 | | 4.9 STSM + double-buf TMA store | STSM 到 `smem_cd[stage]`，TMA store 到 `l2_acts` | `l2_acts` 内容与 reference FP8 quant output 一致 |
| 3.10 `l2_arrival_mask` 设置 | `red_or_rel_gpu(mask, 1 << n_block_idx)` | 单 expert 全部 L1 N-blocks 完成后 mask == expected | | 4.10 `l2_arrival_mask` 设置 | `red_or_rel_gpu(mask, 1 << n_block_idx)` | 单 expert 全部 L1 N-blocks 完成后 mask == expected |
| 3.11 Cooperative sync | `sync_aligned(256)` 等两个 WG 都 TMA store 完成后设置 mask | 无 race condition | | 4.11 Cooperative sync | `sync_aligned(256)` 等两个 WG 都 TMA store 完成后设置 mask | 无 race condition |
**验证里程碑**：单 expert 的 `l2_acts` (FP8) 和 `l2_acts_sf` (float) 与 reference `SwiGLU → FP8_quant` 一致。Multi-expert 场景下 `l2_arrival_mask` 正确触发。 **验证里程碑**：单 expert 的 `l2_acts` (FP8) 和 `l2_acts_sf` (float) 与 reference `SwiGLU → FP8_quant` 一致。Multi-expert 场景下 `l2_arrival_mask` 正确触发。
**依赖**：Phase 2 完成 **依赖**：Phase 3 完成
**关键挑战**： **关键挑战**：
- WGMMA register layout 到 gate/up 列的映射关系（参考 SM100 的 `SM100_TMEM_LOAD_16dp256b1x` 读取模式，SM90 的 register layout 不同） - WGMMA register layout 到 gate/up 列的映射关系（参考 SM100 的 `SM100_TMEM_LOAD_16dp256b1x` 读取模式，SM90 的 register layout 不同）
- Double-buffered TMA store 的 `tma_store_wait<1>` 与 smem_cd 复用 - Double-buffered TMA store 的 `tma_store_wait<1>` 与 smem_cd 复用
@@ -861,26 +885,26 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
--- ---
### Phase 4: L2 GEMM + NVLink Scatter (Day 31-40) ### Phase 5: L2 GEMM + NVLink Scatter (Day 31-40)
**目标**：L2 WGMMA pipeline + per-64 dual-half SF + BF16 epilogue + NVLink scatter 到 remote combine buffer。 **目标**：L2 WGMMA pipeline + per-64 dual-half SF + BF16 epilogue + NVLink scatter 到 remote combine buffer。
| 步骤 | 交付物 | 验证标准 | | 步骤 | 交付物 | 验证标准 |
|------|--------|----------| |------|--------|----------|
| 4.1 L2 TMA wait | TMA A warp wait `l2_arrival_mask == expected` | 正确等待 L1 所有 N-blocks 完成 | | 5.1 L2 TMA wait | TMA A warp wait `l2_arrival_mask == expected` | 正确等待 L1 所有 N-blocks 完成 |
| 4.2 L2 TMA load | TMA load `l2_acts` + 两个 per-64 SFA halves | SMEM 数据正确 | | 5.2 L2 TMA load | TMA load `l2_acts` + 两个 per-64 SFA halves | SMEM 数据正确 |
| 4.3 L2 WGMMA dual-half | K-step 0-1 累积 half0，K-step 2-3 累积 half1 | 两组 partial accum 正确 | | 5.3 L2 WGMMA dual-half | K-step 0-1 累积 half0，K-step 2-3 累积 half1 | 两组 partial accum 正确 |
| 4.4 Dual-half SF apply | `final = half0 * sfa0 * sfb + half1 * sfa1 * sfb` | 与 reference L2 dequant GEMM 一致 | | 5.4 Dual-half SF apply | `final = half0 * sfa0 * sfb + half1 * sfa1 * sfb` | 与 reference L2 dequant GEMM 一致 |
| 4.5 Weight SF `__ldg` | Math WG prefetch L2 weight SF from global | 值正确 | | 5.5 Weight SF `__ldg` | Math WG prefetch L2 weight SF from global | 值正确 |
| 4.6 BF16 cast + STSM | FP32 → BF16 cast，STSM 到 `smem_cd_l2` (BF16 swizzle) | SMEM 内容正确 | | 5.6 BF16 cast + STSM | FP32 → BF16 cast，STSM 到 `smem_cd_l2` (BF16 swizzle) | SMEM 内容正确 |
| 4.7 Source metadata read | `token_src_metadata[pool_token_idx + row]` → {rank, token, topk} | 值与 dispatch 写入一致 | | 5.7 Source metadata read | `token_src_metadata[pool_token_idx + row]` → {rank, token, topk} | 值与 dispatch 写入一致 |
| 4.8 NVLink scatter | `ld_shared` + `sym_buffer.map` + remote store | 单 rank 下本地 combine buffer 数据正确 | | 5.8 NVLink scatter | `ld_shared` + `sym_buffer.map` + remote store | 单 rank 下本地 combine buffer 数据正确 |
| 4.9 Cooperative smem_cd guard | `sync_aligned(256)` 防止下一 tile 覆盖 smem_cd | 多 tile 序列无 data corruption | | 5.9 Cooperative smem_cd guard | `sync_aligned(256)` 防止下一 tile 覆盖 smem_cd | 多 tile 序列无 data corruption |
| 4.10 Multi-expert persistent loop | Scheduler `for_each_block` 遍历多 expert 的 L1→L2 cycle | 全量 expert 输出正确 | | 5.10 Multi-expert persistent loop | Scheduler `for_each_block` 遍历多 expert 的 L1→L2 cycle | 全量 expert 输出正确 |
**验证里程碑**：多 expert、单 rank 下 `combine_token_buffer` 内容与 reference L2 GEMM BF16 输出一致。`l2_arrival_mask` 正确控制 L1→L2 依赖。 **验证里程碑**：多 expert、单 rank 下 `combine_token_buffer` 内容与 reference L2 GEMM BF16 输出一致。`l2_arrival_mask` 正确控制 L1→L2 依赖。
**依赖**：Phase 3 完成 **依赖**：Phase 4 完成
**关键挑战**： **关键挑战**：
- Per-64 dual-half 需要在 K-loop 中区分 half0/half1（k_step 0-1 vs 2-3），需要两组 accumulator 或 intermediate scale → **224 reg budget 最紧张的地方** - Per-64 dual-half 需要在 K-loop 中区分 half0/half1（k_step 0-1 vs 2-3），需要两组 accumulator 或 intermediate scale → **224 reg budget 最紧张的地方**
- NVLink scatter 的 `sym_buffer.map` 在 SM90 上的行为验证（与 SM100 一致） - NVLink scatter 的 `sym_buffer.map` 在 SM90 上的行为验证（与 SM100 一致）
@@ -888,24 +912,24 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
--- ---
### Phase 5: Combine Reduce + 端到端 Fused (Day 41-48) ### Phase 6: Combine Reduce + 端到端 Fused (Day 41-48)
**目标**：完成 combine reduce、NVLink barrier、workspace 清理。实现完整端到端 fused kernel。 **目标**：完成 combine reduce、NVLink barrier、workspace 清理。实现完整端到端 fused kernel。
| 步骤 | 交付物 | 验证标准 | | 步骤 | 交付物 | 验证标准 |
|------|--------|----------| |------|--------|----------|
| 5.1 NVLink barrier | Epilogue grid sync + cross-rank signal | 多 rank 下所有 rank 到达 barrier | | 6.1 NVLink barrier | Epilogue grid sync + cross-rank signal | 多 rank 下所有 rank 到达 barrier |
| 5.2 Dispatch↔Epilogue handshake | `kDispatchWithEpilogueBarrierIdx` sync | Dispatch 和 epilogue 阶段正确交替 | | 6.2 Dispatch↔Epilogue handshake | `kDispatchWithEpilogueBarrierIdx` sync | Dispatch 和 epilogue 阶段正确交替 |
| 5.3 Combine TMA load | Double-buffered TMA 从 combine buffer load BF16 chunks | Load 数据与 scatter 写入一致 | | 6.3 Combine TMA load | Double-buffered TMA 从 combine buffer load BF16 chunks | Load 数据与 scatter 写入一致 |
| 5.4 FP32 accumulate | 逐 top-k slot 累加到 float registers | 累加结果与 reference `sum(topk_weights * expert_outputs)` 一致 | | 6.4 FP32 accumulate | 逐 top-k slot 累加到 float registers | 累加结果与 reference `sum(topk_weights * expert_outputs)` 一致 |
| 5.5 BF16 cast + TMA store | Cast + TMA store 到 output `y` | `y` 与 reference 在 BF16 精度内一致 | | 6.5 BF16 cast + TMA store | Cast + TMA store 到 output `y` | `y` 与 reference 在 BF16 精度内一致 |
| 5.6 Workspace cleanup | Dispatch warps 清理 expert counts、arrival counters、metadata | 下次 kernel 调用前 workspace 已归零 | | 6.6 Workspace cleanup | Dispatch warps 清理 expert counts、arrival counters、metadata | 下次 kernel 调用前 workspace 已归零 |
| 5.7 NVLink cleanup barrier | 跨 rank 确认 workspace 清理完成 | 连续两次 kernel 调用无 data corruption | | 6.7 NVLink cleanup barrier | 跨 rank 确认 workspace 清理完成 | 连续两次 kernel 调用无 data corruption |
| 5.8 端到端 correctness | 完整 `fp8_mega_moe(...)` 调用 | `test_mega_moe_sm90.py` Smoke test 通过 (2 ranks, 4096×2048, topk=6) | | 6.8 端到端 correctness | 完整 `fp8_mega_moe(...)` 调用 | `test_mega_moe_sm90.py` Smoke test 通过 (2 ranks, 4096×2048, topk=6) |
**验证里程碑**：2-rank Smoke test (4 tokens, 8 experts) `allclose(atol=0.05, rtol=0.1)` 通过。 **验证里程碑**：2-rank Smoke test (4 tokens, 8 experts) `allclose(atol=0.05, rtol=0.1)` 通过。
**依赖**：Phase 4 完成 **依赖**：Phase 5 完成
**关键挑战**： **关键挑战**：
- Combine 逻辑直接移植自 SM100 (与 TMEM/UMMA 无关)，但需要确认 SM90 下 TMA 1D load/store 的行为 - Combine 逻辑直接移植自 SM100 (与 TMEM/UMMA 无关)，但需要确认 SM90 下 TMA 1D load/store 的行为
- NVLink barrier 的 phase 计数在多次 kernel 调用间的正确性 - NVLink barrier 的 phase 计数在多次 kernel 调用间的正确性
@@ -913,65 +937,65 @@ PyTorch FP32/BF16 reference：dispatch → L1 dequant GEMM → SwiGLU → FP8 qu
--- ---
### Phase 6: Single-WG 变体 (BLOCK_M=64/32) (Day 49-55) ### Phase 7: Single-WG 变体 (BLOCK_M=64/32) (Day 49-55)
**目标**：实现 `kCooperativeMode=false` 路径，支持小 BLOCK_M 场景。 **目标**：实现 `kCooperativeMode=false` 路径，支持小 BLOCK_M 场景。
| 步骤 | 交付物 | 验证标准 | | 步骤 | 交付物 | 验证标准 |
|------|--------|----------| |------|--------|----------|
| 6.1 模板分支 | `if constexpr (kCooperativeMode)` 控制 WG2 行为 | 两种模式均编译通过 | | 7.1 模板分支 | `if constexpr (kCooperativeMode)` 控制 WG2 行为 | 两种模式均编译通过 |
| 6.2 WG2 idle | WG2 在 `kCooperativeMode=false` 时 `warpgroup_reg_dealloc` 后 idle | 无 deadlock / launch failure | | 7.2 WG2 idle | WG2 在 `kCooperativeMode=false` 时 `warpgroup_reg_dealloc` 后 idle | 无 deadlock / launch failure |
| 6.3 TMA A 调整 | TMA A 只加载 BLOCK_M 行（64 或 32）而非 128 | SMEM A 数据正确 | | 7.3 TMA A 调整 | TMA A 只加载 BLOCK_M 行（64 或 32）而非 128 | SMEM A 数据正确 |
| 6.4 WGMMA 有效行 | `valid_m = min(tokens_in_block, BLOCK_M)` 控制 epilogue 写入范围 | 无 OOB 写入 | | 7.4 WGMMA 有效行 | `valid_m = min(tokens_in_block, BLOCK_M)` 控制 epilogue 写入范围 | 无 OOB 写入 |
| 6.5 Scheduler 适配 | 动态 BLOCK_M 下 scheduler 的 block 分配和 pool_block_offset 正确 | 全量 expert 覆盖，无漏 tile | | 7.5 Scheduler 适配 | 动态 BLOCK_M 下 scheduler 的 block 分配和 pool_block_offset 正确 | 全量 expert 覆盖，无漏 tile |
| 6.6 Heuristics JIT 分发 | 根据 `expected_tokens_per_expert` 选择 BLOCK_M，JIT 编译对应模板 | 正确选择 config 并编译 | | 7.6 Heuristics JIT 分发 | 根据 `expected_tokens_per_expert` 选择 BLOCK_M，JIT 编译对应模板 | 正确选择 config 并编译 |
| 6.7 Small M correctness | BLOCK_M=64: 单 WG 处理完整 64 行 | Smoke test 通过 | | 7.7 Small M correctness | BLOCK_M=64: 单 WG 处理完整 64 行 | Smoke test 通过 |
| 6.8 Extreme decode | BLOCK_M=32: token-per-expert=1-2 | 极端 decode 场景输出正确 | | 7.8 Extreme decode | BLOCK_M=32: token-per-expert=1-2 | 极端 decode 场景输出正确 |
**验证里程碑**：BLOCK_M=64 和 BLOCK_M=32 的 Smoke test 通过。动态 BLOCK_M 选择逻辑根据 token count 正确切换。 **验证里程碑**：BLOCK_M=64 和 BLOCK_M=32 的 Smoke test 通过。动态 BLOCK_M 选择逻辑根据 token count 正确切换。
**依赖**：Phase 5 完成（cooperative 端到端已验证） **依赖**：Phase 6 完成（cooperative 端到端已验证）
**关键挑战**： **关键挑战**：
- BLOCK_M=32 时 WGMMA 计算 64 行但只用 32 行的 register 浪费 - BLOCK_M=32 时 WGMMA 计算 64 行但只用 32 行的 register 浪费
- Single-WG mode 下 WG2 的 named barrier arrive 需要调整（WG2 不参与 epilogue barrier） - Single-WG mode 下 WG2 的 named barrier arrive 需要调整（WG2 不参与 epilogue barrier）
--- ---
### Phase 7: 多 Rank 集成测试 (Day 56-62) ### Phase 8: 多 Rank 集成测试 (Day 56-62)
**目标**：在 8×H200 集群上验证完整 multi-rank 正确性和性能。 **目标**：在 8×H200 集群上验证完整 multi-rank 正确性和性能。
| 步骤 | 交付物 | 验证标准 | | 步骤 | 交付物 | 验证标准 |
|------|--------|----------| |------|--------|----------|
| 7.1 Multi-rank launch | 8-rank torchrun 启动 MegaMoE kernel | 无 hang / crash | | 8.1 Multi-rank launch | 8-rank torchrun 启动 MegaMoE kernel | 无 hang / crash |
| 7.2 Correctness sweep | Token sweep (1-8192) × Shape sweep × BLOCK_M 三档 | 全部 `allclose` 通过 | | 8.2 Correctness sweep | Token sweep (1-8192) × Shape sweep × BLOCK_M 三档 | 全部 `allclose` 通过 |
| 7.3 Edge cases | 0 tokens, unbalanced routing, single-expert-per-token | 无 crash，输出正确 | | 8.3 Edge cases | 0 tokens, unbalanced routing, single-expert-per-token | 无 crash，输出正确 |
| 7.4 Benchmark framework | `bench_mega_moe_sm90.py`：latency, TFLOPS, HBM GB/s | 可生成 performance table | | 8.4 Benchmark framework | `bench_mega_moe_sm90.py`：latency, TFLOPS, HBM GB/s | 可生成 performance table |
| 7.5 nsys timeline | Phase profile: dispatch / L1 / L2 / combine 各阶段时间 | 识别主要 stall 来源 | | 8.5 nsys timeline | Phase profile: dispatch / L1 / L2 / combine 各阶段时间 | 识别主要 stall 来源 |
| 7.6 DeepEP v1 baseline | 对比 DeepEP dispatch + grouped GEMM + combine | 输出 speedup ratio | | 8.6 DeepEP v1 baseline | 对比 DeepEP dispatch + grouped GEMM + combine | 输出 speedup ratio |
| 7.7 Regression test | CI 集成到 `tests/test_mega_moe_sm90.py` | `pytest` 通过 | | 8.7 Regression test | CI 集成到 `tests/test_mega_moe_sm90.py` | `pytest` 通过 |
**验证里程碑**：8-rank × DSV4 Flash shape (4096×2048, E256, topk6)：correctness 通过 + fused latency < 3000us (128 tokens)。 **验证里程碑**：8-rank × DSV4 Flash shape (4096×2048, E256, topk6)：correctness 通过 + fused latency < 3000us (128 tokens)。
**依赖**：Phase 6 完成 **依赖**：Phase 7 完成
--- ---
### Phase 8: 性能优化 (Day 63-75) ### Phase 9: 性能优化 (Day 63-75)
**目标**：基于 nsys timeline 分析，针对瓶颈进行优化迭代。 **目标**：基于 nsys timeline 分析，针对瓶颈进行优化迭代。
| 步骤 | 优化方向 | 预期收益 | 依据 | | 步骤 | 优化方向 | 预期收益 | 依据 |
|------|---------|---------|------| |------|---------|---------|------|
| 8.1 N-major scheduling tuning | 调整 `tokens/expert >= 256` 阈值 | 大 M 下 weight L2 cache 命中率提升 | NCU L2 sector miss 数据 | | 9.1 N-major scheduling tuning | 调整 `tokens/expert >= 256` 阈值 | 大 M 下 weight L2 cache 命中率提升 | NCU L2 sector miss 数据 |
| 8.2 SFA 软件 prefetch | L2 per-64 SFA 的 `__ldg` prefetch 提前一个 K-block | 减少 math stall waiting on SF | nsys timeline 中 SF load 占比 | | 9.2 SFA 软件 prefetch | L2 per-64 SFA 的 `__ldg` prefetch 提前一个 K-block | 减少 math stall waiting on SF | nsys timeline 中 SF load 占比 |
| 8.3 Combine chunk 大小调优 | 根据 hidden dim 选择 1 vs 2 chunks | 减少 combine 阶段 TMA round-trip | combine 占总时间比例 | | 9.3 Combine chunk 大小调优 | 根据 hidden dim 选择 1 vs 2 chunks | 减少 combine 阶段 TMA round-trip | combine 占总时间比例 |
| 8.4 Dispatch pull 优化 | 增加 dispatch warp 内的 token parallelism | 减少 dispatch phase 时间 | nsys dispatch 耗时 | | 9.4 Dispatch pull 优化 | 增加 dispatch warp 内的 token parallelism | 减少 dispatch phase 时间 | nsys dispatch 耗时 |
| 8.5 Register spill 消除 | 检查 SASS 中的 local memory access，调整变量生命周期 | 消除 spill 带来的 DRAM 流量 | cuobjdump --dump-sass | | 9.5 Register spill 消除 | 检查 SASS 中的 local memory access，调整变量生命周期 | 消除 spill 带来的 DRAM 流量 | cuobjdump --dump-sass |
| 8.6 wave_count 调优 | 根据实测 per-shape 最优 experts_per_wave | 减少 tail effect | benchmark sweep | | 9.6 wave_count 调优 | 根据实测 per-shape 最优 experts_per_wave | 减少 tail effect | benchmark sweep |
| 8.7 Optional: cluster=2 探索 | TMA multicast B tile（需解决 amax 跨 CTA sync） | weight HBM 进一步减半 | 若 cooperative 的 HBM 仍是瓶颈 | | 9.7 Optional: cluster=2 探索 | TMA multicast B tile（需解决 amax 跨 CTA sync） | weight HBM 进一步减半 | 若 cooperative 的 HBM 仍是瓶颈 |
**验证里程碑**：相对 Phase 7 baseline，主要 shape 上 latency 降低 10-20%。 **验证里程碑**：相对 Phase 8 baseline，主要 shape 上 latency 降低 10-20%。
--- ---
@@ -984,24 +1008,27 @@ Phase 0 (环境)
Phase 1 (基础设施) Phase 1 (基础设施)
│ │
v v
Phase 2 (Dispatch + L1 GEMM) Phase 2 (Dispatch)
│ │
v v
Phase 3 (L1 Epilogue) Phase 3 (L1 TMA + WGMMA)
│ │
v v
Phase 4 (L2 GEMM + Scatter) Phase 4 (L1 Epilogue)
│ │
v v
Phase 5 (Combine + 端到端) Phase 5 (L2 GEMM + Scatter)
│
v
Phase 6 (Combine + 端到端)
│ │
├──────────────────┐ ├──────────────────┐
v v v v
Phase 6 (Single-WG) Phase 7 (多 Rank 集成) Phase 7 (Single-WG) Phase 8 (多 Rank 集成)
│ │ │ │
└──────┬───────────┘ └──────┬───────────┘
v v
Phase 8 (性能优化) Phase 9 (性能优化)
``` ```
### 时间线总结 ### 时间线总结
@@ -1010,12 +1037,13 @@ Phase 6 (Single-WG) Phase 7 (多 Rank 集成)
|-------|------|------|---------| |-------|------|------|---------|
| 0 | 3 | 3 | 环境 + baseline | | 0 | 3 | 3 | 环境 + baseline |
| 1 | 7 | 10 | 基础设施 (可 JIT 编译) | | 1 | 7 | 10 | 基础设施 (可 JIT 编译) |
| 2 | 10 | 20 | Dispatch + L1 WGMMA 正确 | | 2 | 5 | 15 | Dispatch + L1 pool 正确 |
| 3 | 10 | 30 | L1 SwiGLU + FP8 quant 正确 | | 3 | 5 | 20 | L1 TMA + WGMMA 正确 |
| 4 | 10 | 40 | L2 GEMM + NVLink scatter 正确 | | 4 | 10 | 30 | L1 SwiGLU + FP8 quant 正确 |
| 5 | 8 | 48 | 端到端 fused kernel 正确 (2-rank) | | 5 | 10 | 40 | L2 GEMM + NVLink scatter 正确 |
| 6 | 7 | 55 | 动态 BLOCK_M 全变体正确 | | 6 | 8 | 48 | 端到端 fused kernel 正确 (2-rank) |
| 7 | 7 | 62 | 8-rank 集成 + benchmark | | 7 | 7 | 55 | 动态 BLOCK_M 全变体正确 |
| 8 | 13 | 75 | 性能调优完成 | | 8 | 7 | 62 | 8-rank 集成 + benchmark |
| 9 | 13 | 75 | 性能调优完成 |
**总工期**：约 75 工作日（10-11 周），含充分的验证和调试时间。关键路径是 Phase 2-5 的串行 kernel 开发（40 天）。 **总工期**：约 75 工作日（10-11 周），含充分的验证和调试时间。关键路径是 Phase 2-6 的串行 kernel 开发。

234
csrc/apis/sm90_mega.hpp Normal file
View File

@@ -0,0 +1,234 @@
#pragma once
#include <functional>
#include <pybind11/functional.h>
#if DG_TENSORMAP_COMPATIBLE
#include "../jit/compiler.hpp"
#endif
#include "../jit/device_runtime.hpp"
#include "../jit_kernels/impls/sm90_fp8_mega_moe.hpp"
namespace deep_gemm::mega {
using SM90MegaMoEBufferViews = std::tuple<
torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>;
static int get_token_alignment_for_sm90_mega_moe() {
return layout::kLCMCandidateBlockM;
}
static std::tuple<int64_t, std::function<SM90MegaMoEBufferViews(const torch::Tensor&)>>
get_symm_buffer_size_for_sm90_mega_moe(
const int& num_ranks, const int& num_experts,
const int& num_max_tokens_per_rank, const int& num_topk,
const int& hidden, const int& intermediate_hidden,
const bool& use_fp8_dispatch, const std::string& activation) {
DG_HOST_ASSERT(num_experts % num_ranks == 0);
// Workspace bytes
const auto workspace = layout::Workspace(nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk);
// Layouts
const auto fp8_token_layout = layout::Data(hidden);
const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden);
const auto bf16_token_layout = layout::Data(hidden * 2);
const auto fp8_sf_layout = layout::Data(hidden / 128 * static_cast<int>(sizeof(float)), false);
const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 128 * static_cast<int>(sizeof(float)), false);
const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false);
const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false);
const auto l1_topk_weights_layout = layout::Data(sizeof(float), false);
// Input buffers
const auto input_token_buffer = layout::Buffer(
fp8_token_layout, 1, num_max_tokens_per_rank,
workspace.get_end_ptr());
const auto input_sf_buffer = layout::Buffer(
fp8_sf_layout, 1, num_max_tokens_per_rank,
input_token_buffer.get_end_ptr());
const auto input_topk_idx_buffer = layout::Buffer(
input_topk_idx_layout, 1, num_max_tokens_per_rank,
input_sf_buffer.get_end_ptr());
const auto input_topk_weights_buffer = layout::Buffer(
input_topk_weights_layout, 1, num_max_tokens_per_rank,
input_topk_idx_buffer.get_end_ptr());
// Buffer configs
const auto num_max_pool_tokens = static_cast<int>(workspace.num_max_pool_tokens);
int num_max_padded_sf_pool_tokens = 0;
for (int block_m: layout::kCandidateBlockM) {
num_max_padded_sf_pool_tokens = std::max(
num_max_padded_sf_pool_tokens,
layout::get_num_padded_sf_pool_tokens(num_max_pool_tokens, block_m)
);
}
// L1 input buffer
const auto l1_token_buffer = layout::Buffer(
fp8_token_layout, 1, num_max_pool_tokens,
input_topk_weights_buffer.get_end_ptr());
const auto l1_sf_buffer = layout::Buffer(
fp8_sf_layout, 1, num_max_padded_sf_pool_tokens,
l1_token_buffer.get_end_ptr());
const auto l1_topk_weights_buffer = layout::Buffer(
l1_topk_weights_layout, 1, num_max_pool_tokens,
l1_sf_buffer.get_end_ptr());
// L2 input buffer
const auto l2_token_buffer = layout::Buffer(
fp8_intermediate_token_layout, 1, num_max_pool_tokens,
l1_topk_weights_buffer.get_end_ptr());
const auto l2_sf_buffer = layout::Buffer(
fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens,
l2_token_buffer.get_end_ptr());
// Combine input buffer
const auto combine_token_buffer = layout::Buffer(
bf16_token_layout, num_topk, num_max_tokens_per_rank,
l2_sf_buffer.get_end_ptr());
// Check SF buffer requirements
DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
DG_HOST_ASSERT(num_max_padded_sf_pool_tokens % 4 == 0);
// Slice function: creates input and L1/L2 pool tensor views.
auto slice_input_buffers = [=](const torch::Tensor& buffer) {
auto x = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_token_buffer.base)),
{num_max_tokens_per_rank, hidden},
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
auto x_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_sf_buffer.base)),
{num_max_tokens_per_rank, hidden / 128},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
auto topk_idx = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_topk_idx_buffer.base)),
{num_max_tokens_per_rank, num_topk},
torch::TensorOptions().dtype(torch::kInt64).device(buffer.device()));
auto topk_weights = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_topk_weights_buffer.base)),
{num_max_tokens_per_rank, num_topk},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
auto l1_acts = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_token_buffer.base)),
{num_max_pool_tokens, hidden},
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
auto l1_acts_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_sf_buffer.base)),
{num_max_padded_sf_pool_tokens, hidden / 128},
{1, num_max_padded_sf_pool_tokens},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
auto l2_acts = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_token_buffer.base)),
{num_max_pool_tokens, intermediate_hidden},
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
auto l2_acts_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_sf_buffer.base)),
{num_max_padded_sf_pool_tokens, intermediate_hidden / 128},
{1, num_max_padded_sf_pool_tokens},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
return std::make_tuple(x, x_sf, topk_idx, topk_weights,
l1_acts, l1_acts_sf, l2_acts, l2_acts_sf);
};
return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
}
static void fp8_mega_moe(
const torch::Tensor& y,
const std::tuple<torch::Tensor, torch::Tensor>& l1_weights_tuple,
const std::tuple<torch::Tensor, torch::Tensor>& l2_weights_tuple,
const std::optional<torch::Tensor>& cumulative_local_expert_recv_stats,
const torch::Tensor& sym_buffer,
const std::vector<int64_t>& sym_buffer_ptrs, const int& rank_idx,
const int& num_max_tokens_per_rank,
const int& num_experts, const int& num_topk,
const std::tuple<int, int, int>& recipe,
const std::string& activation,
const std::optional<float>& activation_clamp_opt,
const bool& fast_math
) {
const auto [l1_weights, l1_weights_sf] = l1_weights_tuple;
const auto [l2_weights, l2_weights_sf] = l2_weights_tuple;
// Config checks
const auto num_tokens = static_cast<int>(y.size(0));
const auto [rm, rn, rk] = recipe;
DG_HOST_ASSERT(rm == 1 and rn == 128 and rk == 128);
DG_HOST_ASSERT(activation == "swiglu");
// Activation checks
const auto activation_clamp =
activation_clamp_opt.value_or(std::numeric_limits<float>::infinity());
DG_HOST_ASSERT(activation_clamp >= 0);
// Tensor checks
DG_HOST_ASSERT(get_major_type_ab(l1_weights) == cute::UMMA::Major::K);
DG_HOST_ASSERT(get_major_type_ab(l2_weights) == cute::UMMA::Major::K);
const auto arch_major = device_runtime->get_arch_major();
const auto [num_experts_per_rank, intermediate_hidden_2, hidden] =
check_grouped_ab_fp8_fp4(l1_weights, cute::UMMA::Major::K, arch_major);
const auto [num_experts_per_rank_, hidden_, intermediate_hidden] =
check_grouped_ab_fp8_fp4(l2_weights, cute::UMMA::Major::K, arch_major);
DG_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank);
DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_);
DG_HOST_ASSERT(hidden == hidden_);
DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden);
DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous());
// Check weight SF layout: float, natural MN-major, per-128-N and per-128-K.
constexpr int kGranMN = 128, kGranK = 128;
check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK,
num_experts_per_rank, false, true, torch::kFloat);
check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK,
num_experts_per_rank, false, true, torch::kFloat);
// Check stats counter
if (cumulative_local_expert_recv_stats.has_value()) {
DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt);
DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank);
DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous());
}
// Check buffer bytes
const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
const auto num_experts_ = num_experts_per_rank * num_ranks;
const auto [num_required_bytes, slice] = get_symm_buffer_size_for_sm90_mega_moe(
num_ranks, num_experts,
num_max_tokens_per_rank, num_topk,
hidden, intermediate_hidden,
true, activation);
DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast<size_t>(num_required_bytes));
DG_HOST_ASSERT(num_experts == num_experts_);
// Already registered tensors
const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer);
// Dispatch into SM90 path
DG_HOST_ASSERT(arch_major == 9);
sm90_fp8_mega_moe(y,
l1_acts, l1_acts_sf,
l2_acts, l2_acts_sf,
l1_weights, l2_weights,
l1_weights_sf, l2_weights_sf,
cumulative_local_expert_recv_stats,
sym_buffer_ptrs,
rank_idx, num_max_tokens_per_rank,
num_experts_per_rank,
num_tokens, num_topk,
hidden, intermediate_hidden,
activation_clamp, fast_math);
if (get_env<int>("DG_COMM_KERNEL_DEBUG"))
sym_buffer.zero_();
}
static void register_sm90_apis(pybind11::module_& m) {
#if DG_TENSORMAP_COMPATIBLE
m.def("get_token_alignment_for_sm90_mega_moe", &get_token_alignment_for_sm90_mega_moe);
m.def("get_symm_buffer_size_for_sm90_mega_moe", &get_symm_buffer_size_for_sm90_mega_moe);
m.def("fp8_mega_moe", &fp8_mega_moe);
#endif
}
} // namespace deep_gemm::mega

212
csrc/jit_kernels/heuristics/sm90_mega_moe.hpp Normal file
View File

@@ -0,0 +1,212 @@
#pragma once
#include <algorithm>
#include <unordered_set>
#include <deep_gemm/layout/mega_moe.cuh>
#include "../../utils/exception.hpp"
#include "../../utils/math.hpp"
#include "../../utils/system.hpp"
namespace deep_gemm {
struct SM90MegaMoEConfig {
// Block tiling
int block_m, block_n, block_k;
int store_block_m;
// Pool capacity and SF-padded token count
int num_max_pool_tokens;
int num_padded_sf_pool_tokens;
// Number of experts to process per wave
int num_experts_per_wave;
// Pipeline stages and shared memory
int num_stages, smem_size;
// Thread layout (384 total: 64 dispatch + 64 TMA + 256 math)
int num_dispatch_threads, num_tma_threads, num_math_threads;
// Mode flags
bool cooperative;
bool use_n_major_l2;
friend std::ostream& operator << (std::ostream& os, const SM90MegaMoEConfig& config) {
os << "SM90MegaMoEConfig("
<< "block_m=" << config.block_m << ", block_n=" << config.block_n << ", block_k=" << config.block_k
<< ", store_block_m=" << config.store_block_m
<< ", num_max_pool_tokens=" << config.num_max_pool_tokens
<< ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens
<< ", num_experts_per_wave=" << config.num_experts_per_wave
<< ", num_stages=" << config.num_stages << ", smem_size=" << config.smem_size
<< ", num_dispatch_threads=" << config.num_dispatch_threads
<< ", num_tma_threads=" << config.num_tma_threads
<< ", num_math_threads=" << config.num_math_threads
<< ", cooperative=" << config.cooperative
<< ", use_n_major_l2=" << config.use_n_major_l2 << ")";
return os;
}
};
static std::tuple<int, int, int, bool> get_block_config_for_sm90_mega_moe(
const int& num_ranks, const int& num_experts,
const int& num_max_tokens_per_rank, const int& num_topk,
const int& num_tokens) {
float num_expected_tokens_per_expert =
static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;
if (num_expected_tokens_per_expert <= 16.5) {
// Extreme decode: RL long-tail, large EP
return {32, 16, 256, false};
} else if (num_expected_tokens_per_expert <= 64.5) {
// Medium decode
return {64, 32, 256, false};
} else {
// Large M prefill / large EP decode
return {128, 32, 256, true};
}
}
static int get_num_experts_per_wave_for_sm90_mega_moe(
const int& num_experts_per_rank, const int& num_tokens, const int& num_topk,
const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms) {
float expected_tokens_per_expert = static_cast<float>(num_tokens) * num_topk / num_experts_per_rank;
if (expected_tokens_per_expert < 1) {
return num_experts_per_rank;
}
constexpr int kImbalanceFactor = 2;
const int num_m_blocks = ceil_div(static_cast<int>(std::ceil(expected_tokens_per_expert)), block_m);
const int num_n_blocks = (2 * intermediate_hidden) / block_n;
const int num_l1_blocks_per_expert = num_m_blocks * num_n_blocks;
int num_experts_per_wave = num_l1_blocks_per_expert > 0
? ceil_div(kImbalanceFactor * num_sms, num_l1_blocks_per_expert) : 1;
num_experts_per_wave = std::min(num_experts_per_wave, num_experts_per_rank);
while (num_experts_per_wave < num_experts_per_rank and num_experts_per_rank % num_experts_per_wave != 0)
++ num_experts_per_wave;
return num_experts_per_wave;
}
static std::pair<int, int> get_pipeline_config_for_sm90_mega_moe(
const int& smem_capacity,
const int& num_experts, const int& hidden,
const int& block_m, const int& block_n, const int& block_k, const int& store_block_m,
const int& num_dispatch_threads, const int& num_math_threads,
const bool& cooperative) {
constexpr int kSmemAlignment = 1024;
constexpr int kNumTMAStoreStages = 2;
const int num_dispatch_warps = num_dispatch_threads / 32;
const int num_math_warps = num_math_threads / 32;
// Dispatch region
const int smem_expert_count_size = align(
num_experts * static_cast<int>(sizeof(uint32_t)), kSmemAlignment);
const int smem_send_buffers_size = align(
static_cast<int>(layout::Buffer(layout::Data(hidden), num_dispatch_warps, 1).get_num_bytes()),
kSmemAlignment);
const int smem_dispatch_size = smem_expert_count_size + smem_send_buffers_size;
// C/D output region: max of L1 FP8 (2 TMA stages) and L2 BF16 staging
// L1: store_block_m * (block_n / 2) * kNumTMAStoreStages (FP8 = 1 byte)
// L2: block_m * block_n * sizeof(BF16) (BF16 = 2 bytes)
const int num_math_warpgroups = cooperative ? 2 : 1;
const int smem_cd_l1 = num_math_warpgroups * store_block_m * (block_n / 2) * kNumTMAStoreStages;
const int smem_cd_l2 = block_m * block_n * static_cast<int>(sizeof(nv_bfloat16));
const int smem_cd = std::max(smem_cd_l1, smem_cd_l2);
// Barriers: dispatch + full/empty pipeline (2 per stage) + combine (2 per math warp)
const int smem_barriers = (num_dispatch_warps + 2 * 2 + num_math_warps * 2) * 8;
// Amax reduction
const int smem_amax_reduction = store_block_m * num_math_warps * static_cast<int>(sizeof(float));
// Float SF per stage: align(2 * BLOCK_M * sizeof(float), 128)
const int smem_sfa_per_stage = align(2 * block_m * static_cast<int>(sizeof(float)), 128);
// Per-stage: A tile + B tile + SFA tile + full/empty barriers
const int smem_per_stage = block_m * block_k + block_n * block_k + smem_sfa_per_stage + 2 * 8;
// Fixed total
const int smem_fixed = smem_dispatch_size + smem_cd + smem_amax_reduction + smem_barriers;
const int num_stages = (smem_capacity - smem_fixed) / smem_per_stage;
DG_HOST_ASSERT(num_stages >= 3);
return {num_stages, smem_fixed + num_stages * smem_per_stage};
}
static SM90MegaMoEConfig get_sm90_mega_moe_config(
const int& num_ranks, const int& num_experts, const int& num_experts_per_rank,
const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk,
const int& hidden, const int& intermediate_hidden,
const int& num_padded_sf_pool_tokens) {
const auto [block_m, store_block_m, num_math_threads, cooperative] =
get_block_config_for_sm90_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens);
const int block_n = 128;
const int block_k = 128;
const int num_max_pool_tokens = layout::get_num_max_pool_tokens(
num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank);
// Thread layout: 64 dispatch + 64 TMA + 256 math/epilogue = 384
const int num_dispatch_threads = 64;
const int num_tma_threads = 64;
// Auto N-major L2: enabled when large M (high tokens per expert)
const bool use_n_major_l2 = [&]() {
auto env_val = get_env<int>("DG_SM90_MOE_NMAJOR", -1);
if (env_val != -1)
return env_val > 0;
float expected = static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;
return expected >= 256;
}();
// Waves
const int num_sms = device_runtime->get_num_sms();
const int num_experts_per_wave = get_num_experts_per_wave_for_sm90_mega_moe(
num_experts_per_rank, num_tokens, num_topk,
intermediate_hidden, block_m, block_n, num_sms);
// Pipeline
constexpr int smem_capacity = 232448;
const auto [num_stages, smem_size] = get_pipeline_config_for_sm90_mega_moe(
smem_capacity,
num_experts, hidden,
block_m, block_n, block_k, store_block_m,
num_dispatch_threads, num_math_threads,
cooperative);
const auto config = SM90MegaMoEConfig {
block_m, block_n, block_k,
store_block_m,
num_max_pool_tokens, num_padded_sf_pool_tokens,
num_experts_per_wave,
num_stages, smem_size,
num_dispatch_threads, num_tma_threads, num_math_threads,
cooperative, use_n_major_l2
};
if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("DG_PRINT_CONFIGS")) {
const auto key = fmt::format(
"SM90MegaMoEConfig(num_ranks={}, num_experts={}, hidden={}, intermediate_hidden={}, num_max_tokens_per_rank={}, num_tokens={}, num_topk={})",
num_ranks, num_experts, hidden, intermediate_hidden, num_max_tokens_per_rank, num_tokens, num_topk);
static std::unordered_set<std::string> printed;
if (printed.count(key) == 0) {
std::cout << key << ": " << config << std::endl;
printed.insert(key);
}
}
return config;
}
} // namespace deep_gemm

207
csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp Normal file
View File

@@ -0,0 +1,207 @@
#pragma once
#include <torch/python.h>
#include "../../jit/compiler.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/format.hpp"
#include "runtime_utils.hpp"
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include "../heuristics/sm90_mega_moe.hpp"
namespace deep_gemm {
class SM90FP8MegaMoERuntime final : public LaunchRuntime<SM90FP8MegaMoERuntime> {
public:
struct Args {
// Templated arguments
int num_max_tokens_per_rank;
int hidden, intermediate_hidden;
int num_experts, num_topk;
int num_ranks;
float activation_clamp;
bool fast_math;
SM90MegaMoEConfig config;
// Runtime arguments
void* y;
int* cumulative_local_expert_recv_stats;
int num_tokens;
layout::SymBuffer<> sym_buffer_ptrs;
// Tensormap
CUtensorMap tensor_map_l1_acts;
CUtensorMap tensor_map_l1_acts_sf;
CUtensorMap tensor_map_l1_weights;
CUtensorMap tensor_map_l1_output;
CUtensorMap tensor_map_l2_acts;
CUtensorMap tensor_map_l2_acts_sf;
CUtensorMap tensor_map_l2_weights;
void* l1_weights_sf;
void* l2_weights_sf;
// Launch configs
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_gemm/impls/sm90_fp8_mega_moe.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm90_fp8_mega_moe_impl<
{},
{}, {},
{}, {},
{},
{}, {}, {},
{},
{}, {},
{},
{}, {}, {},
{}, {},
{}, {},
{}, {}
>);
}};
)", args.num_max_tokens_per_rank,
args.hidden, args.intermediate_hidden,
args.num_experts, args.num_topk,
args.config.num_experts_per_wave,
args.config.block_m, args.config.block_n, args.config.block_k,
args.config.store_block_m,
args.config.num_max_pool_tokens,
args.config.num_padded_sf_pool_tokens,
args.config.num_stages,
args.config.num_dispatch_threads, args.config.num_tma_threads, args.config.num_math_threads,
args.config.cooperative ? "true" : "false",
args.config.use_n_major_l2 ? "true" : "false",
args.launch_args.grid_dim.first, args.num_ranks,
to_string(args.activation_clamp),
args.fast_math ? "true" : "false");
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.y,
args.cumulative_local_expert_recv_stats,
args.num_tokens,
args.sym_buffer_ptrs,
args.tensor_map_l1_acts,
args.tensor_map_l1_acts_sf,
args.tensor_map_l1_weights,
args.l1_weights_sf,
args.tensor_map_l1_output,
args.tensor_map_l2_acts,
args.tensor_map_l2_acts_sf,
args.tensor_map_l2_weights,
args.l2_weights_sf
));
}
};
static void sm90_fp8_mega_moe(
const torch::Tensor& y,
const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf,
const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf,
const torch::Tensor& l1_weights, const torch::Tensor& l2_weights,
const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf,
const std::optional<torch::Tensor> cumulative_local_expert_recv_stats,
const std::vector<int64_t>& sym_buffer_ptrs,
const int& rank_idx, const int& num_max_tokens_per_rank,
const int& num_experts_per_rank,
const int& num_tokens, const int& num_topk,
const int& hidden, const int& intermediate_hidden,
const float& activation_clamp,
const bool& fast_math
) {
const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
const auto num_experts = num_experts_per_rank * num_ranks;
const auto num_padded_sf_pool_tokens = static_cast<int>(l1_acts_sf.size(0));
// Heuristics
const auto config = get_sm90_mega_moe_config(
num_ranks, num_experts, num_experts_per_rank,
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_padded_sf_pool_tokens);
// Make tensormap
constexpr int kGranK = 128;
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
hidden, config.num_max_pool_tokens,
config.block_k, config.block_m,
static_cast<int>(l1_acts.stride(-2)),
128);
const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf,
config.num_padded_sf_pool_tokens, hidden,
config.block_m, kGranK,
1, 0);
const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights,
hidden, num_experts_per_rank * intermediate_hidden * 2,
config.block_k, config.block_n,
static_cast<int>(l1_weights.stride(-2)),
128);
// L1 output SwiGLU has half N width
const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts,
intermediate_hidden, config.num_max_pool_tokens,
config.block_n / 2, config.store_block_m,
static_cast<int>(l2_acts.stride(-2)),
64);
const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts,
intermediate_hidden, config.num_max_pool_tokens,
config.block_k, config.block_m,
static_cast<int>(l2_acts.stride(-2)),
128);
const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf,
config.num_padded_sf_pool_tokens, intermediate_hidden,
config.block_m, kGranK,
1, 0);
const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights,
intermediate_hidden, num_experts_per_rank * hidden,
config.block_k, config.block_n,
static_cast<int>(l2_weights.stride(-2)),
128);
// Stats can be optional
int* cumulative_local_expert_recv_stats_ptr = nullptr;
if (cumulative_local_expert_recv_stats.has_value())
cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr<int>();
// Launch
const auto num_sms = device_runtime->get_num_sms();
const int num_threads = config.num_dispatch_threads + config.num_tma_threads + config.num_math_threads;
const SM90FP8MegaMoERuntime::Args args = {
.num_max_tokens_per_rank = num_max_tokens_per_rank,
.hidden = hidden, .intermediate_hidden = intermediate_hidden,
.num_experts = num_experts, .num_topk = num_topk,
.num_ranks = num_ranks,
.activation_clamp = activation_clamp,
.fast_math = fast_math,
.config = config,
.y = y.data_ptr(),
.cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr,
.num_tokens = num_tokens,
.sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx),
.tensor_map_l1_acts = tensor_map_l1_acts,
.tensor_map_l1_acts_sf = tensor_map_l1_acts_sf,
.tensor_map_l1_weights = tensor_map_l1_weights,
.tensor_map_l1_output = tensor_map_l1_output,
.tensor_map_l2_acts = tensor_map_l2_acts,
.tensor_map_l2_acts_sf = tensor_map_l2_acts_sf,
.tensor_map_l2_weights = tensor_map_l2_weights,
.l1_weights_sf = l1_weights_sf.data_ptr(),
.l2_weights_sf = l2_weights_sf.data_ptr(),
.launch_args = LaunchArgs(num_sms, num_threads, config.smem_size)
};
const auto code = SM90FP8MegaMoERuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_mega_moe", code);
SM90FP8MegaMoERuntime::launch(runtime, args);
}
} // namespace deep_gemm

2
csrc/python_api.cpp
View File

@@ -7,6 +7,7 @@
#include "apis/gemm.hpp" #include "apis/gemm.hpp"
#include "apis/layout.hpp" #include "apis/layout.hpp"
#include "apis/mega.hpp" #include "apis/mega.hpp"
#include "apis/sm90_mega.hpp"
#include "apis/runtime.hpp" #include "apis/runtime.hpp"
#ifndef TORCH_EXTENSION_NAME #ifndef TORCH_EXTENSION_NAME
@@ -24,5 +25,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
deep_gemm::gemm::register_apis(m); deep_gemm::gemm::register_apis(m);
deep_gemm::layout::register_apis(m); deep_gemm::layout::register_apis(m);
deep_gemm::mega::register_apis(m); deep_gemm::mega::register_apis(m);
deep_gemm::mega::register_sm90_apis(m);
deep_gemm::runtime::register_apis(m); deep_gemm::runtime::register_apis(m);
} }

1
deep_gemm/__init__.py
View File

@@ -85,6 +85,7 @@ from .mega import (
SymmBuffer, SymmBuffer,
get_symm_buffer_for_mega_moe, get_symm_buffer_for_mega_moe,
transform_weights_for_mega_moe, transform_weights_for_mega_moe,
fp8_mega_moe,
fp8_fp4_mega_moe, fp8_fp4_mega_moe,
) )

58
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh Normal file
View File

@@ -0,0 +1,58 @@
#pragma once
#include <cutlass/cutlass.h>
#include <deep_gemm/common/exception.cuh>
#include <deep_gemm/common/tma_utils.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include <deep_gemm/scheduler/sm90_mega_moe.cuh>
namespace deep_gemm {
template <
uint32_t kNumMaxTokensPerRank,
uint32_t kHidden, uint32_t kIntermediateHidden,
uint32_t kNumExperts, uint32_t kNumTopk,
uint32_t kNumExpertsPerWave,
uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
uint32_t kStoreBlockM,
uint32_t kNumMaxPoolTokens,
uint32_t kNumPaddedSFPoolTokens,
uint32_t kNumStages,
uint32_t kNumDispatchThreads, uint32_t kNumTMAThreads,
uint32_t kNumMathThreads,
bool kCooperativeMode, bool kUseNMajorL2,
uint32_t kNumSMs, uint32_t kNumRanks,
float kActivationClamp, bool kFastMath,
uint32_t L1_SHAPE_N = kIntermediateHidden * 2,
uint32_t L1_SHAPE_K = kHidden,
uint32_t L2_SHAPE_N = kHidden,
uint32_t L2_SHAPE_K = kIntermediateHidden,
uint32_t kNumThreads = kNumDispatchThreads + kNumTMAThreads + kNumMathThreads,
uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks>
CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void
sm90_fp8_mega_moe_impl(void* y,
int* cumulative_local_expert_recv_stats,
const uint32_t num_tokens,
const __grid_constant__ layout::SymBuffer<kNumRanks> sym_buffer,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights,
const void* l1_weights_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights,
const void* l2_weights_sf) {
DG_STATIC_ASSERT(kNumThreads == 384, "SM90 MegaMoE expects 384 threads");
DG_STATIC_ASSERT(BLOCK_N == 128, "SM90 MegaMoE expects BLOCK_N=128");
DG_STATIC_ASSERT(BLOCK_K == 128, "SM90 MegaMoE expects BLOCK_K=128");
DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks");
// Phase 1 only validates the host/JIT/API path and launches an empty kernel.
return;
}
} // namespace deep_gemm

20
deep_gemm/include/deep_gemm/layout/mega_moe.cuh
View File

@@ -108,46 +108,46 @@ struct Workspace {
static constexpr uint32_t kNumMaxGridSyncCounters = 4; static constexpr uint32_t kNumMaxGridSyncCounters = 4;
template <uint32_t kIndex = 0> template <uint32_t kIndex = 0>
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
uint32_t* get_grid_sync_count_ptr() const { uint32_t* get_grid_sync_count_ptr() const {
DG_STATIC_ASSERT(kIndex < kNumMaxGridSyncCounters, "Grid sync index out of bounds"); DG_STATIC_ASSERT(kIndex < kNumMaxGridSyncCounters, "Grid sync index out of bounds");
return static_cast<uint32_t*>(base) + kIndex; return static_cast<uint32_t*>(base) + kIndex;
} }
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
uint32_t* get_nvl_barrier_counter_ptr() const { uint32_t* get_nvl_barrier_counter_ptr() const {
return static_cast<uint32_t*>(base) + kNumMaxGridSyncCounters; return static_cast<uint32_t*>(base) + kNumMaxGridSyncCounters;
} }
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
int* get_nvl_barrier_signal_ptr(const uint32_t& phase) const { int* get_nvl_barrier_signal_ptr(const uint32_t& phase) const {
// NOTES: the signal is signed, as we may minus // NOTES: the signal is signed, as we may minus
return math::advance_ptr<int>(base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int)); return math::advance_ptr<int>(base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int));
} }
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const { uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const {
return math::advance_ptr<uint64_t>(base, kNumBarrierSignalBytes) + expert_idx; return math::advance_ptr<uint64_t>(base, kNumBarrierSignalBytes) + expert_idx;
} }
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
uint64_t* get_expert_recv_count_ptr( uint64_t* get_expert_recv_count_ptr(
const uint32_t& rank_idx = 0, const uint32_t& expert_idx = 0) const { const uint32_t& rank_idx = 0, const uint32_t& expert_idx = 0) const {
return get_expert_send_count_ptr(num_experts) + rank_idx * num_experts_per_rank + expert_idx; return get_expert_send_count_ptr(num_experts) + rank_idx * num_experts_per_rank + expert_idx;
} }
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
uint64_t* get_expert_recv_count_sum_ptr(const uint32_t& expert_idx = 0) const { uint64_t* get_expert_recv_count_sum_ptr(const uint32_t& expert_idx = 0) const {
return get_expert_send_count_ptr(num_experts * 2) + expert_idx; return get_expert_send_count_ptr(num_experts * 2) + expert_idx;
} }
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
uint32_t* get_l1_arrival_count_ptr(const uint32_t& pool_block_idx = 0) const { uint32_t* get_l1_arrival_count_ptr(const uint32_t& pool_block_idx = 0) const {
const auto base = get_expert_recv_count_sum_ptr(num_experts_per_rank); const auto base = get_expert_recv_count_sum_ptr(num_experts_per_rank);
return reinterpret_cast<uint32_t*>(base) + pool_block_idx; return reinterpret_cast<uint32_t*>(base) + pool_block_idx;
} }
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
uint64_t* get_l2_arrival_mask_ptr(const uint32_t& pool_block_idx = 0) const { uint64_t* get_l2_arrival_mask_ptr(const uint32_t& pool_block_idx = 0) const {
// Pad L1 entry count to even so that the `l2_arrival_mask` is 8-byte aligned // Pad L1 entry count to even so that the `l2_arrival_mask` is 8-byte aligned
const auto base = get_l1_arrival_count_ptr(math::align(num_max_pool_blocks, 2u)); const auto base = get_l1_arrival_count_ptr(math::align(num_max_pool_blocks, 2u));
@@ -155,7 +155,7 @@ struct Workspace {
} }
// For dispatch pulling // For dispatch pulling
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
uint32_t* get_src_token_topk_idx_ptr( uint32_t* get_src_token_topk_idx_ptr(
const uint32_t& expert_idx = 0, const uint32_t& rank_idx = 0, const uint32_t& token_idx = 0) const { const uint32_t& expert_idx = 0, const uint32_t& rank_idx = 0, const uint32_t& token_idx = 0) const {
const auto base = get_l2_arrival_mask_ptr(num_max_pool_blocks); const auto base = get_l2_arrival_mask_ptr(num_max_pool_blocks);
@@ -165,7 +165,7 @@ struct Workspace {
} }
// For combine usages // For combine usages
CUTLASS_DEVICE CUTLASS_HOST_DEVICE
TokenSrcMetadata* get_token_src_metadata_ptr(const uint32_t& pool_token_idx = 0) const { TokenSrcMetadata* get_token_src_metadata_ptr(const uint32_t& pool_token_idx = 0) const {
const auto base = reinterpret_cast<TokenSrcMetadata*>(get_src_token_topk_idx_ptr(num_experts_per_rank)); const auto base = reinterpret_cast<TokenSrcMetadata*>(get_src_token_topk_idx_ptr(num_experts_per_rank));
return base + pool_token_idx; return base + pool_token_idx;

199
deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh Normal file
View File

@@ -0,0 +1,199 @@
#pragma once
#include <deep_gemm/common/cute_tie.cuh>
#include <deep_gemm/common/math.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/ptx/ld_st.cuh>
#include <deep_gemm/ptx/utils.cuh>
namespace deep_gemm::sched {
enum class SM90BlockPhase {
None = 0,
Linear1 = 1,
Linear2 = 2
};
template <uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
uint32_t L1_SHAPE_N, uint32_t L1_SHAPE_K,
uint32_t L2_SHAPE_N, uint32_t L2_SHAPE_K,
uint32_t kNumExpertsPerRank,
uint32_t kNumExpertsPerWave,
uint32_t kNumSMs, uint32_t kNumRanks,
bool kUseNMajorL2,
uint32_t kNumExpertsPerLane = math::constexpr_ceil_div(kNumExpertsPerRank, 32u),
uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N,
uint32_t kNumL2BlockNs = L2_SHAPE_N / BLOCK_N,
uint32_t kNumL1BlockKs = L1_SHAPE_K / BLOCK_K,
uint32_t kNumL2BlockKs = L2_SHAPE_K / BLOCK_K>
struct SM90MegaMoEScheduler {
DG_STATIC_ASSERT(L1_SHAPE_N % BLOCK_N == 0, "Invalid shape");
DG_STATIC_ASSERT(L2_SHAPE_N % BLOCK_N == 0, "Invalid shape");
DG_STATIC_ASSERT(L1_SHAPE_K % BLOCK_K == 0, "Invalid shape");
DG_STATIC_ASSERT(L2_SHAPE_K % BLOCK_K == 0, "Invalid shape");
DG_STATIC_ASSERT(kNumExpertsPerRank % kNumExpertsPerWave == 0, "Invalid wave config");
const layout::Workspace& workspace;
SM90BlockPhase next_phase = SM90BlockPhase::Linear1;
uint32_t current_local_expert_idx = 0;
uint32_t current_num_tokens = 0;
uint32_t current_pool_block_offset = 0;
uint32_t block_idx = 0;
uint32_t m_block_idx = 0;
uint32_t n_block_idx = 0;
uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {};
CUTLASS_DEVICE explicit SM90MegaMoEScheduler(const layout::Workspace& workspace): workspace(workspace) {
block_idx = blockIdx.x;
}
CUTLASS_DEVICE uint32_t get_wave_expert_end_idx() const {
return math::align(current_local_expert_idx + 1, kNumExpertsPerWave);
}
CUTLASS_DEVICE uint32_t get_num_tokens(const uint32_t& expert_idx) const {
uint32_t valid_value;
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ?
stored_num_tokens_per_expert[i] : valid_value;
}
return ptx::exchange(valid_value, expert_idx % 32);
}
CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) {
uint32_t num_blocks = 0;
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
if (i * 32 + ptx::get_lane_idx() < expert_idx)
num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], BLOCK_M);
}
return __reduce_add_sync(0xffffffff, num_blocks);
}
CUTLASS_DEVICE void advance_expert_idx() {
current_pool_block_offset += get_current_num_m_blocks();
current_local_expert_idx += 1;
current_num_tokens = get_num_tokens(current_local_expert_idx);
}
CUTLASS_DEVICE void set_expert_idx(const uint32_t& expert_idx) {
current_local_expert_idx = expert_idx;
current_num_tokens = get_num_tokens(expert_idx);
current_pool_block_offset = get_pool_block_offset(expert_idx);
}
CUTLASS_DEVICE uint32_t get_current_pool_block_offset() const {
return current_pool_block_offset;
}
CUTLASS_DEVICE uint32_t get_current_num_m_blocks() const {
return math::ceil_div(current_num_tokens, BLOCK_M);
}
CUTLASS_DEVICE uint32_t get_valid_m() const {
return cute::min(current_num_tokens - m_block_idx * BLOCK_M, BLOCK_M);
}
CUTLASS_DEVICE bool fetch_next_l1_block() {
const auto wave_end_expert_idx = get_wave_expert_end_idx();
while (current_local_expert_idx < wave_end_expert_idx) {
const auto num_m_blocks = get_current_num_m_blocks();
m_block_idx = block_idx / kNumL1BlockNs;
if (m_block_idx < num_m_blocks)
return true;
block_idx -= num_m_blocks * kNumL1BlockNs;
advance_expert_idx();
}
return false;
}
CUTLASS_DEVICE bool fetch_next_l2_block() {
const auto wave_end_expert_idx = get_wave_expert_end_idx();
while (current_local_expert_idx < wave_end_expert_idx) {
const auto num_m_blocks = get_current_num_m_blocks();
if (block_idx < num_m_blocks * kNumL2BlockNs) {
if constexpr (kUseNMajorL2) {
n_block_idx = block_idx / num_m_blocks;
m_block_idx = block_idx - n_block_idx * num_m_blocks;
} else {
m_block_idx = block_idx / kNumL2BlockNs;
n_block_idx = block_idx - m_block_idx * kNumL2BlockNs;
}
return true;
}
block_idx -= num_m_blocks * kNumL2BlockNs;
advance_expert_idx();
}
return false;
}
CUTLASS_DEVICE cute::tuple<SM90BlockPhase, uint32_t, uint32_t, uint32_t> get_next_block() {
while (true) {
if (current_local_expert_idx >= kNumExpertsPerRank)
break;
if (next_phase == SM90BlockPhase::Linear1) {
if (fetch_next_l1_block()) {
n_block_idx = block_idx - m_block_idx * kNumL1BlockNs;
block_idx += kNumSMs;
return {SM90BlockPhase::Linear1, current_local_expert_idx, m_block_idx, n_block_idx};
} else {
next_phase = SM90BlockPhase::Linear2;
set_expert_idx(math::align<uint32_t, false>(current_local_expert_idx - 1, kNumExpertsPerWave));
}
} else {
if (fetch_next_l2_block()) {
if constexpr (not kUseNMajorL2) {
n_block_idx = block_idx - m_block_idx * kNumL2BlockNs;
}
block_idx += kNumSMs;
return {SM90BlockPhase::Linear2, current_local_expert_idx, m_block_idx, n_block_idx};
} else {
next_phase = SM90BlockPhase::Linear1;
}
}
}
return {SM90BlockPhase::None, 0, 0, 0};
}
CUTLASS_DEVICE void fetch_expert_recv_count() {
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
const auto expert_idx = i * 32 + ptx::get_lane_idx();
uint64_t value = 0;
if (expert_idx < kNumExpertsPerRank) {
do {
value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx));
} while (static_cast<uint32_t>(value >> 32) != kNumSMs * kNumRanks);
}
stored_num_tokens_per_expert[i] = static_cast<uint32_t>(value);
}
__syncwarp();
}
template <typename Func>
CUTLASS_DEVICE void for_each_block(Func&& func) {
fetch_expert_recv_count();
set_expert_idx(0);
while (true) {
CUTE_TIE_DECL(get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx);
if (block_phase == SM90BlockPhase::None)
break;
func(block_phase, current_local_expert_idx,
block_phase == SM90BlockPhase::Linear2 ? kNumL2BlockKs : kNumL1BlockKs,
m_block_idx, n_block_idx);
}
}
};
} // namespace deep_gemm::sched

147
deep_gemm/mega/__init__.py
View File

@@ -13,6 +13,14 @@ except Exception as exception:
from .. import _C from .. import _C
def _is_sm90() -> bool:
return torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
def _is_sm100() -> bool:
return torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 10
class SymmBuffer: class SymmBuffer:
def __init__(self, group: dist.ProcessGroup, def __init__(self, group: dist.ProcessGroup,
# MoE arguments # MoE arguments
@@ -28,13 +36,23 @@ class SymmBuffer:
self.hidden = hidden self.hidden = hidden
self.intermediate_hidden = intermediate_hidden self.intermediate_hidden = intermediate_hidden
# Allocate a symmetric buffer # Allocate a symmetric buffer (route by architecture)
num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe( if _is_sm90():
group.size(), num_experts, num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe(
num_max_tokens_per_rank, num_topk, group.size(), num_experts,
hidden, intermediate_hidden, num_max_tokens_per_rank, num_topk,
use_fp8_dispatch, activation hidden, intermediate_hidden,
) use_fp8_dispatch, activation
)
elif _is_sm100():
num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe(
group.size(), num_experts,
num_max_tokens_per_rank, num_topk,
hidden, intermediate_hidden,
use_fp8_dispatch, activation
)
else:
raise RuntimeError('Unsupported architecture for MegaMoE')
self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda') self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')
self.handle = symm_mem.rendezvous(self.buffer, group=group) self.handle = symm_mem.rendezvous(self.buffer, group=group)
self.buffer.zero_() self.buffer.zero_()
@@ -46,6 +64,10 @@ class SymmBuffer:
self.topk_idx, self.topk_weights, self.topk_idx, self.topk_weights,
self.l1_acts, self.l1_acts_sf, self.l1_acts, self.l1_acts_sf,
self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer) self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer)
self.l1_topk_weights = None
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
def destroy(self): def destroy(self):
self.handle = None self.handle = None
@@ -53,6 +75,16 @@ class SymmBuffer:
self.group = None self.group = None
self.x = None self.x = None
self.x_sf = None self.x_sf = None
self.topk_idx = None
self.topk_weights = None
self.l1_acts = None
self.l1_acts_sf = None
self.l1_topk_weights = None
self.l2_acts = None
self.l2_acts_sf = None
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,
@@ -62,7 +94,13 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,
use_fp8_dispatch: bool = True, use_fp8_dispatch: bool = True,
activation: str = 'swiglu') -> SymmBuffer: activation: str = 'swiglu') -> SymmBuffer:
# Token count must be aligned to block sizes # Token count must be aligned to block sizes
num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe()) if _is_sm90():
alignment = _C.get_token_alignment_for_sm90_mega_moe()
elif _is_sm100():
alignment = _C.get_token_alignment_for_mega_moe()
else:
raise RuntimeError('Unsupported architecture for MegaMoE')
num_max_tokens_per_rank = align(num_max_tokens_per_rank, alignment)
return SymmBuffer( return SymmBuffer(
group, num_experts, group, num_experts,
@@ -72,16 +110,17 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,
) )
def _interleave_l1_weights(l1_weights: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]: def _interleave_l1_weight_tensor(t: torch.Tensor, gran: int = 8) -> torch.Tensor:
# [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up] # [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up]
def interleave(t, gran: int = 8) -> torch.Tensor: g, n, *rest = t.shape
g, n, *rest = t.shape half = n // 2
half = n // 2 gate = t[:, :half].reshape(g, half // gran, gran, *rest)
gate = t[:, :half].reshape(g, half // gran, gran, *rest) up = t[:, half:].reshape(g, half // gran, gran, *rest)
up = t[:, half:].reshape(g, half // gran, gran, *rest) return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest))
return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest))
return interleave(l1_weights[0]), interleave(l1_weights[1])
def _interleave_l1_weights(l1_weights: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
return _interleave_l1_weight_tensor(l1_weights[0]), _interleave_l1_weight_tensor(l1_weights[1])
def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor: def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor:
@@ -93,36 +132,82 @@ def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor:
return torch.empty_like(sf).copy_(result) return torch.empty_like(sf).copy_(result)
def transform_weights_for_mega_moe_sm90(
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor]
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
# L1: interleave FP8 gate/up weights only; SM90 float weight SF stays natural MN-major.
l1_weights = (_interleave_l1_weight_tensor(l1_weights[0]), l1_weights[1])
# L2: no transform
return l1_weights, l2_weights
def transform_weights_for_mega_moe( def transform_weights_for_mega_moe(
l1_weights: Tuple[torch.Tensor, torch.Tensor], l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor] l2_weights: Tuple[torch.Tensor, torch.Tensor]
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
# L1: interleave gate/up, then transpose SF for UTCCP if _is_sm90():
return transform_weights_for_mega_moe_sm90(l1_weights, l2_weights)
# SM100: L1 interleave gate/up + UTCCP SF transpose, L2 UTCCP SF transpose
l1_interleaved = _interleave_l1_weights(l1_weights) l1_interleaved = _interleave_l1_weights(l1_weights)
l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1])) l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1]))
# L2: only transpose SF for UTCCP
l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1])) l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1]))
return l1_weights, l2_weights return l1_weights, l2_weights
def fp8_mega_moe(y: torch.Tensor,
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor],
sym_buffer: SymmBuffer,
cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
recipe: Optional[Tuple[int, int, int]] = None,
activation: str = 'swiglu',
activation_clamp: Optional[float] = None,
fast_math: bool = True):
if _is_sm90():
if recipe is None:
recipe = (1, 128, 128)
_C.fp8_mega_moe(
y,
l1_weights, l2_weights,
cumulative_local_expert_recv_stats,
sym_buffer.buffer,
sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
sym_buffer.num_max_tokens_per_rank,
sym_buffer.num_experts, sym_buffer.num_topk,
recipe,
activation, activation_clamp,
fast_math
)
elif _is_sm100():
if recipe is None:
recipe = (1, 1, 32)
_C.fp8_fp4_mega_moe(
y,
l1_weights, l2_weights,
cumulative_local_expert_recv_stats,
sym_buffer.buffer,
sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
sym_buffer.num_max_tokens_per_rank,
sym_buffer.num_experts, sym_buffer.num_topk,
recipe,
activation, activation_clamp,
fast_math
)
else:
raise RuntimeError('Unsupported architecture for MegaMoE')
# Backward-compatible alias
def fp8_fp4_mega_moe(y: torch.Tensor, def fp8_fp4_mega_moe(y: torch.Tensor,
l1_weights: Tuple[torch.Tensor, torch.Tensor], l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor], l2_weights: Tuple[torch.Tensor, torch.Tensor],
sym_buffer: SymmBuffer, sym_buffer: SymmBuffer,
cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
recipe: Tuple[int, int, int] = (1, 1, 32), recipe: Optional[Tuple[int, int, int]] = None,
activation: str = 'swiglu', activation: str = 'swiglu',
activation_clamp: Optional[float] = None, activation_clamp: Optional[float] = None,
fast_math: bool = True): fast_math: bool = True):
_C.fp8_fp4_mega_moe( fp8_mega_moe(y, l1_weights, l2_weights, sym_buffer,
y, cumulative_local_expert_recv_stats, recipe,
l1_weights, l2_weights, activation, activation_clamp, fast_math)
cumulative_local_expert_recv_stats,
sym_buffer.buffer,
sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
sym_buffer.num_max_tokens_per_rank,
sym_buffer.num_experts, sym_buffer.num_topk,
recipe,
activation, activation_clamp,
fast_math
)

110
megamoe_dev_test_scripts/phase1/interface_smoke.py Normal file
View File

@@ -0,0 +1,110 @@
import argparse
import inspect
import os
import pathlib
import sys
from typing import Tuple
import torch
import torch.distributed as dist
REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
import deep_gemm
def init_test_dist(local_rank_arg: int = None) -> Tuple[int, int, dist.ProcessGroup]:
local_rank = local_rank_arg if local_rank_arg is not None else int(os.environ.get('LOCAL_RANK', '0'))
rank = int(os.environ.get('RANK', '0'))
world_size = int(os.environ.get('WORLD_SIZE', '1'))
master_addr = os.environ.get('MASTER_ADDR', '127.0.0.1')
master_port = int(os.environ.get('MASTER_PORT', '8361'))
torch.cuda.set_device(local_rank)
sig = inspect.signature(dist.init_process_group)
params = {
'backend': 'nccl',
'init_method': f'tcp://{master_addr}:{master_port}',
'world_size': world_size,
'rank': rank,
}
if 'device_id' in sig.parameters:
params['device_id'] = torch.device(f'cuda:{local_rank}')
dist.init_process_group(**params)
torch.set_default_device('cuda')
return rank, world_size, dist.new_group(list(range(world_size)))
def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int):
l1_weights = torch.randn(
(num_experts_per_rank, intermediate_hidden * 2, hidden),
dtype=torch.float32, device='cuda').to(torch.float8_e4m3fn)
l2_weights = torch.randn(
(num_experts_per_rank, hidden, intermediate_hidden),
dtype=torch.float32, device='cuda').to(torch.float8_e4m3fn)
l1_weights_sf = torch.ones(
(num_experts_per_rank, (intermediate_hidden * 2 + 127) // 128, hidden // 128),
dtype=torch.float32, device='cuda')
l2_weights_sf = torch.ones(
(num_experts_per_rank, (hidden + 127) // 128, intermediate_hidden // 128),
dtype=torch.float32, device='cuda')
return deep_gemm.transform_weights_for_mega_moe(
(l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf))
def main() -> None:
parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 1 interface smoke')
parser.add_argument('--num-tokens', type=int, default=8)
parser.add_argument('--num-max-tokens-per-rank', type=int, default=384)
parser.add_argument('--hidden', type=int, default=512)
parser.add_argument('--intermediate-hidden', type=int, default=256)
parser.add_argument('--num-experts', type=int, default=16)
parser.add_argument('--num-topk', type=int, default=6)
parser.add_argument('--local-rank', type=int, default=None)
args = parser.parse_args()
local_rank = args.local_rank if args.local_rank is not None else int(os.environ.get('LOCAL_RANK', '0'))
rank_idx, num_ranks, group = init_test_dist(local_rank)
assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
assert args.num_experts % num_ranks == 0
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
group, args.num_experts,
args.num_max_tokens_per_rank, args.num_topk,
args.hidden, args.intermediate_hidden)
assert buffer.x.shape == (buffer.num_max_tokens_per_rank, args.hidden)
assert buffer.x_sf.shape == (buffer.num_max_tokens_per_rank, args.hidden // 128)
assert buffer.x_sf.dtype == torch.float32
assert buffer.l1_acts.shape[1] == args.hidden
assert buffer.l1_acts_sf.shape[1] == args.hidden // 128
assert buffer.l1_acts_sf.dtype == torch.float32
assert buffer.l2_acts.shape[1] == args.intermediate_hidden
assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 128
assert buffer.l2_acts_sf.dtype == torch.float32
num_tokens = args.num_tokens
buffer.x[:num_tokens].copy_(torch.randn((num_tokens, args.hidden), device='cuda').to(torch.float8_e4m3fn))
buffer.x_sf[:num_tokens].fill_(1.0)
buffer.topk_idx[:num_tokens].fill_(rank_idx * (args.num_experts // num_ranks))
buffer.topk_weights[:num_tokens].fill_(1.0)
weights = make_weights(args.num_experts // num_ranks, args.hidden, args.intermediate_hidden)
y = torch.empty((num_tokens, args.hidden), dtype=torch.bfloat16, device='cuda')
stats = torch.zeros((args.num_experts // num_ranks,), dtype=torch.int32, device='cuda')
deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer,
cumulative_local_expert_recv_stats=stats)
torch.cuda.synchronize()
dist.barrier(group=group)
buffer.destroy()
dist.destroy_process_group()
if rank_idx == 0:
print('[PASSED] SM90 MegaMoE Phase 1 interface smoke', flush=True)
if __name__ == '__main__':
main()