fix(legacy): make IBGDA CQ polling concurrency-safe - #702
Conversation
|
|
||
| // A 16-bit WQE counter is ambiguous if `idx` is more than one epoch ahead. | ||
| // Wait until the target has been submitted before interpreting the counter. | ||
| while (ld_na_relaxed(cq->prod_idx) < idx) |
There was a problem hiding this comment.
🔵 suggestion: The producer guard while (ld_na_relaxed(cq->prod_idx) < idx) is an unbounded busy-wait whose correctness depends entirely on cq->prod_idx aliasing the monotonic posted frontier that ibgda_post_send/ibgda_submit_requests update (via mvars->tx_wq.prod_idx in the async path and ready_head otherwise). This wiring lives in the NVSHMEM headers, not in this repo, so the invariant is implicit. Since the PR's own integration note flags this as a hard dependency (especially when stacking with async post-send / #677), please add a brief comment here stating that cq->prod_idx must point at the WQ posted frontier, so a future async-post-send refactor does not silently break the wrap disambiguation. Also consider a debug-only assert that the requested idx is reachable, since if a caller ever passes a stale/over-large target this loop hangs silently.
🤖 v3
| // Reconstruct the full software consumer index from the 16-bit hardware | ||
| // counter and keep the shared consumer monotonic across concurrent pollers. | ||
| ++wqe_counter; | ||
| const uint64_t new_cons_idx = ((idx & ~0xffffULL) | wqe_counter) + ((static_cast<uint16_t>(idx) > wqe_counter) ? 0x10000ULL : 0); |
There was a problem hiding this comment.
🔵 suggestion: The full-index reconstruction is correct but subtle and relies on the newly added NVSHMEMI_IBGDA_MAX_QP_DEPTH <= 32768 bound plus the producer guard to keep wqe_counter+1 within one 16-bit epoch of idx. Please expand the comment to explicitly tie these together (e.g. "safe because idx has been produced (prod_idx >= idx) and depth <= 32768, so the completed counter is within one epoch of idx's low 16 bits"). This documents why the 0x10000 carry correction cannot mis-select an epoch and protects the invariant against future edits to the depth assert or the guard.
🤖 v3
🤖 ds-review-bot Code Reviewv6The concurrency-safe polling logic matches NVSHMEM's established CQ polling algorithm, including wraparound handling and monotonic consumer updates. No actionable regressions were identified. v4v3The change makes Files reviewed: 1 |
2 participants
Summary
Make the legacy IBGDA CQ poller safe when multiple callers observe the same
collapsed CQ, and make its 16-bit completion-counter handling wrap-safe.
The change mirrors the core protocol used by NVSHMEM:
interpreting the 16-bit CQE counter;
wait;
atomicMax;Why
The existing helper explicitly retained an exclusive-poller assumption.
Concurrent callers can otherwise continue spinning after another caller has
already advanced the shared consumer, or write an older target over a newer
consumer index. A target sufficiently far ahead can also make a stale
16-bit WQE counter ambiguous across wraparound.
This extends the initial early-return added by #371 with the producer,
concurrent-consumer, wrap reconstruction, and monotonic update checks from
the NVSHMEM implementation.
Scope
This PR only changes CQ polling in
csrc/kernels/legacy/ibgda_device.cuh. It intentionally does not change:Validation
sm_103with CUDA 13.0.88,NVSHMEM 3.7.2, NCCL 2.30.7, and PyTorch 2.11;
bash ./format.sh;git diff --check upstream/main...HEAD;32 GPUs, with the final correctness check at round 18,921 and without the
prior CQ/queue timeout.
The multi-node validation binary also carried the required NVSHMEM 3.7.2 RC
QP-layout compatibility backport, so that result is integration evidence
rather than a current-main single-variable comparison.
Integration note
This patch relies on
cq->prod_idxremaining the monotonic posted frontier.Any async post-send changes, including when stacking with #677, must preserve
that invariant; the producer guard is required to disambiguate 16-bit
counter wraparound.