docs: record sm90 megamoe phase4 dev log
This commit is contained in:
@@ -118,3 +118,51 @@
|
||||
- 后续待办:
|
||||
- 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。
|
||||
|
||||
Reference in New Issue
Block a user