Files
DeepGEMM/MEGAMOE_SM90_DEV.md
2026-06-18 17:42:09 +08:00

284 lines
32 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# 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/SM100SM90 使用 `(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 个 viewsSM100 分支保持 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 consumerTMA 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 branchesPhase 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 epiloguegate/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 多 expertmulti-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 reduceLinear2 scatter 全部完成后执行 system fence、跨 rank NVLink barrier然后按 token/hidden 直接从 `combine_acts[topk, token, hidden]` 读取 BF16 slotFP32 累加有效 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。