Fix hybrid combine top-k weight gradients with single reduction

[Gstdioh]

With hybrid combine, do_expand=False and allow_multiple_reduction=False, the sender forwards a separate top-k weight row for each destination GPU. The reduce epilogue currently selects the master slot by scale-out rank, so routes on different GPUs in the same scale-out domain can read the wrong row.

For example, receiver-local score-gradient rows [11, 0] and [0, 22] should produce [11, 22], but selecting one domain-wide row produces [0, 22]. This affects combine(..., topk_weights=...), including its use for dispatch backward. A plain dispatch/combine weight roundtrip can hide the bug because the dispatched weights are duplicated.

This change selects the score source master slot by global GPU rank for this mode, matching the forwarded row layout. The existing selection is preserved for other modes, with a compile-time assertion for slot layout. The original header formatting is preserved; the production change is limited to this one block (+9/-1 lines).

A standalone distributed regression uses synthetic receiver-local gradients and covers same-GPU, same-domain/different-GPU, and cross-domain routes, including invalid top-k slots. It supports --allow-multiple-reduction=1 as a control.

Validation: Python syntax, Ruff 0.6.5, YAPF 0.40.2, and git diff --check pass. The header already differs from clang-format on the base branch; whole-file formatting is intentionally excluded to keep this fix focused. Real multi-node GPU validation completed for commit 655d77e6bfcdcf46b065a4596104f6f2d4338503: two nodes with 8 NVIDIA H200 GPUs each (16 ranks), with fresh processes and separate JIT caches for --allow-multiple-reduction=0 and =1. All three cases passed for all ranks in both modes, and both node launchers exited with code 0 for each run. No test or production-code changes were needed. This run used PyTorch 2.9.0 (CUDA 12.4) and nvcc 12.8.61; it does not validate the README's PyTorch >=2.10 baseline or add distributed CI coverage. See the GPU validation results.

Run on each node (replace the rendezvous placeholders):

torchrun --nnodes=2 --nproc-per-node=2 --node-rank=NODE_RANK \
    --master-addr=MASTER_ADDR --master-port=29500 \
    tests/elastic/test_combine_topk_weights.py

The test asserts at least two logical scale-out domains and at least two GPUs per domain. Repeat with --allow-multiple-reduction=1 to exercise the control path.

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

未发现需要修复的问题。新的权重来源选择与发送端按 GPU 转发的布局一致，其他模式保持原有行为。Python 语法检查和 diff 检查通过；当前环境缺少 PyTorch，未运行多节点 CUDA 回归测试。

v5

The MR fixes a real correctness bug in the hybrid combine reduce epilogue: for do_expand=False + allow_multiple_reduction=False, the sender (hybrid_combine.cuh) writes one top-k weight row per destination GPU into slot master_topk_idx and the forwarder deduplicates by scale-up rank and forwards slot k to scaleout_recv_buffer.get_rank_buffer(k) without merging rows, so the receiver must select the master slot by global GPU rank (stored_dst_expert_idx / kNumExpertsPerRank) rather than by scale-out rank. I cross-checked this against the sender path and the selection now matches the forwarded row layout exactly; the old scale-out-rank match would pick one slot per domain and read the wrong row for other GPUs in that domain. The change is gated with if constexpr so all other modes keep the existing behaviour, and the EP_STATIC_ASSERT(not kUseRankLayout) is sound because use_rank_layout() returns false unconditionally when multiple reduction is disabled. The new key also matches the deduplicate_key already used for activation reduction in this mode, making weights and activations consistent. The standalone distributed regression is well designed: receiver-local gradients ([0,11,0,22] masked to local experts) mean each GPU contributes a distinct partial row so duplicated dispatched weights cannot mask the bug; it covers same-GPU / same-domain-different-GPU / cross-domain routes, invalid -1 slots (exercising masked lanes in ptx::match), a MIN all-reduce of pass flags, and an --allow-multiple-reduction=1 control. API usage matches deep_ep/buffers/elastic.py and deep_ep.topk_idx_t. git diff --check a56d615 HEAD is clean, the working tree is clean, and the clang-format reformatting is whitespace-only. Remaining risk: the distributed test has not been executed on real multi-node hardware and is not wired into any automated suite. Overall the change is small, targeted, and correct; approve with minor non-blocking suggestions.

v4

This MR fixes the hybrid combine path used for top-k weight gradients when do_expand=False and allow_multiple_reduction=False. In that mode the sender forwards a separate top-k weight row per destination GPU, but the reduce epilogue selected the source slot's master lane using the scale-out rank, so multiple GPUs inside the same scale-out domain could read the same (wrong) row. The fix computes the weight-source rank as the global GPU rank (stored_dst_expert_idx / kNumExpertsPerRank) for this mode, keeps the previous selection for all other modes, and adds a compile-time assertion that the per-GPU layout is the non-rank layout. A standalone multi-node regression (tests/elastic/test_combine_topk_weights.py) covers same-GPU, same-scale-out/different-GPU, cross-scale-out, and invalid top-k slots, with --allow-multiple-reduction=1 as a control. Verification notes: (1) the source selection is correct because stored_dst_expert_idx / kNumExpertsPerRank equals the global GPU rank and, with kUseRankLayout false, compute_topk_slots uses master lane indices as buffer slots, so the same key selects the forwarded per-GPU row; invalid lanes (>= kNumTopk) keep stored_dst_expert_idx = -1 and never dereference the pointer. (2) The test's local_routes mask is correct because the dispatch copy epilogue stores rank-local expert indices (dst_expert_idx - expert_start_idx), so local experts are always in [0, experts_per_rank). (3) The synthetic receiver-local gradient is a good design choice, since a plain weight roundtrip would be masked by duplicated dispatched weights. No blocking issues found; the only actionable items are the three suggestions below.

Files reviewed: 2
Issues found: 🟡 1 warning | 🔵 6 suggestion
Inline comments posted: 7