diff --git a/MEGAMOE_SM90_DEV.md b/MEGAMOE_SM90_DEV.md index bbcf811..d11e240 100644 --- a/MEGAMOE_SM90_DEV.md +++ b/MEGAMOE_SM90_DEV.md @@ -236,3 +236,48 @@ - `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::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::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::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。