Fix the CPU-assisted (GDRCopy) IBGDA path in the legacy kernels - #677
Fix the CPU-assisted (GDRCopy) IBGDA path in the legacy kernels#677tommy85 wants to merge 2 commits into
Conversation
docs/nvshmem.md documents the CPU-assisted IBGDA mode (NVSHMEM_IBGDA_NIC_HANDLER=cpu, backed by GDRCopy) as the supported alternative when modifying the driver regkeys is not an option, but the async post-send branch of ibgda_submit_requests never worked: 1. It CASed on qp->tx_wq.prod_idx, which NVSHMEM wires to a device-side management counter the CPU proxy never reads (the proxy polls a separate producer-slot array), so submissions were invisible and the first dispatch hung waiting for a doorbell that was never rung. 2. It skipped the ready_head CAS queue that NVSHMEM's own submit path orders on. DeepEP and NVSHMEM share the QP management memory, so after a DeepEP workload ready_head was left stale and the first NVSHMEM-side submission on a shared QP - the nvshmem_barrier_all inside Buffer::destroy - spun forever on a CAS that can never succeed (that spin has no timeout). Order every submission through ready_head and, in the async mode, publish the new producer index with a system-scope atomicMax on qp->tx_wq.bf, which NVSHMEM wires to the proxy-polled slot in the CPU handler modes - mirroring NVSHMEM's own ibgda_proxy_post_send. The default GPU-handler path is unchanged. Validated with NVSHMEM 3.4.5 and NVSHMEM_IBGDA_NIC_HANDLER=cpu: the full tests/test_low_latency.py run including teardown now exits 0 (previously: hang at the first dispatch; with only the pointer fixed, deadlock in Buffer::destroy).
| (unsigned long long int)new_wqe_idx); | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
🟡 warning: In async mode nvshmemi_ibgda_quiet reads qp->tx_wq.prod_idx (wired to &mvars->tx_wq.prod_idx) to decide how many completions to wait for, but the patched submit path (ibgda_submit_requests, :156-177) no longer writes that counter: it now CASes ready_head and publishes bf, and the only other writer of mvars->tx_wq.prod_idx is ibgda_post_send (:135), which runs only in the GPU path. Before the patch the async CAS targeted qp->tx_wq.prod_idx and so advanced it. Result: in async mode mvars->tx_wq.prod_idx is never updated, ibgda_poll_cq() polls against a stale index and returns immediately, and the QP is never actually quieted. That breaks the documented 'wait for all previous inflight wrs to complete' drain that internode.cu relies on before rewriting cleared rdma buffers (internode.cu:142-149 and 197-199), so under NVSHMEM_IBGDA_NIC_HANDLER=cpu those kernels can rewrite still-in-flight buffers and reuse SQ slots before the NIC drains them. The low_latency path the author tested does not call DeepEP's quiet (it uses NVSHMEM barriers + receiver-side flag polling), which is why validation passed.
🤖 v4
There was a problem hiding this comment.
Fixed in 5c4f2b8. nvshmemi_ibgda_quiet now reads mvars->tx_wq.ready_head unconditionally: with the submit path always ordering through ready_head, it is the submitted frontier in both modes, and it is exactly what the GPU mode already polled against. You are right that the low-latency test drains via receiver-side flags rather than this quiet, which is why our validation did not catch it — thanks for the catch.
| (unsigned long long int)new_wqe_idx); | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
🟡 warning: Potential correctness gap outside the touched lines but directly coupled to this change. nvshmemi_ibgda_quiet still reads qp->tx_wq.prod_idx when state->use_async_postsend is true: uint64_t prod_idx = state->use_async_postsend ? ld_na_relaxed(qp->tx_wq.prod_idx) : ld_na_relaxed(&qp->mvars.tx_wq.ready_head);. After this patch, DeepEP publishes the producer index to qp->tx_wq.bf and no longer to prod_idx. Unless the NVSHMEM CPU proxy writes the completed/producer index back into the location qp->tx_wq.prod_idx points to, this counter can remain stale (e.g. 0), causing ibgda_poll_cq to return immediately (*cq->cons_idx >= idx) without actually waiting for completions. This path is exercised by the normal internode kernels (internode.cu lines ~143/199/1350), not just the low-latency test that was validated. Please confirm whether the proxy updates the prod_idx target in cpu / cpu_host_memory modes, or whether quiet should instead poll based on ready_head (which is now always current). If it is safe (proxy writes back prod_idx), a short comment here explaining why the quiet read target differs from the publish target would prevent future confusion.
🤖 v3
There was a problem hiding this comment.
Confirmed: the NVSHMEM CPU proxy only reads the producer slots (nvshmemt_ibgda_progress copies them out and rings doorbells); it never writes back to the location qp->tx_wq.prod_idx points to. So polling ready_head — always current after this patch — is the correct fix rather than relying on a write-back. Done in 5c4f2b8 with a comment explaining why the read target is ready_head and not the publish target.
| // mode (the UAR is only mapped for the GPU handler), so a system-scope monotonic | ||
| // max mirrors NVSHMEM's own `ibgda_proxy_post_send`. The old code CASed on | ||
| // `tx_wq.prod_idx`, which NVSHMEM wires to a device-side counter the proxy never | ||
| // reads, so submissions were invisible to the proxy and the first dispatch hung |
There was a problem hiding this comment.
🔵 suggestion: The async publish uses atomicMax_system((unsigned long long int*)qp->tx_wq.bf, (unsigned long long int)new_wqe_idx). qp->tx_wq.bf is the same pointer written as a 64-bit doorbell record in ibgda_ring_db (via reinterpret_cast<uint64_t*>(qp->tx_wq.bf)), so its 8-byte alignment for a 64-bit atomic is implied, but consider adding a brief assertion or comment noting the required alignment/width, since a mis-sized atomic on bf would silently corrupt the proxy slot. Also confirm the intended width matches NVSHMEM's proxy-polled slot type; if NVSHMEM defines that slot as a narrower type in some builds, the 64-bit atomicMax_system could overwrite adjacent state.
🤖 v3
There was a problem hiding this comment.
Comment added in 5c4f2b8 documenting the width/alignment contract: NVSHMEM allocates the proxy-polled slot array as dedicated, naturally aligned uint64_t entries (ibgda_create_qp_shared_objects allocates sizeof(uint64_t) * num_eps and nvshmemt_ibgda_progress reads the slots back as uint64_t), so the 64-bit atomicMax_system matches the slot width in all NVSHMEM builds; there is no narrower-slot configuration.
🤖 ds-review-bot Code Reviewv4The patch fixes the CPU-assisted (GDRCopy) IBGDA submit path in csrc/kernels/legacy/ibgda_device.cuh. It (a) orders every submission through mvars->tx_wq.ready_head (fixing the NVSHMEM-side barrier deadlock in Buffer::destroy, Bug 2), and (b) in async mode publishes the new producer index to qp->tx_wq.bf with a system-scope atomicMax (fixing the proxy-visibility hang, Bug 1). The publish target (bf, the proxy-polled slot in the CPU modes), the use of atomicMax_system for host visibility, and the switch to an always-ready_head ordering are all correct, and the GPU-direct path is untouched. The default tests/test_low_latency.py path the author validated should pass. The main residual issue is that the async submit path no longer advances mvars->tx_wq.prod_idx while nvshmemi_ibgda_quiet still reads it, which makes the IBGDA QP-drain a no-op in async mode for the normal (internode.cu) IBGDA kernels that the validation did not cover. v3The change fixes the CPU-assisted (GDRCopy) IBGDA path in csrc/kernels/legacy/ibgda_device.cuh by (1) always ordering submissions through mvars->tx_wq.ready_head instead of conditionally CASing on qp->tx_wq.prod_idx in the async branch, and (2) in async (use_async_postsend) mode publishing the new producer index with a system-scope atomicMax on qp->tx_wq.bf, mirroring NVSHMEM's ibgda_proxy_post_send. Both fixes match the described defects: submissions are now visible to the CPU proxy (Bug 1) and shared-QP teardown no longer deadlocks because ready_head is kept consistent (Bug 2). The default GPU-handler path (!use_async_postsend) is functionally unchanged and still calls ibgda_post_send as before. The diff is minimal, correctly scoped to a single file, and the reasoning is well-documented in comments; the change is protocol logic only and should be NIC-agnostic. Hoisting the ready_head CAS out of the async/non-async conditional is a good, clearly-commented improvement that correctly resolves the teardown deadlock on shared QPs. Two follow-up concerns remain: nvshmemi_ibgda_quiet still reads prod_idx in async mode (no longer updated by DeepEP's publish path), and the width/alignment of the atomicMax_system on tx_wq.bf should be confirmed against NVSHMEM's proxy slot type. Files reviewed: 1 📍 未定位到 diff 的评论🔵 suggestion |
…ange Review follow-up: nvshmemi_ibgda_quiet still read qp->tx_wq.prod_idx in the async mode, but the fixed submit path no longer advances that counter (the async publish goes to tx_wq.bf), which would have turned the QP drain into a no-op under the CPU handlers - exactly where the normal internode kernels rely on it before rewriting cleared rdma buffers. ready_head is the submitted frontier in both modes now that ibgda_submit_requests always orders through it, so read it unconditionally; this also matches what the GPU mode already did. Also document the width/alignment contract of the atomicMax_system on tx_wq.bf (the proxy slot is a dedicated naturally-aligned uint64_t).
|
Thanks @ds-review-bot — the main warning was a real gap. Addressed in 5c4f2b8:
|
|
|
2 participants
docs/nvshmem.md (§ "Install GDRCopy and load the gdrdrv kernel module") documents the CPU-assisted IBGDA mode as the supported alternative "when modifying the driver regkeys is not an option". In practice that mode (
NVSHMEM_IBGDA_NIC_HANDLER=cpu) has never worked with DeepEP: the async post-send branch inibgda_submit_requests(csrc/kernels/legacy/ibgda_device.cuh) has two independent defects, either of which is fatal.Bug 1 — submissions are invisible to the CPU proxy (first dispatch hangs).
The async branch CASes on
qp->tx_wq.prod_idx. NVSHMEM's host code wires that pointer to a device-side management counter (ibgda_setup_gpu_state:rc_h[i].tx_wq.prod_idx = base_mvars_d_addr + prod_idx_offset), but the proxy thread polls a separate producer-slot array (nvshmemt_ibgda_progressreadsqp_shared_object.prod_idx_mobject). The value DeepEP publishes is therefore never observed, no doorbell is ever rung, and the very first dispatch spins forever waiting for data (surfacing as the NCCL watchdog SIGABRT after 600 s).Bug 2 — teardown deadlock on shared QPs.
The async branch also skips the
ready_headCAS queue. NVSHMEM's own submit path orders every submission throughmvars->tx_wq.ready_head, and DeepEP shares the same QP management memory with NVSHMEM. After a DeepEP workload,resv_headis far ahead whileready_headis still 0; the first NVSHMEM-side submission on a shared QP — thenvshmem_barrier_all()insideBuffer::destroy()— reserves a slot and then spins forever onatomicCAS(ready_head==0, expect resv_head). That spin has no timeout, so the process hangs incuStreamSynchronizewith the GPU at 100%.Fix. Always order submissions through
ready_head(matching NVSHMEM's protocol), and in the async mode publish the new producer index with a system-scopeatomicMaxonqp->tx_wq.bf— in the CPU-assisted modes NVSHMEM wiresbfto the proxy-polled producer slot (seeibgda_get_device_qp), so this mirrors NVSHMEM's ownibgda_proxy_post_sendexactly and needs no NVSHMEM-side change. System scope keeps the publish visible to the CPU observer also forNVSHMEM_IBGDA_NIC_HANDLER=cpu_host_memory, where the slot lives in host memory.Validation. On our rig (NVSHMEM 3.4.5, 2 nodes, 4 RC QPs per rank,
NVSHMEM_IBGDA_NIC_HANDLER=cpuwith GDRCopy)tests/test_low_latency.pypreviously hung at the first dispatch; with only the pointer rewired it then deadlocked inBuffer::destroy(). With this patch the full run — benchmark plus teardown (barrier → free → finalize) — completes and exits 0, with benchmark numbers unchanged. The defect is pure protocol logic, independent of the NIC; on ConnectX it should reproduce by simply exportingNVSHMEM_IBGDA_NIC_HANDLER=cpu(NVSHMEM built with GDRCopy) and runningtests/test_low_latency.py.The GPU-direct doorbell path (
NVSHMEM_IBGDA_NIC_HANDLER=gpu, the default) is unaffected: its submit branch already ordered throughready_headand does not take the async path.