diff --git a/MEGAMOE_SM90_DEV.md b/MEGAMOE_SM90_DEV.md index 88376e1..98a2ce0 100644 --- a/MEGAMOE_SM90_DEV.md +++ b/MEGAMOE_SM90_DEV.md @@ -75,3 +75,46 @@ - 后续待办: - 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 映射。