[Bugfix] Reset MegaMoE combine readiness for CUDA graph replay

[foreverrookie]

CUDA graph replay can reuse %gridid, allowing MegaMoE to accept stale combine-readiness tags and read old or partially updated results.

Reset local peer-readiness flags before the existing pull barrier in both FP8/FP4 and BF16 kernels. The workspace layout and existing barriers are unchanged.

Validation:

• FP8: vLLM adapter tests on 8 B300 ranks, including input staging, unequal/empty source ranks, and changed-input graph replay: 3 tests passed per rank.
• BF16: standalone tests on 2 B300 GPUs, with balanced, unequal, and empty-source cases. Eager execution passed before and after; graph replay failed before the fix, and all 12 changed-input replays per case match eager output exactly after it.

AI assistance was used for the fix and validation.

[ds-review-bot]

🔵 suggestion: combine readiness 的推送、清零、notify、等待逻辑在 bf16 与 fp8_fp4 两个 impl 中逐行重复，本次修复恰好只落在其中一个文件，说明重复代码已造成一致性问题。可选的后续改进：将 “push peer grid idx + reset local readiness” 与 “notify/wait” 抽成 comm/ 或 layout/ 下的小工具函数（例如 MegaMoEWorkspace 上的成员），两个 kernel 共用，避免后续修复再次只覆盖一处。

_{🤖 v5}

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

未发现本次变更新引入的缺陷。FP8/FP4 路径在现有 pull barrier 前重置本地 readiness，能够隔离旧标记与当前调用的完成通知，且未改变工作区布局。当前环境缺少 CUDA 工具和 GPU，未运行多卡回归测试。

v5

本次修复在 sm100_fp8_fp4_mega_moe.cuh 的 dispatch 阶段（SM 0、pull barrier 之前）将本 rank 的 combine_ready_grid_idx[peer] 清零，并同步更新了 layout/mega_moe.cuh 中的注释。修复思路正确，已验证：(1) 清零位于 kBeforeDispatchPullBarrierTag NVLink barrier 之前，peer 本轮的 st_rel_sys notify 只能在其通过该 barrier 之后（依赖本 rank 已完成清零）发出，不会被清零覆盖；(2) 上一轮的 notify 由 kDispatchWithEpilogueBarrierIdx 同步排在 kAfterWorkspaceCleanBarrierTag 之前，在下一次 launch 开始前已全部落地，不会晚于本轮清零到达；(3) 本 rank epilogue 的 wait 位于 epilogue grid sync 与 kDispatchWithEpilogueBarrierIdx 之后，严格晚于 SM 0 dispatch 的清零；(4) 重置值 0 与 grid_idx + 1 (>= 1) 不冲突，workspace 布局（MegaMoESignals 偏移 128 的 static_assert）与现有 barrier 均未改动；(5) layout/mega_moe.cuh 注释准确描述了 FP8/FP4 路径的新语义。但修复不完整：sm100_bf16_mega_moe.cuh 使用完全相同的 peer_grid_idx 推送 / st_rel_sys notify / ld_acq_sys == grid_idx 等待协议，却没有加入同样的清零，在 CUDA graph replay 复用 %gridid 时仍会读到过期或部分更新的 combine 结果。

v4

The FP8/FP4 part of this change is correct. Clearing each rank's local combine_ready_grid_idx[*] entries before the dispatch pull barrier is the right fix for CUDA-graph replay, because a replayed graph can reuse %gridid and the readiness protocol (combine_ready_grid_idx[peer] == own grid index) would otherwise match tags left over from the previous invocation. The placement is sound: the reset is performed by the same SM 0 threads that release the NVLink barrier signal, so it is ordered before peers' current-invocation completion notifications and cannot erase them, while stale tags are removed. The workspace layout and barrier set are unchanged and the extra stores are negligible. One gap remains: the BF16 MegaMoE kernel implements the exact same readiness protocol but was not updated, so it is still vulnerable to the same replay bug.

Files reviewed: 2
Issues found: 🔴 1 critical | 🟡 1 warning | 🔵 1 suggestion
Inline comments posted: 1
General comments (无法定位到 diff): 2

---

📍 未定位到 diff 的评论

🔴 critical deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh:L349: sm100_bf16_mega_moe.cuh 第 349-353 行仍是修复前的代码：仅通过 sym_buffer.map(workspace.get_peer_grid_idx_ptr(...), thread_idx) 推送 grid_idx + 1，没有对本地 workspace.get_combine_ready_grid_idx_ptr(thread_idx) 清零。其 epilogue（第 1186-1220 行）与 FP8/FP4 版本一样，用 ptx::ld_acq_sys(peer_ready_ptr) == grid_idx 判断 peer 的 L2 写入完成。CUDA graph replay 复用 %gridid 时，上一轮遗留的 tag 会被误判为本轮完成，导致 combine 读取过期/部分更新的数据——正是本 MR 声称要修复的问题。MR 标题与描述均为通用的 “MegaMoE”，且 layout/mega_moe.cuh 的新注释宣称 “Dispatch resets readiness before the pull barrier”，但对 BF16 kernel 而言该陈述并不成立。

建议在 sm100_bf16_mega_moe.cuh 的 if (sm_idx == 0) 块中做与 FP8/FP4 完全一致的修改：

if (thread_idx < kNumRanks) {
    ptx::st_rel_sys(workspace.get_combine_ready_grid_idx_ptr(thread_idx), uint64_t(0));
    *sym_buffer.map(workspace.get_peer_grid_idx_ptr(sym_buffer.rank_idx), thread_idx) = ptx::get_grid_idx() + 1;
}

并同步更新注释；同时将 BF16 路径也纳入 CUDA graph replay 回归测试（同一 graph、更新输入后连续 replay）。 _{🤖 v5}

🟡 warning deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh:L352: This BF16 MegaMoE kernel has the identical CUDA-graph-replay stale-readiness bug that the MR fixes for FP8/FP4, but it was not included in the fix. Like the FP8/FP4 kernel, dispatch pushes ptx::get_grid_idx() + 1 into every peer's peer_grid_idx (this line), and combine waits for ptx::ld_acq_sys(combine_ready_grid_idx[peer]) == grid_idx (line 1219), but the local combine_ready_grid_idx entries are never cleared before the pull barrier. On graph replay the reused %gridid will therefore match the previous invocation's readiness tags and the BF16 combine can read stale or partially written results. This path is reachable (deep_gemm.bf16_mega_moe is exposed and tests/test_mega_moe.py --mma-type bf16xbf16 exercises it). Please apply the same reset here, i.e. ptx::st_rel_sys(workspace.get_combine_ready_grid_idx_ptr(thread_idx), uint64_t(0)); before pushing the peer grid index, so both MegaMoE variants are covered. _{🤖 v4}

[foreverrookie]

Added the same BF16 reset in 7e924a8. On 2 B300 GPUs, changed-input graph replay failed before the fix in balanced, unequal, and empty-source cases; after the fix, all 12 replays per case match eager output exactly.