diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..9b24abb --- /dev/null +++ b/AGENTS.md @@ -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 :/deep_gemm/include/deep_gemm/impls/ + + # 示例：传输整个目录 + scp -r deep_gemm/include/deep_gemm/impls/ :/deep_gemm/include/deep_gemm/impls/ + ``` + +3. **ssh 远程测试** — 在远程服务器上编译并运行测试，长耗时或可能 hang 的命令必须设置合理 timeout + ```bash + ssh "cd && timeout 60s python tests/test_fp8_fp4.py" + ``` + +4. **迭代** — 根据远程测试结果，在本地继续修改，重复步骤 2-3 + +### 连接配置 +- 远程主机: ``（从 `~/.ssh/config` 或环境变量读取） +- 远程路径: ``（项目在远程服务器的根目录） +- 传输前确认远程路径存在: `ssh "ls "` + +### 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/` + +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/` 并随 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/.cuh :/deep_gemm/include/deep_gemm/impls/ + +# 传输整个 impls 目录 +scp -r deep_gemm/include/deep_gemm/impls/ :/deep_gemm/include/deep_gemm/impls/ + +# 传输 host 端代码 +scp csrc/**/*.hpp :/csrc/ + +# 远程运行测试 +ssh "cd && timeout 60s python tests/test_fp8_fp4.py" + +# 远程运行单个测试函数 +ssh "cd && timeout 60s python -c 'from tests.test_fp8_fp4 import test_fp8_gemm_nt; test_fp8_gemm_nt()'" + +# 远程 build（如需重新编译 _C.so） +ssh "cd && timeout 60s bash develop.sh" + +# 远程多 rank 测试 +ssh "cd && timeout 60s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py" +``` diff --git a/MEGAMOE_SM90_DESIGN.md b/MEGAMOE_SM90_DESIGN.md index 277b900..6b55c3b 100644 --- a/MEGAMOE_SM90_DESIGN.md +++ b/MEGAMOE_SM90_DESIGN.md @@ -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.2 Barrier 初始化 | 所有 mbarrier (full/empty/dispatch/combine) 初始化 + `__syncthreads` | 无 launch failure | -| 2.3 Dispatch phase | 64-thread dispatch: expert count、source metadata、grid sync、NVLink barrier、token pull | 单 rank 模式下 `l1_token_buffer` 中数据正确 | -| 2.4 TMA A/SFA load | w2 TMA A warp：wait `l1_arrival_count`，TMA load L1 acts + float SFA to SMEM | dump SMEM A 区域，与 global 数据逐字节对比 | -| 2.5 TMA B load | w3 TMA B warp：TMA load L1 FP8 weights to SMEM | dump SMEM B 区域对比 | -| 2.6 WGMMA pipeline | WG1 math warp：`warpgroup_arrive`, 4× K-step `WGMMA::wgmma`, `warpgroup_commit_batch`, `warpgroup_wait<0>` | 单 tile (1 expert, 1 block) 的 accumulator 与 PyTorch FP32 GEMM 结果对比（未 scale） | -| 2.7 SF scaling | accum *= sfa * sfb (sfb via `__ldg`) | Scale 后的 accumulator 与 reference scaled GEMM 一致 | -| 2.8 Empty barrier | WGMMA 完成后 arrive empty_barrier → TMA 可重用 SMEM | 多 K-block 的 pipeline 不死锁 | +| 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 Dispatch barrier / workspace 初始化 | dispatch mbarrier、grid sync counter、expert count、arrival counter 初始化 | 连续两次调用无残留 workspace 污染 | +| 2.3 Expert count | 统计本 rank token/top-k 到 local experts 的 recv count，并写入 workspace | expert token counts 与 PyTorch routing reference 一致 | +| 2.4 Source metadata | 写入 `token_src_metadata[pool_token_idx] = {rank, token_idx, topk_idx}` | metadata 与 routing reference 逐项一致 | +| 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 Cross-rank dispatch smoke | NVLink barrier + remote pull 路径打通 | 2-rank / 8-rank 下无 hang，pool 内容与 reference 一致 | +| 2.7 `l1_arrival_count` release | 每个 pool block dispatch 完成后 release-add arrival count | arrival count 达到 block 内 token 数，TMA consumer 可安全等待 | + +**验证里程碑**：单 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)`。 -**依赖**：Phase 1 完成 +**依赖**：Phase 2 完成 **关键挑战**： - WGMMA descriptor 构造（参考 `deep_gemm/mma/sm90.cuh` 中的 `make_gmma_desc`） - 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 写入一致 | -| 3.2 SwiGLU 计算 | gate/up deinterleave + `silu(gate) * up * topk_weight` in registers | 与 reference SwiGLU 一致（精度：atol=0.02） | -| 3.3 Activation clamp | `kActivationClamp` 模板参数控制 clamp | clamp 后无超范围值 | -| 3.4 fast_math SwiGLU | `__expf` + `fast_rcp` 快速路径 | 与精确版相差 < 0.5% | -| 3.5 Per-row amax | `warp_reduce<4>(max(abs(...)))` 跨 WG 内 4 warps reduce | amax 值与 PyTorch `abs().max(dim=-1)` 一致 | -| 3.6 Cooperative amax | 两个 WG 独立 amax（各自 64 行），SMEM scratch 区域用于跨 warp reduce | 每行 amax 正确 | -| 3.7 FP8 quantize | `sf = amax / 448.0`，`fp8 = cast(result / sf)` | 反量化后与 reference 误差 < 1% | -| 3.8 Float SF write | 写入 `l2_acts_sf` (MN-major, per-64-K layout) | SF 值正确、layout 与 L2 TMA load 兼容 | -| 3.9 STSM + double-buf TMA store | STSM 到 `smem_cd[stage]`，TMA store 到 `l2_acts` | `l2_acts` 内容与 reference FP8 quant output 一致 | -| 3.10 `l2_arrival_mask` 设置 | `red_or_rel_gpu(mask, 1 << n_block_idx)` | 单 expert 全部 L1 N-blocks 完成后 mask == expected | -| 3.11 Cooperative sync | `sync_aligned(256)` 等两个 WG 都 TMA store 完成后设置 mask | 无 race condition | +| 4.1 Top-k weight load | 从 `l1_topk_weights_buffer` load weight per row | 值与 dispatch 写入一致 | +| 4.2 SwiGLU 计算 | gate/up deinterleave + `silu(gate) * up * topk_weight` in registers | 与 reference SwiGLU 一致（精度：atol=0.02） | +| 4.3 Activation clamp | `kActivationClamp` 模板参数控制 clamp | clamp 后无超范围值 | +| 4.4 fast_math SwiGLU | `__expf` + `fast_rcp` 快速路径 | 与精确版相差 < 0.5% | +| 4.5 Per-row amax | `warp_reduce<4>(max(abs(...)))` 跨 WG 内 4 warps reduce | amax 值与 PyTorch `abs().max(dim=-1)` 一致 | +| 4.6 Cooperative amax | 两个 WG 独立 amax（各自 64 行），SMEM scratch 区域用于跨 warp reduce | 每行 amax 正确 | +| 4.7 FP8 quantize | `sf = amax / 448.0`，`fp8 = cast(result / sf)` | 反量化后与 reference 误差 < 1% | +| 4.8 Float SF write | 写入 `l2_acts_sf` (MN-major, per-64-K layout) | SF 值正确、layout 与 L2 TMA load 兼容 | +| 4.9 STSM + double-buf TMA store | STSM 到 `smem_cd[stage]`，TMA store 到 `l2_acts` | `l2_acts` 内容与 reference FP8 quant output 一致 | +| 4.10 `l2_arrival_mask` 设置 | `red_or_rel_gpu(mask, 1 << n_block_idx)` | 单 expert 全部 L1 N-blocks 完成后 mask == expected | +| 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` 正确触发。 -**依赖**：Phase 2 完成 +**依赖**：Phase 3 完成 **关键挑战**： - WGMMA register layout 到 gate/up 列的映射关系（参考 SM100 的 `SM100_TMEM_LOAD_16dp256b1x` 读取模式，SM90 的 register layout 不同） - 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。 | 步骤 | 交付物 | 验证标准 | |------|--------|----------| -| 4.1 L2 TMA wait | TMA A warp wait `l2_arrival_mask == expected` | 正确等待 L1 所有 N-blocks 完成 | -| 4.2 L2 TMA load | TMA load `l2_acts` + 两个 per-64 SFA halves | SMEM 数据正确 | -| 4.3 L2 WGMMA dual-half | K-step 0-1 累积 half0，K-step 2-3 累积 half1 | 两组 partial accum 正确 | -| 4.4 Dual-half SF apply | `final = half0 * sfa0 * sfb + half1 * sfa1 * sfb` | 与 reference L2 dequant GEMM 一致 | -| 4.5 Weight SF `__ldg` | Math WG prefetch L2 weight SF from global | 值正确 | -| 4.6 BF16 cast + STSM | FP32 → BF16 cast，STSM 到 `smem_cd_l2` (BF16 swizzle) | SMEM 内容正确 | -| 4.7 Source metadata read | `token_src_metadata[pool_token_idx + row]` → {rank, token, topk} | 值与 dispatch 写入一致 | -| 4.8 NVLink scatter | `ld_shared` + `sym_buffer.map` + remote store | 单 rank 下本地 combine buffer 数据正确 | -| 4.9 Cooperative smem_cd guard | `sync_aligned(256)` 防止下一 tile 覆盖 smem_cd | 多 tile 序列无 data corruption | -| 4.10 Multi-expert persistent loop | Scheduler `for_each_block` 遍历多 expert 的 L1→L2 cycle | 全量 expert 输出正确 | +| 5.1 L2 TMA wait | TMA A warp wait `l2_arrival_mask == expected` | 正确等待 L1 所有 N-blocks 完成 | +| 5.2 L2 TMA load | TMA load `l2_acts` + 两个 per-64 SFA halves | SMEM 数据正确 | +| 5.3 L2 WGMMA dual-half | K-step 0-1 累积 half0，K-step 2-3 累积 half1 | 两组 partial accum 正确 | +| 5.4 Dual-half SF apply | `final = half0 * sfa0 * sfb + half1 * sfa1 * sfb` | 与 reference L2 dequant GEMM 一致 | +| 5.5 Weight SF `__ldg` | Math WG prefetch L2 weight SF from global | 值正确 | +| 5.6 BF16 cast + STSM | FP32 → BF16 cast，STSM 到 `smem_cd_l2` (BF16 swizzle) | SMEM 内容正确 | +| 5.7 Source metadata read | `token_src_metadata[pool_token_idx + row]` → {rank, token, topk} | 值与 dispatch 写入一致 | +| 5.8 NVLink scatter | `ld_shared` + `sym_buffer.map` + remote store | 单 rank 下本地 combine buffer 数据正确 | +| 5.9 Cooperative smem_cd guard | `sync_aligned(256)` 防止下一 tile 覆盖 smem_cd | 多 tile 序列无 data corruption | +| 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 依赖。 -**依赖**：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 最紧张的地方** - 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。 | 步骤 | 交付物 | 验证标准 | |------|--------|----------| -| 5.1 NVLink barrier | Epilogue grid sync + cross-rank signal | 多 rank 下所有 rank 到达 barrier | -| 5.2 Dispatch↔Epilogue handshake | `kDispatchWithEpilogueBarrierIdx` sync | Dispatch 和 epilogue 阶段正确交替 | -| 5.3 Combine TMA load | Double-buffered TMA 从 combine buffer load BF16 chunks | Load 数据与 scatter 写入一致 | -| 5.4 FP32 accumulate | 逐 top-k slot 累加到 float registers | 累加结果与 reference `sum(topk_weights * expert_outputs)` 一致 | -| 5.5 BF16 cast + TMA store | Cast + TMA store 到 output `y` | `y` 与 reference 在 BF16 精度内一致 | -| 5.6 Workspace cleanup | Dispatch warps 清理 expert counts、arrival counters、metadata | 下次 kernel 调用前 workspace 已归零 | -| 5.7 NVLink cleanup barrier | 跨 rank 确认 workspace 清理完成 | 连续两次 kernel 调用无 data corruption | -| 5.8 端到端 correctness | 完整 `fp8_mega_moe(...)` 调用 | `test_mega_moe_sm90.py` Smoke test 通过 (2 ranks, 4096×2048, topk=6) | +| 6.1 NVLink barrier | Epilogue grid sync + cross-rank signal | 多 rank 下所有 rank 到达 barrier | +| 6.2 Dispatch↔Epilogue handshake | `kDispatchWithEpilogueBarrierIdx` sync | Dispatch 和 epilogue 阶段正确交替 | +| 6.3 Combine TMA load | Double-buffered TMA 从 combine buffer load BF16 chunks | Load 数据与 scatter 写入一致 | +| 6.4 FP32 accumulate | 逐 top-k slot 累加到 float registers | 累加结果与 reference `sum(topk_weights * expert_outputs)` 一致 | +| 6.5 BF16 cast + TMA store | Cast + TMA store 到 output `y` | `y` 与 reference 在 BF16 精度内一致 | +| 6.6 Workspace cleanup | Dispatch warps 清理 expert counts、arrival counters、metadata | 下次 kernel 调用前 workspace 已归零 | +| 6.7 NVLink cleanup barrier | 跨 rank 确认 workspace 清理完成 | 连续两次 kernel 调用无 data corruption | +| 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)` 通过。 -**依赖**：Phase 4 完成 +**依赖**：Phase 5 完成 **关键挑战**： - Combine 逻辑直接移植自 SM100 (与 TMEM/UMMA 无关)，但需要确认 SM90 下 TMA 1D load/store 的行为 - 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 场景。 | 步骤 | 交付物 | 验证标准 | |------|--------|----------| -| 6.1 模板分支 | `if constexpr (kCooperativeMode)` 控制 WG2 行为 | 两种模式均编译通过 | -| 6.2 WG2 idle | WG2 在 `kCooperativeMode=false` 时 `warpgroup_reg_dealloc` 后 idle | 无 deadlock / launch failure | -| 6.3 TMA A 调整 | TMA A 只加载 BLOCK_M 行（64 或 32）而非 128 | SMEM A 数据正确 | -| 6.4 WGMMA 有效行 | `valid_m = min(tokens_in_block, BLOCK_M)` 控制 epilogue 写入范围 | 无 OOB 写入 | -| 6.5 Scheduler 适配 | 动态 BLOCK_M 下 scheduler 的 block 分配和 pool_block_offset 正确 | 全量 expert 覆盖，无漏 tile | -| 6.6 Heuristics JIT 分发 | 根据 `expected_tokens_per_expert` 选择 BLOCK_M，JIT 编译对应模板 | 正确选择 config 并编译 | -| 6.7 Small M correctness | BLOCK_M=64: 单 WG 处理完整 64 行 | Smoke test 通过 | -| 6.8 Extreme decode | BLOCK_M=32: token-per-expert=1-2 | 极端 decode 场景输出正确 | +| 7.1 模板分支 | `if constexpr (kCooperativeMode)` 控制 WG2 行为 | 两种模式均编译通过 | +| 7.2 WG2 idle | WG2 在 `kCooperativeMode=false` 时 `warpgroup_reg_dealloc` 后 idle | 无 deadlock / launch failure | +| 7.3 TMA A 调整 | TMA A 只加载 BLOCK_M 行（64 或 32）而非 128 | SMEM A 数据正确 | +| 7.4 WGMMA 有效行 | `valid_m = min(tokens_in_block, BLOCK_M)` 控制 epilogue 写入范围 | 无 OOB 写入 | +| 7.5 Scheduler 适配 | 动态 BLOCK_M 下 scheduler 的 block 分配和 pool_block_offset 正确 | 全量 expert 覆盖，无漏 tile | +| 7.6 Heuristics JIT 分发 | 根据 `expected_tokens_per_expert` 选择 BLOCK_M，JIT 编译对应模板 | 正确选择 config 并编译 | +| 7.7 Small M correctness | BLOCK_M=64: 单 WG 处理完整 64 行 | Smoke test 通过 | +| 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 正确切换。 -**依赖**：Phase 5 完成（cooperative 端到端已验证） +**依赖**：Phase 6 完成（cooperative 端到端已验证） **关键挑战**： - BLOCK_M=32 时 WGMMA 计算 64 行但只用 32 行的 register 浪费 - 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 正确性和性能。 | 步骤 | 交付物 | 验证标准 | |------|--------|----------| -| 7.1 Multi-rank launch | 8-rank torchrun 启动 MegaMoE kernel | 无 hang / crash | -| 7.2 Correctness sweep | Token sweep (1-8192) × Shape sweep × BLOCK_M 三档 | 全部 `allclose` 通过 | -| 7.3 Edge cases | 0 tokens, unbalanced routing, single-expert-per-token | 无 crash，输出正确 | -| 7.4 Benchmark framework | `bench_mega_moe_sm90.py`：latency, TFLOPS, HBM GB/s | 可生成 performance table | -| 7.5 nsys timeline | Phase profile: dispatch / L1 / L2 / combine 各阶段时间 | 识别主要 stall 来源 | -| 7.6 DeepEP v1 baseline | 对比 DeepEP dispatch + grouped GEMM + combine | 输出 speedup ratio | -| 7.7 Regression test | CI 集成到 `tests/test_mega_moe_sm90.py` | `pytest` 通过 | +| 8.1 Multi-rank launch | 8-rank torchrun 启动 MegaMoE kernel | 无 hang / crash | +| 8.2 Correctness sweep | Token sweep (1-8192) × Shape sweep × BLOCK_M 三档 | 全部 `allclose` 通过 | +| 8.3 Edge cases | 0 tokens, unbalanced routing, single-expert-per-token | 无 crash，输出正确 | +| 8.4 Benchmark framework | `bench_mega_moe_sm90.py`：latency, TFLOPS, HBM GB/s | 可生成 performance table | +| 8.5 nsys timeline | Phase profile: dispatch / L1 / L2 / combine 各阶段时间 | 识别主要 stall 来源 | +| 8.6 DeepEP v1 baseline | 对比 DeepEP dispatch + grouped GEMM + combine | 输出 speedup ratio | +| 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)。 -**依赖**：Phase 6 完成 +**依赖**：Phase 7 完成 --- -### Phase 8: 性能优化 (Day 63-75) +### Phase 9: 性能优化 (Day 63-75) **目标**：基于 nsys timeline 分析，针对瓶颈进行优化迭代。 | 步骤 | 优化方向 | 预期收益 | 依据 | |------|---------|---------|------| -| 8.1 N-major scheduling tuning | 调整 `tokens/expert >= 256` 阈值 | 大 M 下 weight L2 cache 命中率提升 | NCU L2 sector miss 数据 | -| 8.2 SFA 软件 prefetch | L2 per-64 SFA 的 `__ldg` prefetch 提前一个 K-block | 减少 math stall waiting on SF | nsys timeline 中 SF load 占比 | -| 8.3 Combine chunk 大小调优 | 根据 hidden dim 选择 1 vs 2 chunks | 减少 combine 阶段 TMA round-trip | combine 占总时间比例 | -| 8.4 Dispatch pull 优化 | 增加 dispatch warp 内的 token parallelism | 减少 dispatch phase 时间 | nsys dispatch 耗时 | -| 8.5 Register spill 消除 | 检查 SASS 中的 local memory access，调整变量生命周期 | 消除 spill 带来的 DRAM 流量 | cuobjdump --dump-sass | -| 8.6 wave_count 调优 | 根据实测 per-shape 最优 experts_per_wave | 减少 tail effect | benchmark sweep | -| 8.7 Optional: cluster=2 探索 | TMA multicast B tile（需解决 amax 跨 CTA sync） | weight HBM 进一步减半 | 若 cooperative 的 HBM 仍是瓶颈 | +| 9.1 N-major scheduling tuning | 调整 `tokens/expert >= 256` 阈值 | 大 M 下 weight L2 cache 命中率提升 | NCU L2 sector miss 数据 | +| 9.2 SFA 软件 prefetch | L2 per-64 SFA 的 `__ldg` prefetch 提前一个 K-block | 减少 math stall waiting on SF | nsys timeline 中 SF load 占比 | +| 9.3 Combine chunk 大小调优 | 根据 hidden dim 选择 1 vs 2 chunks | 减少 combine 阶段 TMA round-trip | combine 占总时间比例 | +| 9.4 Dispatch pull 优化 | 增加 dispatch warp 内的 token parallelism | 减少 dispatch phase 时间 | nsys dispatch 耗时 | +| 9.5 Register spill 消除 | 检查 SASS 中的 local memory access，调整变量生命周期 | 消除 spill 带来的 DRAM 流量 | cuobjdump --dump-sass | +| 9.6 wave_count 调优 | 根据实测 per-shape 最优 experts_per_wave | 减少 tail effect | benchmark sweep | +| 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 (基础设施) │ v -Phase 2 (Dispatch + L1 GEMM) +Phase 2 (Dispatch) │ v -Phase 3 (L1 Epilogue) +Phase 3 (L1 TMA + WGMMA) │ v -Phase 4 (L2 GEMM + Scatter) +Phase 4 (L1 Epilogue) │ v -Phase 5 (Combine + 端到端) +Phase 5 (L2 GEMM + Scatter) + │ + v +Phase 6 (Combine + 端到端) │ ├──────────────────┐ v v -Phase 6 (Single-WG) Phase 7 (多 Rank 集成) +Phase 7 (Single-WG) Phase 8 (多 Rank 集成) │ │ └──────┬───────────┘ v - Phase 8 (性能优化) + Phase 9 (性能优化) ``` ### 时间线总结 @@ -1010,12 +1037,13 @@ Phase 6 (Single-WG) Phase 7 (多 Rank 集成) |-------|------|------|---------| | 0 | 3 | 3 | 环境 + baseline | | 1 | 7 | 10 | 基础设施 (可 JIT 编译) | -| 2 | 10 | 20 | Dispatch + L1 WGMMA 正确 | -| 3 | 10 | 30 | L1 SwiGLU + FP8 quant 正确 | -| 4 | 10 | 40 | L2 GEMM + NVLink scatter 正确 | -| 5 | 8 | 48 | 端到端 fused kernel 正确 (2-rank) | -| 6 | 7 | 55 | 动态 BLOCK_M 全变体正确 | -| 7 | 7 | 62 | 8-rank 集成 + benchmark | -| 8 | 13 | 75 | 性能调优完成 | +| 2 | 5 | 15 | Dispatch + L1 pool 正确 | +| 3 | 5 | 20 | L1 TMA + WGMMA 正确 | +| 4 | 10 | 30 | L1 SwiGLU + FP8 quant 正确 | +| 5 | 10 | 40 | L2 GEMM + NVLink scatter 正确 | +| 6 | 8 | 48 | 端到端 fused kernel 正确 (2-rank) | +| 7 | 7 | 55 | 动态 BLOCK_M 全变体正确 | +| 8 | 7 | 62 | 8-rank 集成 + benchmark | +| 9 | 13 | 75 | 性能调优完成 | -**总工期**：约 75 工作日（10-11 周），含充分的验证和调试时间。关键路径是 Phase 2-5 的串行 kernel 开发（40 天）。 +**总工期**：约 75 工作日（10-11 周），含充分的验证和调试时间。关键路径是 Phase 2-6 的串行 kernel 开发。 diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp new file mode 100644 index 0000000..a68a554 --- /dev/null +++ b/csrc/apis/sm90_mega.hpp @@ -0,0 +1,234 @@ +#pragma once + +#include +#include + +#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> +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(sizeof(float)), false); + const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 128 * static_cast(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(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(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(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(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(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(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(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(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(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(combine_token_buffer.get_end_ptr()), slice_input_buffers}; +} + +static void fp8_mega_moe( + const torch::Tensor& y, + const std::tuple& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, const int& rank_idx, + const int& num_max_tokens_per_rank, + const int& num_experts, const int& num_topk, + const std::tuple& recipe, + const std::string& activation, + const std::optional& 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(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::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(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(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("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 diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp new file mode 100644 index 0000000..0c92096 --- /dev/null +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -0,0 +1,212 @@ +#pragma once + +#include +#include + +#include + +#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 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(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(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(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 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(sizeof(uint32_t)), kSmemAlignment); + const int smem_send_buffers_size = align( + static_cast(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(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(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(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("DG_SM90_MOE_NMAJOR", -1); + if (env_val != -1) + return env_val > 0; + float expected = static_cast(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("DG_JIT_DEBUG") or get_env("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 printed; + if (printed.count(key) == 0) { + std::cout << key << ": " << config << std::endl; + printed.insert(key); + } + } + return config; +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp new file mode 100644 index 0000000..c826118 --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -0,0 +1,207 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "runtime_utils.hpp" + +#include +#include + +#include "../heuristics/sm90_mega_moe.hpp" + +namespace deep_gemm { + +class SM90FP8MegaMoERuntime final : public LaunchRuntime { +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 + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&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 cumulative_local_expert_recv_stats, + const std::vector& 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(sym_buffer_ptrs.size()); + const auto num_experts = num_experts_per_rank * num_ranks; + const auto num_padded_sf_pool_tokens = static_cast(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(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(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(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(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(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(); + + // 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 diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp index a966afe..55c0fa2 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -7,6 +7,7 @@ #include "apis/gemm.hpp" #include "apis/layout.hpp" #include "apis/mega.hpp" +#include "apis/sm90_mega.hpp" #include "apis/runtime.hpp" #ifndef TORCH_EXTENSION_NAME @@ -24,5 +25,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { deep_gemm::gemm::register_apis(m); deep_gemm::layout::register_apis(m); deep_gemm::mega::register_apis(m); + deep_gemm::mega::register_sm90_apis(m); deep_gemm::runtime::register_apis(m); } diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index a9542e2..aad7df2 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -85,6 +85,7 @@ from .mega import ( SymmBuffer, get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe, + fp8_mega_moe, fp8_fp4_mega_moe, ) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh new file mode 100644 index 0000000..74cbd72 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -0,0 +1,58 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +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 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 diff --git a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh index 13520c6..c66eecd 100644 --- a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh @@ -108,46 +108,46 @@ struct Workspace { static constexpr uint32_t kNumMaxGridSyncCounters = 4; template - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE uint32_t* get_grid_sync_count_ptr() const { DG_STATIC_ASSERT(kIndex < kNumMaxGridSyncCounters, "Grid sync index out of bounds"); return static_cast(base) + kIndex; } - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE uint32_t* get_nvl_barrier_counter_ptr() const { return static_cast(base) + kNumMaxGridSyncCounters; } - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE int* get_nvl_barrier_signal_ptr(const uint32_t& phase) const { // NOTES: the signal is signed, as we may minus return math::advance_ptr(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 { return math::advance_ptr(base, kNumBarrierSignalBytes) + expert_idx; } - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE uint64_t* get_expert_recv_count_ptr( 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; } - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE 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; } - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE 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); return reinterpret_cast(base) + pool_block_idx; } - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE 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 const auto base = get_l1_arrival_count_ptr(math::align(num_max_pool_blocks, 2u)); @@ -155,7 +155,7 @@ struct Workspace { } // For dispatch pulling - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE 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 auto base = get_l2_arrival_mask_ptr(num_max_pool_blocks); @@ -165,7 +165,7 @@ struct Workspace { } // For combine usages - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE TokenSrcMetadata* get_token_src_metadata_ptr(const uint32_t& pool_token_idx = 0) const { const auto base = reinterpret_cast(get_src_token_topk_idx_ptr(num_experts_per_rank)); return base + pool_token_idx; diff --git a/deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh b/deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh new file mode 100644 index 0000000..4d546e7 --- /dev/null +++ b/deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh @@ -0,0 +1,199 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace deep_gemm::sched { + +enum class SM90BlockPhase { + None = 0, + Linear1 = 1, + Linear2 = 2 +}; + +template +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 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(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(value >> 32) != kNumSMs * kNumRanks); + } + stored_num_tokens_per_expert[i] = static_cast(value); + } + __syncwarp(); + } + + template + 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 diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index e624ecf..4c7e19d 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -13,6 +13,14 @@ except Exception as exception: 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: def __init__(self, group: dist.ProcessGroup, # MoE arguments @@ -28,13 +36,23 @@ class SymmBuffer: self.hidden = hidden self.intermediate_hidden = intermediate_hidden - # Allocate a symmetric buffer - 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 - ) + # Allocate a symmetric buffer (route by architecture) + if _is_sm90(): + num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe( + group.size(), num_experts, + num_max_tokens_per_rank, num_topk, + 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.handle = symm_mem.rendezvous(self.buffer, group=group) self.buffer.zero_() @@ -46,6 +64,10 @@ class SymmBuffer: self.topk_idx, self.topk_weights, self.l1_acts, self.l1_acts_sf, 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): self.handle = None @@ -53,6 +75,16 @@ class SymmBuffer: self.group = None self.x = 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, @@ -62,7 +94,13 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, use_fp8_dispatch: bool = True, activation: str = 'swiglu') -> SymmBuffer: # 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( 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] - def interleave(t, gran: int = 8) -> torch.Tensor: - g, n, *rest = t.shape - half = n // 2 - gate = 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)) + g, n, *rest = t.shape + half = n // 2 + gate = 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 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: @@ -93,36 +132,82 @@ def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor: 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( 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 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_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])) 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, 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: Tuple[int, int, int] = (1, 1, 32), + recipe: Optional[Tuple[int, int, int]] = None, activation: str = 'swiglu', activation_clamp: Optional[float] = None, fast_math: bool = True): - _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 - ) + fp8_mega_moe(y, l1_weights, l2_weights, sym_buffer, + cumulative_local_expert_recv_stats, recipe, + activation, activation_clamp, fast_math) diff --git a/megamoe_dev_test_scripts/phase1/interface_smoke.py b/megamoe_dev_test_scripts/phase1/interface_smoke.py new file mode 100644 index 0000000..b118456 --- /dev/null +++ b/megamoe_dev_test_scripts/phase1/interface_smoke.py @@ -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()