Fix LL dispatch RC QP assert for all-P2P topologies

[xiayan0118]

Summary

Low-latency dispatch currently asserts num_rc_per_pe >= num_local_experts unconditionally on SM 0. With NVSHMEM 3.6, optional IBGDA may leave num_rc_per_pe == 0 even when every remote rank is directly reachable through P2P. That poisons single-host / all-local jobs (e.g. B200) even though every send/count path already falls back to the P2P pointer when nvshmemi_get_p2p_ptr(...) != 0.

Gate the RC QP requirement on per-peer P2P reachability: require RC QPs only if at least one non-self peer has no P2P pointer.

Related reports: #134, #488 (same num_rc_per_pe family; this PR specifically covers the LL dispatch all-P2P case).

Change

In csrc/kernels/legacy/internode_ll.cu dispatch:

• Scan non-self peers with nvshmemi_get_p2p_ptr(rdma_recv_count, rank, dst_rank)
• Assert num_rc_per_pe >= num_local_experts only when needs_ibgda is true

Test plan

• Single-node / all-local LL dispatch with IBGDA unavailable (num_rc_per_pe == 0) succeeds when all peers are P2P-reachable
• Multi-node / non-P2P topology still fails the assert when RC QPs are insufficient
• Existing LL dispatch correctness on NVLink all-local ranks

Made with Cursor

[ds-review-bot]

🟡 warning: The gate is dispatch-only. The sibling barrier/clean_low_latency_buffer and combine kernels in the same file also depend on RC QPs for non-P2P peers (e.g. via ibgda_quiet/nvshmemi_ibgda_put_nbi_warp/nvshmemi_ibgda_amo_nonfetch_add) without analogous per-peer P2P gating. With num_rc_per_pe == 0, a non-P2P/mixed run passing this assert (because the dispatcher only requires full RC capacity, same as before) can still crash in those kernels. Confirm this MR is intentionally shipping alongside the referenced PRs #134/#488 and won't be merged alone relying on num_rc_per_pe == 0.

_{🤖 v4}

[ds-review-bot]

🟡 warning: The per-peer probe/self-exclusion relies on subtle invariants that are only implicit: self (dst_rank == rank) is excluded because nvshmemi_get_p2p_ptr maps self to the local pointer (always nonzero), and a single representative heap address rdma_recv_count is used to probe reachability. A zero peer_heap_base_p2p means no P2P (RC mandatory), a nonzero one means NVLink P2P. If any local symmetric buffer used by the real send/count/recv paths is ever not in the same heap (or the self-mapping changes), the gate would be wrong. Consider documenting those invariants or deriving the flag from the actual buffers/state.

_{🤖 v4}

[ds-review-bot]

🔵 suggestion: The P2P scan loop is executed on the single-thread path (warp_id == num_warps - 1, sm_id == 0) but not restricted to a single lane, so every lane of that warp redundantly walks the full num_ranks loop and calls nvshmemi_get_p2p_ptr per peer. This is functionally correct (the assert is uniform across the warp) but does O(num_ranks) redundant work per lane. Consider guarding the scan with lane_id == 0 and broadcasting needs_ibgda, or documenting that the redundancy is intentional and negligible.

_{🤖 v3}

[ds-review-bot]

🔵 suggestion: The reachability scan probes only rdma_recv_count. If any other buffer used on the IBGDA path (e.g. rdma_recv_x / rdma_x send targets) could have different P2P reachability than rdma_recv_count, this single-buffer probe might not capture all IBGDA needs. In practice these buffers share the same NVSHMEM symmetric heap and peer topology, so a single representative probe is sufficient — consider adding a brief comment making that assumption explicit for future readers.

_{🤖 v3}

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

The conditional QP assertion accurately matches the existing P2P fallback paths without introducing an observable regression.

v4

The LL dispatch asserts num_rc_per_pe >= num_local_experts unconditionally on SM 0. With NVSHMEM 3.6 the optional IBGDA transport may leave num_rc_per_pe == 0 even when every remote rank is P2P/NVLink-reachable, which poisons single-node/all-local (e.g. B200) jobs while every send/count/barrier path already falls back to the P2P pointer. This MR gates the assert on per-peer P2P reachability: it requires RC QPs only if at least one non-self peer has no P2P pointer (nvshmemi_get_p2p_ptr == 0). The gating is correct and sufficient for the dispatch all-P2P case: self/self-traffic maps to the local pointer (always nonzero) and single-rank/self never needs RC QPs; the assert runs on the original SM 0/last-warp location, which is representative because per-peer P2P-reachability is a rank-global flag; the preserved threshold num_rc_per_pe >= num_local_experts ignores the num_devices_initialized factor used in the QP-index selection (id % (num_rc_per_pe * num_devices_initialized)), but that is pre-existing/conservative and out of scope. One scope risk: the sibling barrier/clean_low_latency_buffer and combine kernels in the same file use the same RC QPs for non-P2P peers without analogous gating, so a non-P2P/multi-node run with num_rc_per_pe == 0 would still fail there (the description refers this to the related PRs #134/#488).

v3

The change correctly gates the RC QP requirement in the legacy low-latency dispatch kernel on per-peer P2P reachability. On SM 0, it now scans all non-self peers with nvshmemi_get_p2p_ptr(rdma_recv_count, rank, dst_rank) and only asserts num_rc_per_pe >= num_local_experts when at least one non-self peer is not directly reachable through P2P (needs_ibgda). This resolves the false-positive assert on all-P2P / single-host jobs (e.g. B200) where NVSHMEM 3.6 zeroes num_rc_per_pe because optional IBGDA is unavailable, while every send/count path already falls back to the P2P pointer. Multi-node / non-P2P topologies still trigger the assert when RC QPs are insufficient. The implementation is consistent with existing nvshmemi_get_p2p_ptr usages (same (uint64_t ptr, rank, dst_rank) signature at lines 44, 263, 346), and all referenced variables (rank, num_ranks, num_local_experts, rdma_recv_count) are in scope. The scope is limited to the single file described in the PR.

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