284 lines
32 KiB
Markdown
284 lines
32 KiB
Markdown
# SM90 MegaMoE 开发日志
|
||
|
||
## 2026-06-17 23:55:32 CST
|
||
|
||
- 对应 clean 工作 commit:`1dafd0f17376d8730bebb472dc8e15d9645c6b52`
|
||
- 用户请求摘要:重做 Phase 1 接口/基础设施提交,保证它作为独立 clean commit 在线性历史中出现,并通过远端测试。
|
||
- 本次提交的核心改动:
|
||
- 新增 SM90 MegaMoE host API:`get_symm_buffer_size_for_sm90_mega_moe`、`get_token_alignment_for_sm90_mega_moe`、`fp8_mega_moe`。
|
||
- 新增 SM90 MegaMoE JIT heuristics、runtime launch glue、scheduler 和空 kernel stub。
|
||
- Python MegaMoE wrapper 支持按当前 CUDA device 区分 SM90/SM100,SM90 使用 `(1, 128, 128)` recipe 和 float SF 路径。
|
||
- SM90 symmetric buffer Phase 1 接口返回 8 个常规 view:`x`、`x_sf`、`topk_idx`、`topk_weights`、`l1_acts`、`l1_acts_sf`、`l2_acts`、`l2_acts_sf`。
|
||
- 新增 Phase 1 smoke 测试脚本:`megamoe_dev_test_scripts/phase1/interface_smoke.py`。
|
||
- 更新 `AGENTS.md`,记录 phase 测试脚本提交规范、开发日志规范和 60 秒左右 timeout 规则。
|
||
- 更新 `MEGAMOE_SM90_DESIGN.md`,将 Dispatch 与 L1 TMA/WGMMA 拆成独立 phase。
|
||
- 关键文件列表:
|
||
- `AGENTS.md`
|
||
- `MEGAMOE_SM90_DESIGN.md`
|
||
- `csrc/apis/sm90_mega.hpp`
|
||
- `csrc/jit_kernels/heuristics/sm90_mega_moe.hpp`
|
||
- `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp`
|
||
- `csrc/python_api.cpp`
|
||
- `deep_gemm/__init__.py`
|
||
- `deep_gemm/mega/__init__.py`
|
||
- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh`
|
||
- `deep_gemm/include/deep_gemm/layout/mega_moe.cuh`
|
||
- `deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh`
|
||
- `megamoe_dev_test_scripts/phase1/interface_smoke.py`
|
||
- 详细开发流程:
|
||
- 本地从原混合 Phase 1/Phase 2 commit 回退到 Phase 0 基线,保留工作区改动并重新拆分 commit。
|
||
- 本地将 SM90 kernel 降到 Phase 1 空 kernel stub,仅保留 ABI、JIT 模板参数和 launch 静态检查。
|
||
- 本地将 `get_symm_buffer_size_for_sm90_mega_moe` 降到 Phase 1 接口,返回 8 个常规 buffer views,不暴露 Phase 2 dispatch 内部状态。
|
||
- 本地新增 `megamoe_dev_test_scripts/phase1/interface_smoke.py`,验证 SM90 buffer views、float SF dtype、weight transform 和空 kernel launch。
|
||
- 本地静态检查:`git diff --check`、`python3 -m py_compile megamoe_dev_test_scripts/phase1/interface_smoke.py`。
|
||
- 远端同步:`rsync -avR AGENTS.md MEGAMOE_SM90_DESIGN.md csrc/python_api.cpp deep_gemm/__init__.py deep_gemm/mega/__init__.py deep_gemm/include/deep_gemm/layout/mega_moe.cuh csrc/apis/sm90_mega.hpp csrc/jit_kernels/heuristics/sm90_mega_moe.hpp csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh megamoe_dev_test_scripts/phase1/interface_smoke.py g0034:/mnt/beegfs/lxy/DeepGEMM/`。
|
||
- 远端 build:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s bash develop.sh'"`,通过。
|
||
- 远端 single-rank smoke:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s python3 megamoe_dev_test_scripts/phase1/interface_smoke.py'"`,通过。
|
||
- 远端 2-rank smoke:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && CUDA_VISIBLE_DEVICES=0,1 timeout 60s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase1/interface_smoke.py'"`,通过。
|
||
- 测试结果和已知问题:
|
||
- Phase 1 build 和 smoke 测试均通过。
|
||
- single-rank 仍出现 PyTorch symmetric memory warning:`cuMulticastCreate failed with: "invalid argument"`,当前不影响 smoke。
|
||
- Phase 1 kernel 是空 kernel,不产生 dispatch/GEMM 输出;正确性只覆盖接口、buffer views、JIT 编译与 launch。
|
||
- 后续待办:
|
||
- 在线性历史中继续提交 Phase 2 dispatch-only 实现。
|
||
- Phase 3 开始实现 L1 TMA + 单 tile WGMMA,并增加 accumulator/reference 对比。
|
||
|
||
## 2026-06-18 00:00:45 CST
|
||
|
||
- 对应 clean 工作 commit:`540e5aeadc1b21dca8222a838ed6e5b6c4de93e8`
|
||
- 用户请求摘要:在线性 Phase 1 commit 之后提交 Phase 2 dispatch-only 实现,并重新通过远端 single-rank、2-rank、8-rank correctness 测试。
|
||
- 本次提交的核心改动:
|
||
- 将 SM90 kernel 从 Phase 1 空 kernel 扩展为 Phase 2 dispatch-only kernel。
|
||
- 实现 dispatch workspace 清理、expert count、source token-topk index 写入、NVLink barrier、remote/local token/SF/top-k weight pull、token source metadata 写入、`l1_arrival_count` release。
|
||
- 扩展 SM90 symmetric buffer views,额外暴露 `l1_topk_weights`、`expert_recv_count_sum`、`l1_arrival_count`、`token_src_metadata` 供 Phase 2 correctness 检查。
|
||
- 更新 Python `SymmBuffer`,SM90 分支解包 12 个 views,SM100 分支保持 8 个 views。
|
||
- 新增 Phase 2 correctness 脚本 `megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py`,使用 torchrun-compatible distributed init。
|
||
- 关键文件列表:
|
||
- `csrc/apis/sm90_mega.hpp`
|
||
- `deep_gemm/mega/__init__.py`
|
||
- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh`
|
||
- `megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py`
|
||
- 详细开发流程:
|
||
- 在 Phase 1 commit 和 Phase 1 日志 commit 之后,从原 dispatch-only 实现恢复 Phase 2 相关文件,保证历史顺序为 Phase 1 -> Phase 1 log -> Phase 2。
|
||
- 本地静态检查:`python3 -m py_compile megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py`、`git diff --check`。
|
||
- 远端同步:`rsync -avR csrc/apis/sm90_mega.hpp deep_gemm/mega/__init__.py deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py g0034:/mnt/beegfs/lxy/DeepGEMM/`。
|
||
- 远端 build:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s bash develop.sh'"`,通过。
|
||
- 远端 single-rank correctness:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s python3 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py --tokens-list 0,8,48,192'"`,通过。
|
||
- 远端 2-rank correctness:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && CUDA_VISIBLE_DEVICES=0,1 timeout 60s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py --tokens-list 0,8,48,192'"`,通过。
|
||
- 远端 8-rank correctness:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 timeout 60s torchrun --standalone --nproc_per_node=8 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py --tokens-list 0,8,48,192'"`,通过。
|
||
- 确认远端无残留 `dispatch_only_correctness` / `interface_smoke` / `torchrun` 进程。
|
||
- 测试结果和已知问题:
|
||
- Phase 2 build 通过。
|
||
- Phase 2 dispatch-only correctness 在 single-rank、2-rank、8-rank 下通过。
|
||
- 覆盖 `num_tokens=0,8,48,192`,并覆盖 `BLOCK_M=32/64/128`。
|
||
- single-rank 仍出现 PyTorch symmetric memory warning:`cuMulticastCreate failed with: "invalid argument"`,当前不影响 correctness。
|
||
- 后续待办:
|
||
- Phase 3:在 Phase 2 dispatch 产物上实现 L1 TMA + 单 tile WGMMA。
|
||
- 为 Phase 3 增加 accumulator / scaled GEMM reference 对比测试。
|
||
|
||
## 2026-06-18 00:46:51 CST
|
||
|
||
- 对应 clean 工作 commit:`842124b1c588d78a9fc4ff7fbdaf2b7b80e5287c`
|
||
- 用户请求摘要:确认 Phase 3 verification buffer 方案后,实施 SM90 MegaMoE Phase 3 L1 TMA + 单 tile WGMMA,并完成远端 build/correctness/regression 验证。
|
||
- 本次提交的核心改动:
|
||
- 扩展 SM90 symmetric buffer views,新增 `l1_accum_debug: float32[128, 128]`,用于 Phase 3 单 tile FP32 accumulator correctness 验证。
|
||
- SM90 Python `SymmBuffer` 暴露 `buffer.l1_accum_debug`,SM100 分支保持兼容的 `None` 属性。
|
||
- SM90 JIT/runtime 增加 `l1_accum_debug` kernel 参数。
|
||
- SM90 kernel 保留 Phase 2 dispatch 路径,并新增 `BLOCK_M == 128` 下的 L1 TMA producer 和 WGMMA consumer:TMA load L1 acts/SFA/L1 weights,执行 FP8xFP8 WGMMA,应用 activation SF 和 weight SF,并写出首个 L1 tile 的 scaled FP32 accumulator。
|
||
- Phase 3 barrier/debug 初始化仅在 `BLOCK_M == 128` 启用,dispatch named barrier 从 `0` 调整到 `1`,避免与 full-CTA `__syncthreads()` 使用的 barrier 冲突并修复 Phase 2 小 `BLOCK_M` 回归。
|
||
- 新增 Phase 3 correctness 脚本 `megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py`,基于实际 dispatch pool 计算 PyTorch FP32 reference 并对比 `l1_accum_debug`。
|
||
- 关键文件列表:
|
||
- `csrc/apis/sm90_mega.hpp`
|
||
- `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp`
|
||
- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh`
|
||
- `deep_gemm/mega/__init__.py`
|
||
- `megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py`
|
||
- 详细开发流程:
|
||
- 本地扩展 SM90 buffer layout 和 Python wrapper,新增 committed verification view `l1_accum_debug`。
|
||
- 本地在 SM90 kernel 中移除 Phase 2 的非 dispatch thread 早退,改为 dispatch/TMA/math role branches;Phase 3 computation 只在 `BLOCK_M == 128` 路径启用。
|
||
- 本地实现 Phase 3 shared memory layout、full/empty barrier 初始化、L1 acts/SFA/weights TMA load、SM90 FP8 WGMMA、float SF scaling 和 debug accumulator 写出。
|
||
- 本地新增 `megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py`,固定 `num_tokens=128`、`num_topk=1`,从 `buffer.l1_acts` / `buffer.l1_acts_sf` / transformed L1 weights 计算 reference。
|
||
- 本地静态检查:`python3 -m py_compile megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py`、`git diff --check`。
|
||
- 远端路径确认:`ssh g0034 "docker exec sglang-lxy bash -lc 'ls /mnt/beegfs/lxy/DeepGEMM'"`。
|
||
- 远端同步:`rsync -avz --relative csrc/apis/sm90_mega.hpp csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh deep_gemm/mega/__init__.py megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py g0034:/mnt/beegfs/lxy/DeepGEMM/`。
|
||
- 远端首次 build 暴露远端 `csrc/apis/mega.hpp` stale copy 与本地不一致,随后同步本地 `csrc/apis/mega.hpp` 到远端用于恢复远端 workspace 一致性;该文件本次 clean work commit 未改动。
|
||
- 远端 build:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s bash develop.sh'"`,通过。
|
||
- 远端 Phase 3 single-rank correctness:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s python3 megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py'"`,通过,`max_diff=0.006541`。
|
||
- 远端 Phase 3 2-rank correctness/no-hang:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && CUDA_VISIBLE_DEVICES=0,1 timeout 60s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py'"`,通过,rank0 输出 `max_diff=0.006541`。
|
||
- 远端 Phase 2 regression 初次发现小 `BLOCK_M` dispatch-only case 在新增 full-CTA barrier 后出现 illegal memory access;本地修复为 Phase 3 barrier/debug 初始化只在 `BLOCK_M == 128` 启用,并将 dispatch named barrier 改为 `1`。
|
||
- 修复后重新同步 kernel 文件,并重新执行远端 build,通过。
|
||
- 远端 Phase 2 single-rank regression:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s python3 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py --tokens-list 192'"`,通过。
|
||
- 远端 Phase 2 2-rank regression:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && CUDA_VISIBLE_DEVICES=0,1 timeout 60s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py --tokens-list 192'"`,通过。
|
||
- 测试结果和已知问题:
|
||
- Phase 3 build 通过。
|
||
- Phase 3 single-rank 和 2-rank correctness 均通过,FP32 accumulator 最大误差 `0.006541`,满足 `atol=1e-2, rtol=5e-2`。
|
||
- Phase 2 dispatch-only `BLOCK_M=128` single-rank 和 2-rank regression 通过。
|
||
- single-rank 仍出现 PyTorch symmetric memory warning:`cuMulticastCreate failed with: "invalid argument"`,当前不影响 correctness。
|
||
- Phase 3 仅验证 L1 单 tile WGMMA/scaling/debug output,不包含 L1 SwiGLU/quant epilogue、L2 GEMM 或 final combine。
|
||
- 后续待办:
|
||
- Phase 4:在 Phase 3 accumulator 上实现 L1 epilogue,包括 gate/up deinterleave、SwiGLU、top-k weight、FP8 quantize 和 `l2_arrival_mask` release。
|
||
- 扩展 Phase 3/4 测试覆盖更多 N-block、multi-expert 和非首 tile accumulator 映射。
|
||
|
||
## 2026-06-18 01:10:02 CST
|
||
|
||
- 对应 clean 工作 commit:`f3553f976ccd7c2e8978a4c78f7635dca5c14a28`
|
||
- 用户请求摘要:继续 SM90 MegaMoE Phase 4,完成 L1 epilogue 的 SwiGLU、top-k weight、FP8 quant 输出和 `l2_arrival_mask` release,并进行远端 correctness/regression 验证。
|
||
- 本次提交的核心改动:
|
||
- 将 SM90 `l2_acts_sf` 从 per-128 改为 per-64 layout,使每个 L1 `BLOCK_N=128` tile 经过 gate/up SwiGLU 后独立产生 64 列输出 scale。
|
||
- 扩展 SM90 symmetric buffer views,新增 `l2_arrival_mask` Python 可见 view,用于 Phase 4 correctness 检查 L1 epilogue 到 L2 的 tile readiness。
|
||
- 在 SM90 kernel 的 `BLOCK_M == 128` math path 中实现 L1 epilogue:gate/up deinterleave、可选 activation clamp、SwiGLU、top-k weight、per-row amax、E4M3 quantize、`l2_acts`/`l2_acts_sf` global store 和 `l2_arrival_mask` release。
|
||
- 修正 Phase 4 amax reduction 方向为 `math::warp_reduce<4, false>`,按同一 row 的 4 个 lane 做 max reduction。
|
||
- 新增 Phase 4 correctness 脚本 `megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py`,基于实际 dispatch pool 和 transformed weights 计算 PyTorch reference,并验证 `l2_acts`、`l2_acts_sf` 和 `l2_arrival_mask`。
|
||
- 更新 Phase 1 interface smoke 对 `l2_acts_sf` shape 的断言为 `intermediate_hidden // 64`。
|
||
- 关键文件列表:
|
||
- `csrc/apis/sm90_mega.hpp`
|
||
- `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp`
|
||
- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh`
|
||
- `deep_gemm/mega/__init__.py`
|
||
- `megamoe_dev_test_scripts/phase1/interface_smoke.py`
|
||
- `megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py`
|
||
- 详细开发流程:
|
||
- 本地阅读 `MEGAMOE_SM90_DESIGN.md` Phase 4 目标,确认本阶段只实现 L1 epilogue 产物,不实现 L2 GEMM/final combine。
|
||
- 本地扩展 SM90 host buffer layout,将 `l2_acts_sf` 改为 per-64,并暴露 `l2_arrival_mask` view 到 Python `SymmBuffer`。
|
||
- 本地更新 SM90 JIT TMA SF descriptor granularity 为 `kL2ActsGranK = 64`。
|
||
- 本地在 kernel math branch 中添加 L1 epilogue 逻辑,保留 Phase 3 `l1_accum_debug` 写出,新增 `l2_acts`、`l2_acts_sf` 和 `l2_arrival_mask` 写出。
|
||
- 本地新增 `megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py`,固定 `num_tokens=128`、`num_topk=1`,验证两个 L1 N-block 的 64 列 SwiGLU/quant 输出。
|
||
- 本地静态检查:`python3 -m py_compile megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py megamoe_dev_test_scripts/phase1/interface_smoke.py`,通过。
|
||
- 本地静态检查:`git diff --check`,通过。
|
||
- 远端同步:`rsync -avz --relative csrc/apis/sm90_mega.hpp csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh deep_gemm/mega/__init__.py megamoe_dev_test_scripts/phase1/interface_smoke.py megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py g0034:/mnt/beegfs/lxy/DeepGEMM/`。
|
||
- 远端 build:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s bash develop.sh'"`,通过。
|
||
- 远端 Phase 4 single-rank 首次 correctness 暴露 SF mismatch,定位为 row amax reduction 方向错误;本地修复 `warp_reduce<4, true>` 为 `warp_reduce<4, false>` 后重新同步 kernel。
|
||
- 远端 rebuild + Phase 4 single-rank correctness:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s bash develop.sh && timeout 60s python3 megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py'"`,SF 通过,仅剩 2 个 E4M3 rounding-boundary dequant 差异。
|
||
- 本地将 Phase 4 verifier 改为在固定 `atol/rtol` 外允许一档 worst-case E4M3 quantization step,避免 CUDA `expf` 与 PyTorch `silu` 的微小差异跨 FP8 rounding boundary。
|
||
- 远端 Phase 4 single-rank correctness:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s python3 megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py'"`,通过。
|
||
- 远端 Phase 4 2-rank correctness:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && CUDA_VISIBLE_DEVICES=0,1 timeout 60s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py'"`,通过。
|
||
- 远端 Phase 1 single-rank/2-rank regression:`interface_smoke.py`,通过。
|
||
- 远端 Phase 2 single-rank/2-rank regression:`dispatch_only_correctness.py --tokens-list 192`,通过。
|
||
- 远端 Phase 3 single-rank/2-rank regression:`l1_wgmma_single_tile.py`,通过,`max_diff=0.006541`。
|
||
- 提交前本地检查:`python3 -m py_compile ...`、`git diff --check`、`git diff --staged --check`,均通过;检查 touched files 未发现 `// DEBUG` 或 `# DEBUG` debug 代码。
|
||
- 测试结果和已知问题:
|
||
- Phase 4 build 通过。
|
||
- Phase 4 L1 epilogue correctness single-rank 和 2-rank 均通过。
|
||
- Phase 1 interface、Phase 2 dispatch-only、Phase 3 L1 WGMMA regression 在 single-rank 和 2-rank 均通过。
|
||
- single-rank 仍出现 PyTorch symmetric memory warning:`cuMulticastCreate failed with: "invalid argument"`,当前不影响 correctness。
|
||
- Phase 4 当前使用 global store 写 `l2_acts`/`l2_acts_sf`,尚未实现设计中的 STSM + double-buffered TMA store。
|
||
- Phase 4 当前验证单 expert、first pool block、`BLOCK_M=128` L1 epilogue;尚未覆盖 L2 GEMM、final combine、多 expert 完整路径和非首 pool block correctness。
|
||
- 后续待办:
|
||
- Phase 5:实现 L2 GEMM consumer,消费 `l2_arrival_mask` 和 Phase 4 产出的 FP8 intermediate activations。
|
||
- 后续优化:将 Phase 4 global store 替换为 STSM + TMA store,并扩展 correctness 覆盖多 expert、非首 pool block 和更多 shapes。
|
||
|
||
## 2026-06-18 15:17:31 CST
|
||
|
||
- 对应 clean 工作 commit:`9bd05196051dd3f87ebc1eb2ec0c417160ed0418`
|
||
- 用户请求摘要:阅读 Phase 5 计划并执行,实现 SM90 MegaMoE L2 GEMM + scatter 的首个可验证版本。
|
||
- 本次提交的核心改动:
|
||
- 扩展 SM90 symmetric buffer views,新增 `combine_acts: bf16[num_topk, num_max_tokens_per_rank, hidden]`,用于 Phase 5 验证 L2 scatter 输出。
|
||
- SM90 kernel TMA producer 支持 `SM90BlockPhase::Linear2`:等待 `l2_arrival_mask == ((1 << (2 * num_k_blocks)) - 1)` 后加载 `l2_acts` 和 `l2_weights`。
|
||
- SM90 math path 支持 L2 WGMMA:对每个 K=128 block 拆分两个 per-64 half,分别累积 raw FP8 WGMMA,再应用 `l2_acts_sf` 和 `l2_weights_sf` 得到 FP32 L2 output。
|
||
- 新增 BF16 direct scatter:根据 `token_src_metadata` 将 L2 output 写入本地/远端 `combine_acts` 对应 top-k slot 和 token 位置。
|
||
- 新增 Phase 5 correctness 脚本 `megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py`,验证 single-rank 多 expert 下 `combine_acts` 与 PyTorch L2 GEMM BF16 reference 一致。
|
||
- 更新 Phase 1 interface smoke,检查 `combine_acts` shape 和 dtype。
|
||
- 关键文件列表:
|
||
- `csrc/apis/sm90_mega.hpp`
|
||
- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh`
|
||
- `deep_gemm/mega/__init__.py`
|
||
- `megamoe_dev_test_scripts/phase1/interface_smoke.py`
|
||
- `megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py`
|
||
- 详细开发流程:
|
||
- 本地阅读 `MEGAMOE_SM90_DESIGN.md` Phase 5 计划,确认目标为 L2 WGMMA pipeline、per-64 dual-half SF、BF16 epilogue 和 scatter 到 combine buffer。
|
||
- 本地对照 SM100 MegaMoE 的 Linear2 TMA wait/load、L2 BF16 epilogue 和 combine buffer scatter 逻辑,确认 SM90 首版采用 direct global scatter,先不实现 STSM staging。
|
||
- 本地扩展 SM90 host buffer slicing 和 Python `SymmBuffer`,暴露 `combine_acts`。
|
||
- 本地在 SM90 kernel 中重建 `combine_token_buffer` layout,并在 TMA producer 中为 Linear2 增加 `l2_arrival_mask` wait 和 L2 acts/weights TMA load。
|
||
- 本地在 math branch 中保留 Linear1 L1 epilogue,同时新增 Linear2 dual-half WGMMA/scaling/BF16 scatter 分支。
|
||
- 本地新增 `megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py`,固定 `num_tokens=256`、`num_experts=2`、`num_topk=1`,让两个 expert 各接收 128 tokens,覆盖 multi-expert single-rank Phase 5 milestone。
|
||
- 本地静态检查:`python3 -m py_compile megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py megamoe_dev_test_scripts/phase1/interface_smoke.py`,通过。
|
||
- 本地静态检查:`git diff --check`,通过。
|
||
- 远端同步:`rsync -avz --relative csrc/apis/sm90_mega.hpp deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh deep_gemm/mega/__init__.py megamoe_dev_test_scripts/phase1/interface_smoke.py megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py g0034:/mnt/beegfs/lxy/DeepGEMM/`。
|
||
- 远端 build:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s bash develop.sh'"`,通过。
|
||
- 远端 Phase 5 single-rank 首次 correctness 触发 JIT 后失败在测试 reference:脚本直接使用 `expert_recv_count_sum`,未 mask 低 32 位 token count;本地修复为 `count = value & 0xffffffff` 后重新同步测试脚本。
|
||
- 远端 Phase 5 single-rank correctness:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && timeout 60s python3 megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py'"`,通过。
|
||
- 远端 Phase 1 single-rank/2-rank regression:`interface_smoke.py`,通过。
|
||
- 远端 Phase 2 single-rank/2-rank regression:`dispatch_only_correctness.py --tokens-list 192`,通过。
|
||
- 远端 Phase 3 single-rank/2-rank regression:`l1_wgmma_single_tile.py`,通过,`max_diff=0.006541`。
|
||
- 远端 Phase 4 single-rank/2-rank regression:`l1_epilogue_correctness.py`,通过。
|
||
- 远端 Phase 4 8-rank regression:`ssh g0034 "docker exec sglang-lxy bash -lc 'cd /mnt/beegfs/lxy/DeepGEMM && CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 timeout 60s torchrun --standalone --nproc_per_node=8 megamoe_dev_test_scripts/phase4/l1_epilogue_correctness.py'"`,通过。
|
||
- 提交前本地检查:`python3 -m py_compile ...`、`git diff --check`、`git diff --staged --check`,均通过;检查 touched files 未发现 `// DEBUG` 或 `# DEBUG` debug 代码。
|
||
- 测试结果和已知问题:
|
||
- Phase 5 build 通过。
|
||
- Phase 5 single-rank multi-expert L2 GEMM scatter correctness 通过。
|
||
- Phase 1/2/3/4 single-rank 和 2-rank regressions 通过。
|
||
- Phase 4 8-rank regression 通过,说明新增 Linear2 路径未导致现有 multi-rank Phase 4 correctness hang/crash。
|
||
- single-rank 仍出现 PyTorch symmetric memory warning:`cuMulticastCreate failed with: "invalid argument"`,当前不影响 correctness。
|
||
- Phase 5 当前采用 direct global BF16 scatter,没有实现计划中的 BF16 STSM 到 `smem_cd_l2` 再 shared-load scatter。
|
||
- Phase 5 correctness 当前覆盖 single-rank 多 expert;multi-rank L2 scatter 数值 reference、top-k>1、非首 pool block、多 K block/更多 shape 尚未覆盖。
|
||
- Phase 6 combine reduce 和最终 `y` 写回尚未实现,因此端到端 fused output 仍未完成。
|
||
- 后续待办:
|
||
- Phase 6:实现 combine reduce、NVLink barrier、workspace cleanup 和最终 `y` 写回。
|
||
- 扩展 Phase 5 测试到 multi-rank scatter reference、top-k>1、更多 hidden/intermediate shape 和非首 pool block。
|
||
- 后续性能优化:将 direct global scatter 优化为 BF16 STSM/shared-load scatter,降低 store 指令和改善 coalescing。
|
||
|
||
## 2026-06-18 15:54:10 CST
|
||
|
||
- 对应 clean 工作 commit:`71f012a356b7d5d82d15d2354dbf02612a637eba`
|
||
- 用户请求摘要:将 SM90 MegaMoE kernel dispatch 分支缩进修复提交到 git。
|
||
- 本次提交的核心改动:
|
||
- 修正 `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` 中 dispatch 分支 `nvlink_barrier` 之后到 token pull loop 结束的缩进层级。
|
||
- 该 commit 仅包含 whitespace/formatting 变化,不改变代码逻辑。
|
||
- 关键文件列表:
|
||
- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh`
|
||
- 详细开发流程:
|
||
- 本地检查 `git status --short`,确认仅 `sm90_fp8_mega_moe.cuh` 修改。
|
||
- 本地检查 `git diff --check`,通过。
|
||
- 本地检查 `git diff --ignore-all-space -- deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh`,无输出,确认为纯缩进/空白变化。
|
||
- 本地暂存并提交缩进修复:`git commit -m "style: fix sm90 megamoe dispatch indentation"`。
|
||
- 测试结果和已知问题:
|
||
- 未运行远端 build/correctness;本次为 whitespace-only 修复。
|
||
- `git diff --check` 和 staged whitespace-only 检查通过。
|
||
- 后续待办:
|
||
- 继续 Phase 6 combine reduce 和最终 `y` 写回实现。
|
||
|
||
## 2026-06-18 17:41:00 CST
|
||
|
||
- 对应 clean 工作 commit:`453fc7b04654f149d7c5b5cd543e83c74a58e925`
|
||
- 用户请求摘要:完成 SM90 MegaMoE Phase 6 combine reduce 和最终 `y` 写回,并在新提供的 2x H200 远端机器上只跑 rank-2 correctness/reliability 测试后完成收尾提交。
|
||
- 本次提交的核心改动:
|
||
- 在 SM90 kernel 的 `BLOCK_M == 128` 路径中实现 Phase 6 combine reduce:Linear2 scatter 全部完成后执行 system fence、跨 rank NVLink barrier,然后按 token/hidden 直接从 `combine_acts[topk, token, hidden]` 读取 BF16 slot,FP32 累加有效 top-k,写回最终 BF16 `y`。
|
||
- 新增 dispatch+math CTA barrier `kDispatchWithMathBarrierIdx`,确保 dispatch pull 阶段完全结束后 math threads 才进入 combine NVLink barrier,避免多个 phase 复用同一个 NVLink barrier counter 时发生交叠 hang。
|
||
- 将 SM90 L1 epilogue activation clamp 的 compile-time 判断从 `cute::numeric_limits<float>::infinity()` 改为有限阈值 `kActivationClamp < 1.0e30f`,避免 NVRTC 编译时该 `numeric_limits` 表达式不满足 constexpr。
|
||
- SM90 host API 增加 `y` 的 BF16、二维 shape、hidden size 和 contiguous 检查。
|
||
- 新增 Phase 6 端到端 correctness 脚本 `megamoe_dev_test_scripts/phase6/end_to_end_correctness.py`,覆盖 L1 dequant GEMM、SwiGLU、per-64 FP8 requant、L2 dequant GEMM、BF16 per-slot scatter 语义和最终 top-k reduce reference,并默认执行 2 次迭代验证复用可靠性。
|
||
- 关键文件列表:
|
||
- `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh`
|
||
- `csrc/apis/sm90_mega.hpp`
|
||
- `megamoe_dev_test_scripts/phase6/end_to_end_correctness.py`
|
||
- 详细开发流程:
|
||
- 本地阅读 Phase 6 设计和 SM100 combine reduce 参考,确认首版采用 correctness-first direct global load/reduce/store,不实现 TMA combine load/store 优化。
|
||
- 本地在 SM90 kernel 的 math role 中追加 Linear2 scatter 后的 combine 阶段,并在 host API 增加输出 tensor 检查。
|
||
- 本地新增 Phase 6 端到端 PyTorch reference 测试脚本,参数默认 `num_tokens=128`、`hidden=256`、`intermediate_hidden=128`、`num_topk=2`、`iterations=2`,并断言 case 使用 `BLOCK_M=128`。
|
||
- 本地-only 检查:`python3 -B -m py_compile megamoe_dev_test_scripts/phase6/end_to_end_correctness.py`,通过。
|
||
- 本地-only 检查:`git diff --check`,通过。
|
||
- 远端连接:`ssh root@95.133.252.182`,确认项目路径 `/root/DeepGEMM`,机器包含 2 张 NVIDIA H200。
|
||
- 远端默认 `python3` 无 PyTorch,且没有运行中的 Docker 镜像;创建独立 venv:`python3 -m venv /root/dg-venv`。
|
||
- 远端安装依赖:`/root/dg-venv/bin/python -m pip install -U pip packaging wheel setuptools ninja`、`/root/dg-venv/bin/python -m pip install --index-url https://download.pytorch.org/whl/cu128 torch`、`/root/dg-venv/bin/python -m pip install numpy`。
|
||
- 远端环境确认:PyTorch `2.11.0+cu128`、CUDA `12.8`、2 张 H200、`torch.distributed._symmetric_memory` 可 import。
|
||
- 首次只同步 3 个 Phase 6 文件后远端 build 暴露远端 stale workspace:`csrc/python_api.cpp` 仍引用旧 namespace/signature;随后按单向同步规则执行 `rsync -az csrc/ deep_gemm/ megamoe_dev_test_scripts/ root@95.133.252.182:/root/DeepGEMM/...`,使远端代码与本地一致。
|
||
- 远端 build 使用持久化脚本 `/root/phase6_build.sh` 执行 `timeout 180s bash develop.sh > /root/phase6_build.log 2>&1`,最终 `/root/phase6_build.status` 为 `0`,build 通过。
|
||
- 远端 Phase 6 2-rank 首次尝试使用默认 NVCC JIT,`timeout 180s` 和后续 `timeout 600s` 均停在两 rank 并行 NVCC 编译 `sm90_fp8_mega_moe`,未进入 kernel。
|
||
- 改用 `DG_JIT_USE_NVRTC=1` 后暴露 `cute::numeric_limits<float>::infinity()` constexpr 问题;本地改为 `kActivationClamp < 1.0e30f` 并同步 kernel。
|
||
- 修复 constexpr 后,Phase 6 和 Phase 4 2-rank 仍在 cubin 已生成后 timeout,定位为新增 combine NVLink barrier 可能与 dispatch pull NVLink barrier 交叠;本地新增 dispatch+math CTA barrier 后同步 kernel。
|
||
- 远端 Phase 1 2-rank interface smoke:`CUDA_VISIBLE_DEVICES=0,1 timeout 90s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase1/interface_smoke.py`,通过。
|
||
- 远端 Phase 6 2-rank 单迭代验证:`CUDA_VISIBLE_DEVICES=0,1 DG_JIT_USE_NVRTC=1 DG_JIT_CACHE_DIR=/root/dg_jit_cache_nvrtc_phase6_barrier_fix timeout 300s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase6/end_to_end_correctness.py --iterations 1 --activation-clamp 10.0`,通过,`max_diff=0.008789`。
|
||
- 远端 Phase 6 2-rank 两迭代 reliability 验证:`CUDA_VISIBLE_DEVICES=0,1 DG_JIT_USE_NVRTC=1 DG_JIT_CACHE_DIR=/root/dg_jit_cache_nvrtc_phase6_barrier_fix timeout 120s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase6/end_to_end_correctness.py --activation-clamp 10.0`,通过,`iteration=0 max_diff=0.008789`、`iteration=1 max_diff=0.007812`。
|
||
- 提交前本地检查:`python3 -B -m py_compile megamoe_dev_test_scripts/phase6/end_to_end_correctness.py`、`git diff --check`、`git diff --staged --check`,均通过;检查 touched files 未发现临时 `// DEBUG` 或 `# DEBUG` debug 代码。
|
||
- 测试结果和已知问题:
|
||
- 远端 build 通过。
|
||
- Phase 6 2-rank correctness/reliability 通过,覆盖 top-k=2、2 次连续 kernel 调用和最终 `y` reference 对比。
|
||
- 当前远端机器只有 2 张 H200,应用户要求仅执行 rank-2 测试;4-rank/8-rank 尚未在本轮机器上验证。
|
||
- 该远端 CUDA 12.8 的 NVCC JIT 编译 `sm90_fp8_mega_moe` 超过 600 秒;本轮 correctness 使用 `DG_JIT_USE_NVRTC=1` 和有限 `--activation-clamp 10.0` 完成验证。
|
||
- 默认 `activation_clamp=None` 在 NVRTC 路径仍会由 host 生成 `cute::numeric_limits<float>::infinity()` 作为模板实参,可能继续触发 NVRTC constexpr 问题;NVCC 路径或显式 finite clamp 不受该测试限制影响。
|
||
- Phase 6 首版 combine 使用 direct global BF16 load/reduce/store,尚未实现设计中的 chunked TMA combine reduce、workspace cleanup/cleanup NVLink barrier 性能优化路径。
|
||
- 后续待办:
|
||
- 在 4/8 GPU 或多机环境上补跑 Phase 6 4-rank/8-rank correctness/reliability。
|
||
- 处理 NVRTC 下 infinite activation clamp 模板实参生成问题,例如统一调整 host `to_string(float inf)` 策略或 Phase 6 测试默认使用 finite clamp。
|
||
- 后续性能优化:将 direct combine reduce 替换为 chunked/TMA/shared-memory combine,并补充 workspace cleanup 和 cleanup barrier。
|