libfabric: PT-owns-endpoint with MPSC lock-free ring#1949
Conversation
Move fi_writemsg/fi_readmsg posting into the Progress Thread via a lock-free MPSC ring buffer, eliminating ep_mutex_ contention on the data path. Main thread enqueues to per-rail MpscRing; PT dequeues, posts, and polls completions in a single thread. - MpscRing<T,2048>: CAS multi-producer, single-consumer ring - drainPostQueue: PT posts with FI_MORE batching, interleaves CQ - cudaSetDevice in PT for multi-GPU VRAM buffer validation - Graceful error: completion callback on failure, continue draining
|
👋 Hi oa-aws! Thank you for contributing to ai-dynamo/nixl. Your PR reviewers will review your contribution then trigger the CI to test your changes. 🚀 |
📝 WalkthroughWalkthroughLibfabric adds a lock-free progress-thread post queue, routes device-aware transfers through queued submission when enabled, and drains queued RDMA operations alongside completion-queue processing. CUDA context guards and transfer submission arguments are updated accordingly. ChangesLibfabric post queue flow
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant RailManager
participant nixlLibfabricRail
participant PostRing
participant LibfabricCQ
RailManager->>nixlLibfabricRail: enqueuePost(post request)
nixlLibfabricRail->>PostRing: push post request
nixlLibfabricRail->>PostRing: pop queued posts
nixlLibfabricRail->>LibfabricCQ: fi_writemsg or fi_readmsg
LibfabricCQ-->>nixlLibfabricRail: completion or -FI_EAGAIN
nixlLibfabricRail->>LibfabricCQ: drain and process CQ
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/utils/libfabric/libfabric_rail_manager.h (1)
202-223: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocument the new
device_idcontract.Specify that
device_idis the CUDA device ordinal for VRAM and-1for non-device memory, especially because omission silently selects the default.Proposed documentation
* `@param` base_offset Pre-reserved round-robin offset from reserveBaseOffset() + * `@param` device_id CUDA device ordinal for VRAM, or -1 for non-device memory🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/utils/libfabric/libfabric_rail_manager.h` around lines 202 - 223, Update the parameter documentation for prepareAndSubmitTransfer’s device_id argument to state that it is the CUDA device ordinal for VRAM transfers and must be -1 for non-device memory; also note that omitting it silently selects the default device.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/utils/libfabric/libfabric_rail.cpp`:
- Around line 1189-1195: Update the ret != 0 failure branch in drainPostQueue so
failed posts are not reported through the success-only completion_callback.
Record the asynchronous failure, release pr.req back to its request pool, and
invoke the request’s failure-aware completion notification contract while
preserving the existing error log and continue behavior.
- Around line 1127-1133: Update drainPostQueue() to switch to pr.device_id
before posting each queued VRAM request, reusing the existing CUDA
device-tracking logic and only calling cudaSetDevice when the queued device
differs from the current device. Preserve the existing behavior for non-VRAM
requests and error handling.
- Around line 1178-1185: Update the retry-path CQ draining around fi_cq_read and
the periodic drain to propagate negative CQ read failures and
processCompletionQueueEntry failures instead of ignoring them. Reuse the
existing error-aware CQ drain helper if available; otherwise return the
encountered error from the enclosing operation so CQ faults terminate the loop
rather than causing repeated retries.
In `@src/utils/libfabric/libfabric_rail.h`:
- Around line 299-300: Rename the nested enum type in nixlLibfabricPostRequest
to lower_case naming, and update all references to it accordingly. Also rename
the related using alias at the additional location to lower_case with a _t
suffix, updating every use so the code remains consistent.
- Around line 247-259: Replace the modulo-based head_/tail_ reservation logic in
the ring queue with monotonic sequence numbers, preferably per-slot sequence
values, so producers distinguish wrapped generations and cannot reserve stale
slots after a full cycle. Update the enqueue and corresponding dequeue/commit
checks to use the sequence protocol while preserving bounded capacity and
correct visibility of queued entries.
---
Outside diff comments:
In `@src/utils/libfabric/libfabric_rail_manager.h`:
- Around line 202-223: Update the parameter documentation for
prepareAndSubmitTransfer’s device_id argument to state that it is the CUDA
device ordinal for VRAM transfers and must be -1 for non-device memory; also
note that omitting it silently selects the default device.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 6631e0e2-01fe-40ba-a3f7-58a7a522fed4
📒 Files selected for processing (5)
src/plugins/libfabric/libfabric_backend.cppsrc/utils/libfabric/libfabric_rail.cppsrc/utils/libfabric/libfabric_rail.hsrc/utils/libfabric/libfabric_rail_manager.cppsrc/utils/libfabric/libfabric_rail_manager.h
| [[maybe_unused]] int current_cuda_device = -1; | ||
|
|
||
| while (post_ring_.pop(pr)) { | ||
| int ret = -FI_EAGAIN; | ||
|
|
||
| while (ret == -FI_EAGAIN) { | ||
| if (pr.type == nixlLibfabricPostRequest::WRITE) { |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C3 'device_id|current_cuda_device|cudaSetDevice' \
src/utils/libfabric/libfabric_rail.cpp \
src/utils/libfabric/libfabric_rail_manager.cpp \
src/plugins/libfabric/libfabric_backend.cppRepository: ai-dynamo/nixl
Length of output: 14213
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1124,1215p' src/utils/libfabric/libfabric_rail.cpp
printf '\n====\n'
sed -n '1215,1295p' src/utils/libfabric/libfabric_rail.cpp
printf '\n====\n'
sed -n '460,490p' src/utils/libfabric/libfabric_rail_manager.cpp
printf '\n====\n'
sed -n '1128,1178p' src/plugins/libfabric/libfabric_backend.cppRepository: ai-dynamo/nixl
Length of output: 10441
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C3 'struct nixlLibfabricPostRequest|device_id|current_cuda_device|cudaSetDevice|enqueuePost|drainPostQueue' src/utils/libfabric src/plugins/libfabricRepository: ai-dynamo/nixl
Length of output: 36704
Switch to the queued request’s CUDA device before posting it.
pr.device_id is copied into the queue, but drainPostQueue() never applies it. The direct post path already calls cudaSetDevice() for VRAM transfers, so the progress thread can post on the wrong GPU unless it switches first.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/utils/libfabric/libfabric_rail.cpp` around lines 1127 - 1133, Update
drainPostQueue() to switch to pr.device_id before posting each queued VRAM
request, reusing the existing CUDA device-tracking logic and only calling
cudaSetDevice when the queued device differs from the current device. Preserve
the existing behavior for non-VRAM requests and error handling.
| if (ret != 0) { | ||
| NIXL_ERROR << "drainPostQueue: post failed ret=" << ret << " (" << fi_strerror(-ret) | ||
| << ") on rail " << rail_id; | ||
| if (pr.req && pr.req->completion_callback) { | ||
| pr.req->completion_callback(); | ||
| } | ||
| continue; |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift
Do not report failed posts as successful completions.
The backend callback only increments completed_requests, so invoking it here makes an operation that was never posted appear successful. The request is also never returned to its pool. Record an asynchronous failure, release pr.req, then notify completion through a failure-aware contract.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/utils/libfabric/libfabric_rail.cpp` around lines 1189 - 1195, Update the
ret != 0 failure branch in drainPostQueue so failed posts are not reported
through the success-only completion_callback. Record the asynchronous failure,
release pr.req back to its request pool, and invoke the request’s failure-aware
completion notification contract while preserving the existing error log and
continue behavior.
There was a problem hiding this comment.
I think this is incorrect, otherwise the call to check whether a multi-part/segment request finished will never declare success, as the counters will never match up if a completion is missing, even if failed. In addition to that, maybe a fix should be added to tell the total request status, in case some sub-segments posting failed.
| size_t head, next; | ||
| do { | ||
| head = head_.load(std::memory_order_relaxed); | ||
| next = (head + 1) % Capacity; | ||
| if (next == tail_.load(std::memory_order_acquire)) { | ||
| return false; | ||
| } | ||
| } while (!head_.compare_exchange_weak( | ||
| head, next, std::memory_order_acq_rel, std::memory_order_relaxed)); | ||
| // head slot is now reserved for us | ||
| ring_[head] = item; | ||
| // Signal that this slot is written (use a separate committed array) | ||
| committed_[head].store(true, std::memory_order_release); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate queue usages and existing concurrency/stress coverage.
rg -n -C4 '\bMpscRing\b|\bpost_ring_\b' src test 2>/dev/nullRepository: ai-dynamo/nixl
Length of output: 2846
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Outline =="
ast-grep outline src/utils/libfabric/libfabric_rail.h --view expanded || true
echo
echo "== Relevant lines around MpscRing =="
sed -n '230,330p' src/utils/libfabric/libfabric_rail.h | cat -n
echo
echo "== Post-ring usage =="
sed -n '1110,1145p' src/utils/libfabric/libfabric_rail.cpp | cat -n
echo
echo "== Search for additional ring/queue logic =="
rg -n -C3 'head_|tail_|committed_|MpscRing|PostRing|enqueuePost|dequeuePost|progress thread' src/utils/libfabric/libfabric_rail.h src/utils/libfabric/libfabric_rail.cppRepository: ai-dynamo/nixl
Length of output: 15669
Use monotonic sequence numbers for the ring indices.
head_/tail_ are tracked only modulo Capacity, so a stalled producer can CAS a wrapped head_ value after the ring cycles and corrupt the queue state. That can make head_ == tail_ while entries are still queued, hiding data and breaking the buffer. Use per-slot sequence numbers or a proven bounded MPSC queue here.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/utils/libfabric/libfabric_rail.h` around lines 247 - 259, Replace the
modulo-based head_/tail_ reservation logic in the ring queue with monotonic
sequence numbers, preferably per-slot sequence values, so producers distinguish
wrapped generations and cannot reserve stale slots after a full cycle. Update
the enqueue and corresponding dequeue/commit checks to use the sequence protocol
while preserving bounded capacity and correct visibility of queued entries.
| struct nixlLibfabricPostRequest { | ||
| enum Type { WRITE, READ, SEND } type; |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Apply the enforced enum and type-alias naming rules.
Proposed naming fix
struct nixlLibfabricPostRequest {
- enum Type { WRITE, READ, SEND } type;
+ enum post_type { WRITE, READ, SEND } type;
...
- using PostRing = MpscRing<nixlLibfabricPostRequest, POST_RING_CAPACITY>;
- PostRing post_ring_;
+ using post_ring_t = MpscRing<nixlLibfabricPostRequest, POST_RING_CAPACITY>;
+ post_ring_t post_ring_;Based on learnings, enum types use lower_case, while using aliases use lower_case with a _t suffix.
Also applies to: 478-478
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/utils/libfabric/libfabric_rail.h` around lines 299 - 300, Rename the
nested enum type in nixlLibfabricPostRequest to lower_case naming, and update
all references to it accordingly. Also rename the related using alias at the
additional location to lower_case with a _t suffix, updating every use so the
code remains consistent.
Source: Learnings
fengjica
left a comment
fengjica
left a comment
There was a problem hiding this comment.
Reviewed the change except for MPSC ring implementation.
The major question is still that when one PT drains all rings (of multiple rails), the unbounded while loop in draining can block it from progressing other rails.
| int desc_count, | ||
| size_t base_offset); | ||
| size_t base_offset, | ||
| int device_id = -1); |
| // Determine striping strategy | ||
| bool use_striping = shouldUseStriping(transfer_size) && selected_rails.size() > 1; | ||
| NIXL_DEBUG << "use_striping=" << use_striping; | ||
| if (!use_striping) { |
There was a problem hiding this comment.
This change now is for non-striping case. Is there a reason that we should apply them different for striping case below?
There was a problem hiding this comment.
Yes, we can add that too for non-striping. Actually I think the performance issue should be more acute in this case. Fixing.
| mutable std::mutex ep_mutex_; | ||
|
|
||
| // Lock-free MPSC ring: main thread pushes, PT pops and posts | ||
| static constexpr size_t POST_RING_CAPACITY = 2048; |
There was a problem hiding this comment.
Is this big enough to hold the buffer if the draining thread (1 PT today across multiple rails) cannot catch up with the speed of the posting thread? And there can be multiple posting threads (e.g. https://github.com/ai-dynamo/nixl/blob/main/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp#L1685).
Some of our workloads uses very large batch size, which translates to large post request sizes. Maybe worth running one perf test, or just nixlbench with a big batch size.
There was a problem hiding this comment.
I agree. I actually don't like it being compile time constant. Maybe it is better to have it configurable with some large default according to perf tests?
| #include "libfabric_common.h" | ||
|
|
||
| #ifdef HAVE_CUDA | ||
| #include <cuda_runtime.h> |
There was a problem hiding this comment.
Is this included here because some new cuda runtimes functions are added?
There was a problem hiding this comment.
I just took it form the original PR. I think it was in preparation for setting cuda context per deferred request in the progress thread (whereas now this is done in nixlLibfabricEngine::postXferDescriptors before calling rail_manager_.prepareAndSubmitTransfer, which assumes request execution in the current thread). This change is also required (i.e. when PT=1 do not set CUDA context in the engine, but rather let the PT do that).
| void | ||
| nixlLibfabricRail::enqueuePost(const nixlLibfabricPostRequest &req) { | ||
| while (!post_ring_.push(req)) { | ||
| std::this_thread::yield(); |
There was a problem hiding this comment.
Is it better to yield and retry, or abort to let let users know that something is full (and let user retry)?
There was a problem hiding this comment.
I changed the implementation due to code rabbit comments, but the question remains the same: what should we do in case of a full ring buffer? I think the best option for now is to retry a few times until some request timeout is reached and let user know we failed, so the ring buffer size can be configured to larger value for the workload in question.
| nixlLibfabricRail::drainPostQueue() { | ||
| nixlLibfabricPostRequest pr; | ||
| int posted = 0; | ||
| [[maybe_unused]] int current_cuda_device = -1; |
There was a problem hiding this comment.
It should be used, see comments above.
| int cq_ret = fi_cq_read(cq, cq_buf, 16); | ||
| if (cq_ret > 0) { | ||
| for (int c = 0; c < cq_ret; c++) { | ||
| processCompletionQueueEntry(&cq_buf[c]); |
There was a problem hiding this comment.
Could we defer processing them so that post can be drained faster? Just a thought, not necessary for this PR.
There was a problem hiding this comment.
This eventually increments an atomic integer. Unless there is other heavy computation, I don't think we can gain much. The question is, is there a use case where a completion callback would be heavy (i.e. someone developing on top of the libfabric plugin)? If so, then this should have a design, because I don't think it is trivial (e.g. does someone rely on counters to catch up to tell request is finished? If so, would deferring cause unwanted delays? Maybe user should declare whether the callback should be deferred? etc.).
There was a problem hiding this comment.
You are right.
Upon checking the callbacks today in the code, I also notice a bug. Because I was thinking a notif needs to be sent out after all completions are done, for READ. I only found one place in the code in checkXfer() for that. It means, if the user (initiator) does not poll status through checkXfer() into our code, we do not send the notif to the target. https://github.com/ai-dynamo/nixl/blob/main/src/plugins/libfabric/libfabric_backend.cpp#L1452-L1454. I was expecting when PT=1, this notif should be sent to target earlier. If you also think this is the expected behavior, I'll open a bug.
No need to address it in this PR, too.
| continue; | ||
| } | ||
| posted++; | ||
| if (posted % 32 == 0) { |
There was a problem hiding this comment.
I guess this 32 is something that can be tuned later to find an optimal ratio of draining vs reading cq, is that the case?
There was a problem hiding this comment.
I think so. This is coming from the original PR, and I was wondering when to benchmark fine-tuning of it. We should also (as you mentioned above) enforce some batch-size limit so that post requests from other rails are executed in a timely manner. This budgeting is not trivial though, so I suggest for this PR we put some constants that work reasonably, and for next PRs we improve on that based on perf test results.
| int posted = 0; | ||
| [[maybe_unused]] int current_cuda_device = -1; | ||
|
|
||
| while (post_ring_.pop(pr)) { |
There was a problem hiding this comment.
This while loop will not break until the post_ring is empty. With one PT, it's possible that it's stuck here, and do not get a chance to progress other rails, right?
There was a problem hiding this comment.
Correct. I was worried about that too. This is how it was done in the original PR, and we need to think about a solution. See previous comment. Currently adding some compile time constant as a batch limit.
2 participants
libfabric: Move fi_writemsg/fi_readmsg posting into the Progress Thread (when progress thread is enabled) via a lock-free MPSC ring buffer, eliminating ep_mutex_ contention on the data path, resulting in performance improvement with larger batch sizes. Main thread enqueues to per-rail MpscRing; PT dequeues, posts, and polls completions in a single thread.
Summary by CodeRabbit
Performance
Reliability