feat: add sm90 megamoe phase1 interfaces
This commit is contained in:
Xinyi Liu
254
AGENTS.md
Normal file
254
AGENTS.md
Normal 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"
|
||||
```
|
||||
@@ -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 开发。
|
||||
|
||||
234
csrc/apis/sm90_mega.hpp
Normal file
234
csrc/apis/sm90_mega.hpp
Normal 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
212
csrc/jit_kernels/heuristics/sm90_mega_moe.hpp
Normal 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
207
csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp
Normal 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
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
58
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
Normal file
58
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
Normal 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
|
||||
@@ -108,46 +108,46 @@ struct Workspace {
|
||||
static constexpr uint32_t kNumMaxGridSyncCounters = 4;
|
||||
|
||||
template <uint32_t kIndex = 0>
|
||||
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<uint32_t*>(base) + kIndex;
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
CUTLASS_HOST_DEVICE
|
||||
uint32_t* get_nvl_barrier_counter_ptr() const {
|
||||
return static_cast<uint32_t*>(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<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 {
|
||||
return math::advance_ptr<uint64_t>(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<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 {
|
||||
// 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<TokenSrcMetadata*>(get_src_token_topk_idx_ptr(num_experts_per_rank));
|
||||
return base + pool_token_idx;
|
||||
|
||||
199
deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh
Normal file
199
deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh
Normal 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
|
||||
@@ -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)
|
||||
|
||||
110
megamoe_dev_test_scripts/phase1/interface_smoke.py
Normal file
110
megamoe_dev_test_scripts/phase1/interface_smoke.py
Normal 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()
|
||||
Reference in New Issue
Block a user