From 2b8fa5983efa918df429370b444bb89f642978bf Mon Sep 17 00:00:00 2001 From: Oren Amor Date: Wed, 15 Jul 2026 13:16:42 +0000 Subject: [PATCH] libfabric: PT-owns-endpoint with MPSC lock-free ring 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: 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 --- src/plugins/libfabric/libfabric_backend.cpp | 31 +++- src/utils/libfabric/libfabric_rail.cpp | 139 ++++++++++++++++-- src/utils/libfabric/libfabric_rail.h | 106 ++++++++++++- .../libfabric/libfabric_rail_manager.cpp | 34 ++++- src/utils/libfabric/libfabric_rail_manager.h | 3 +- 5 files changed, 285 insertions(+), 28 deletions(-) diff --git a/src/plugins/libfabric/libfabric_backend.cpp b/src/plugins/libfabric/libfabric_backend.cpp index 663b63af8f..23e1aee945 100644 --- a/src/plugins/libfabric/libfabric_backend.cpp +++ b/src/plugins/libfabric/libfabric_backend.cpp @@ -259,16 +259,28 @@ nixlLibfabricCudaCtx::cudaUpdateCtxPtr(void *address, int expected_dev, bool &wa was_updated = false; - if (expected_dev == -1) return -1; - if (myDevId_ != -1 && expected_dev != myDevId_) return -1; + if (expected_dev == -1) { + return -1; + } + if (myDevId_ != -1 && expected_dev != myDevId_) { + return -1; + } ret = cudaQueryAddr(address, is_dev, dev, ctx, pci_bus_id); - if (ret) return ret; - if (!is_dev) return 0; - if (dev != expected_dev) return -1; + if (ret) { + return ret; + } + if (!is_dev) { + return 0; + } + if (dev != expected_dev) { + return -1; + } if (pthrCudaCtx_) { - if (pthrCudaCtx_ != ctx) return -1; + if (pthrCudaCtx_ != ctx) { + return -1; + } return 0; } @@ -282,7 +294,9 @@ nixlLibfabricCudaCtx::cudaUpdateCtxPtr(void *address, int expected_dev, bool &wa int nixlLibfabricCudaCtx::cudaSetCtx() { CUresult result; - if (NULL == pthrCudaCtx_) return 0; + if (NULL == pthrCudaCtx_) { + return 0; + } result = cuCtxSetCurrent(pthrCudaCtx_); return (CUDA_SUCCESS == result); @@ -1183,7 +1197,8 @@ nixlLibfabricEngine::postXferDescriptors(nixlLibfabricReq::OpType op_type, desc_submitted_count, desc_idx, desc_count, - xfer_base_offset); + xfer_base_offset, + local[desc_idx].devId); if (status != NIXL_SUCCESS) { NIXL_ERROR << "prepareAndSubmitTransfer failed for descriptor " << desc_idx diff --git a/src/utils/libfabric/libfabric_rail.cpp b/src/utils/libfabric/libfabric_rail.cpp index 536f25ab78..05fa4a402d 100644 --- a/src/utils/libfabric/libfabric_rail.cpp +++ b/src/utils/libfabric/libfabric_rail.cpp @@ -21,6 +21,10 @@ #include "serdes/serdes.h" #include "libfabric_common.h" +#ifdef HAVE_CUDA +#include +#endif + #include #include #include @@ -168,7 +172,9 @@ ControlRequestPool::~ControlRequestPool() { void ControlRequestPool::cleanup() { - if (buffer_chunks_.empty()) return; // Already cleaned up + if (buffer_chunks_.empty()) { + return; // Already cleaned up + } for (auto &chunk : buffer_chunks_) { if (chunk.mr) { @@ -712,8 +718,8 @@ nixlLibfabricRail::setXferIdCallback(std::function callback) { // Per-rail completion processing - handles one rail's CQ with configurable blocking behavior nixl_status_t -nixlLibfabricRail::progressCompletionQueue() const { - // Completion processing +nixlLibfabricRail::progressCompletionQueue() { + // Completion processing FIRST — process arrivals before posting new items struct fi_cq_data_entry completions[NIXL_LIBFABRIC_CQ_BATCH_SIZE]; int ret; @@ -747,7 +753,8 @@ nixlLibfabricRail::progressCompletionQueue() const { // CQ lock released here - completion is now local data if (ret == -FI_EAGAIN) { - return NIXL_IN_PROG; // No completions available + // SQ full - do inline CQ progress to free slots, then retry + // No completions - fall through to drainPostQueue } if (ret > 0) { @@ -762,10 +769,14 @@ nixlLibfabricRail::progressCompletionQueue() const { } NIXL_DEBUG << "Processed " << ret << " completions on rail " << rail_id; - return NIXL_SUCCESS; } - return NIXL_ERR_BACKEND; // Unexpected case + // PT-owns-endpoint: drain pending posts AFTER processing completions + if (progress_thread_enabled_) { + drainPostQueue(); + } + + return (ret > 0) ? NIXL_SUCCESS : NIXL_IN_PROG; } // Route completion to appropriate handler (rail-specific) @@ -1018,9 +1029,7 @@ nixlLibfabricRail::postRecv(nixlLibfabricReq *req) const { } nixl_status_t -nixlLibfabricRail::postSend(uint64_t immediate_data, - fi_addr_t dest_addr, - nixlLibfabricReq *req) const { +nixlLibfabricRail::postSend(uint64_t immediate_data, fi_addr_t dest_addr, nixlLibfabricReq *req) { if (req->buffer_size == 0 || req->buffer_size > NIXL_LIBFABRIC_SEND_RECV_BUFFER_SIZE) { NIXL_ERROR << "Invalid message size=" << req->buffer_size << " (max: " << NIXL_LIBFABRIC_SEND_RECV_BUFFER_SIZE << ")"; @@ -1063,6 +1072,7 @@ nixlLibfabricRail::postSend(uint64_t immediate_data, } if (ret == -FI_EAGAIN) { + // SQ full - do inline CQ progress to free slots, then retry // Resource temporarily unavailable - retry indefinitely for all providers attempt++; @@ -1099,6 +1109,111 @@ nixlLibfabricRail::postSend(uint64_t immediate_data, return NIXL_ERR_BACKEND; } +/**************************************** + * PT-owns-endpoint: post queue methods + ****************************************/ + +void +nixlLibfabricRail::enqueuePost(const nixlLibfabricPostRequest &req) { + while (!post_ring_.push(req)) { + std::this_thread::yield(); + } +} + +nixl_status_t +nixlLibfabricRail::drainPostQueue() { + nixlLibfabricPostRequest pr; + int posted = 0; + [[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) { + struct iovec iov{}; + iov.iov_base = pr.local_addr; + iov.iov_len = pr.length; + + struct fi_rma_iov rma_iov{}; + rma_iov.addr = pr.remote_addr; + rma_iov.len = pr.length; + rma_iov.key = pr.remote_key; + + void *desc = pr.local_desc; + struct fi_msg_rma msg = {}; + msg.msg_iov = &iov; + msg.desc = &desc; + msg.iov_count = 1; + msg.addr = pr.dest_addr; + msg.rma_iov = &rma_iov; + msg.rma_iov_count = 1; + msg.context = &pr.req->ctx; + msg.data = pr.immediate_data; + + ret = fi_writemsg(endpoint, &msg, pr.fi_flags | FI_REMOTE_CQ_DATA); + } else { + struct iovec iov{}; + iov.iov_base = pr.local_addr; + iov.iov_len = pr.length; + + struct fi_rma_iov rma_iov{}; + rma_iov.addr = pr.remote_addr; + rma_iov.len = pr.length; + rma_iov.key = pr.remote_key; + + void *desc = pr.local_desc; + struct fi_msg_rma msg = {}; + msg.msg_iov = &iov; + msg.desc = &desc; + msg.iov_count = 1; + msg.addr = pr.dest_addr; + msg.rma_iov = &rma_iov; + msg.rma_iov_count = 1; + msg.context = &pr.req->ctx; + + ret = fi_readmsg(endpoint, &msg, pr.fi_flags); + } + + if (ret == -FI_EAGAIN) { + struct fi_cq_data_entry cq_buf[16]; + 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]); + } + } + } + } + + 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; + } + posted++; + if (posted % 32 == 0) { + struct fi_cq_data_entry cq_buf[16]; + 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]); + } + } else if (cq_ret < 0 && cq_ret != -FI_EAGAIN) { + struct fi_cq_err_entry err_entry = {}; + fi_cq_readerr(cq, &err_entry, 0); + NIXL_ERROR << "CQ error in drain interleave on rail " << rail_id << ": " + << fi_strerror(err_entry.err); + } + } + } + + return NIXL_SUCCESS; +} + nixl_status_t nixlLibfabricRail::postWrite(const void *local_buffer, size_t length, @@ -1108,7 +1223,7 @@ nixlLibfabricRail::postWrite(const void *local_buffer, uint64_t remote_addr, uint64_t remote_key, nixlLibfabricReq *req, - uint64_t fi_flags) const { + uint64_t fi_flags) { // Validation if (!req) { NIXL_ERROR << "Invalid request for write on rail " << rail_id; @@ -1160,6 +1275,7 @@ nixlLibfabricRail::postWrite(const void *local_buffer, } if (ret == -FI_EAGAIN) { + // SQ full - do inline CQ progress to free slots, then retry // Resource temporarily unavailable - retry indefinitely for all providers attempt++; @@ -1203,7 +1319,7 @@ nixlLibfabricRail::postRead(void *local_buffer, fi_addr_t dest_addr, uint64_t remote_addr, uint64_t remote_key, - nixlLibfabricReq *req) const { + nixlLibfabricReq *req) { // Validation if (!req) { NIXL_ERROR << "Invalid request for read on rail " << rail_id; @@ -1242,6 +1358,7 @@ nixlLibfabricRail::postRead(void *local_buffer, } if (ret == -FI_EAGAIN) { + // SQ full - do inline CQ progress to free slots, then retry // Resource temporarily unavailable - retry indefinitely for all providers attempt++; diff --git a/src/utils/libfabric/libfabric_rail.h b/src/utils/libfabric/libfabric_rail.h index 3492bdda0b..47b4a17b54 100644 --- a/src/utils/libfabric/libfabric_rail.h +++ b/src/utils/libfabric/libfabric_rail.h @@ -23,6 +23,8 @@ #include #include #include +#include +#include #include #include @@ -232,6 +234,83 @@ operator<<(std::ostream &os, const ConnectionState &state) { } } +/** + * @brief Lock-free Multi-Producer Single-Consumer ring buffer. + * Multiple producer threads (postXfer) push concurrently. + * Single consumer (progress thread) pops and posts. + * Uses CAS on head for multi-producer safety. + */ +template class MpscRing { +public: + bool + push(const T &item) { + 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); + return true; + } + + bool + pop(T &item) { + size_t tail = tail_.load(std::memory_order_relaxed); + if (tail == head_.load(std::memory_order_acquire)) { + return false; + } + // Wait for the slot to be committed (producer might still be writing) + if (!committed_[tail].load(std::memory_order_acquire)) { + return false; + } + item = ring_[tail]; + committed_[tail].store(false, std::memory_order_release); + tail_.store((tail + 1) % Capacity, std::memory_order_release); + return true; + } + + bool + empty() const { + size_t tail = tail_.load(std::memory_order_acquire); + if (tail == head_.load(std::memory_order_acquire)) { + return true; + } + return !committed_[tail].load(std::memory_order_acquire); + } + +private: + alignas(64) std::atomic head_{0}; + alignas(64) std::atomic tail_{0}; + T ring_[Capacity]; + std::atomic committed_[Capacity] = {}; +}; + +/** + * @brief Queued post request for PT-owns-endpoint model + * Main thread enqueues; progress thread dequeues and posts. + */ +struct nixlLibfabricPostRequest { + enum Type { WRITE, READ, SEND } type; + + void *local_addr; + size_t length; + void *local_desc; + uint64_t immediate_data; // WRITE only + fi_addr_t dest_addr; + uint64_t remote_addr; + uint64_t remote_key; + nixlLibfabricReq *req; + uint64_t fi_flags; + int device_id; // CUDA device ordinal for cudaSetDevice in PT +}; + /** Individual libfabric rail managing fabric, domain, endpoint, CQ, and AV */ class nixlLibfabricRail { public: @@ -303,7 +382,7 @@ class nixlLibfabricRail { /** Post send operation with immediate data */ nixl_status_t - postSend(uint64_t immediate_data, fi_addr_t dest_addr, nixlLibfabricReq *req) const; + postSend(uint64_t immediate_data, fi_addr_t dest_addr, nixlLibfabricReq *req); /** Post RDMA write operation with immediate data */ nixl_status_t @@ -315,7 +394,7 @@ class nixlLibfabricRail { uint64_t remote_addr, uint64_t remote_key, nixlLibfabricReq *req, - uint64_t fi_flags = 0) const; + uint64_t fi_flags = 0); /** Post RDMA read operation */ nixl_status_t @@ -325,11 +404,19 @@ class nixlLibfabricRail { fi_addr_t dest_addr, uint64_t remote_addr, uint64_t remote_key, - nixlLibfabricReq *req) const; + nixlLibfabricReq *req); /** Process completion queue with batching support */ nixl_status_t - progressCompletionQueue() const; + progressCompletionQueue(); + + /** Enqueue a post request for the progress thread to execute */ + void + enqueuePost(const nixlLibfabricPostRequest &req); + + /** Drain pending posts (called by progress thread) */ + nixl_status_t + drainPostQueue(); // Callback registration methods /** Set callback for notification message processing */ @@ -345,6 +432,12 @@ class nixlLibfabricRail { void setProgressThreadEnabled(bool enabled); + /** Check if progress thread is enabled for this rail */ + bool + isProgressThreadEnabled() const { + return progress_thread_enabled_; + } + /** Set callback for XFER_ID tracking */ void setXferIdCallback(std::function callback); @@ -380,6 +473,11 @@ class nixlLibfabricRail { // EP mutex to protect endpoint and CQ operations mutable std::mutex ep_mutex_; + // Lock-free MPSC ring: main thread pushes, PT pops and posts + static constexpr size_t POST_RING_CAPACITY = 2048; + using PostRing = MpscRing; + PostRing post_ring_; + // Whether a progress thread is handling CQ draining bool progress_thread_enabled_ = false; diff --git a/src/utils/libfabric/libfabric_rail_manager.cpp b/src/utils/libfabric/libfabric_rail_manager.cpp index 4445b5afa5..55ca5ea352 100644 --- a/src/utils/libfabric/libfabric_rail_manager.cpp +++ b/src/utils/libfabric/libfabric_rail_manager.cpp @@ -395,7 +395,8 @@ nixlLibfabricRailManager::prepareAndSubmitTransfer( size_t &submitted_count_out, int desc_idx, int desc_count, - size_t base_offset) { + size_t base_offset, + int device_id) { // Initialize output parameter submitted_count_out = 0; @@ -448,9 +449,32 @@ nixlLibfabricRailManager::prepareAndSubmitTransfer( req->local_mr = local_mrs[rail_id]; req->remote_key = remote_keys[remote_ep_id]; req->rail_id = rail_id; - // Submit immediately + // Submit: either enqueue for PT, or post directly nixl_status_t status; - if (op_type == nixlLibfabricReq::WRITE) { + if (rails_[rail_id]->isProgressThreadEnabled()) { + // PT-owns-endpoint: enqueue for progress thread to post + uint8_t seq_id = + (op_type == nixlLibfabricReq::WRITE) ? LibfabricUtils::getNextSeqId() : 0; + uint64_t imm_data = (op_type == nixlLibfabricReq::WRITE) ? + NIXL_MAKE_IMM_DATA(NIXL_LIBFABRIC_MSG_TRANSFER, agent_idx, xfer_id, seq_id) : + 0; + nixlLibfabricPostRequest pr{}; + pr.type = (op_type == nixlLibfabricReq::WRITE) ? nixlLibfabricPostRequest::WRITE : + nixlLibfabricPostRequest::READ; + pr.local_addr = req->local_addr; + pr.length = req->chunk_size; + pr.local_desc = fi_mr_desc(req->local_mr); + pr.immediate_data = imm_data; + pr.dest_addr = dest_addrs.at(rail_id)[remote_ep_id]; + pr.remote_addr = req->remote_addr; + pr.remote_key = req->remote_key; + pr.req = req; + pr.fi_flags = fi_flags; + pr.device_id = device_id; + rails_[rail_id]->enqueuePost(pr); + status = NIXL_SUCCESS; + } else if (op_type == nixlLibfabricReq::WRITE) { + // Direct post (PT OFF path) // Generate next SEQ_ID for this specific write operation uint8_t seq_id = LibfabricUtils::getNextSeqId(); uint64_t imm_data = @@ -500,7 +524,9 @@ nixlLibfabricRailManager::prepareAndSubmitTransfer( remote_selected_endpoints[i % remote_selected_endpoints.size()]; NIXL_DEBUG << "rail " << rail_id << ", remote_ep_id=" << remote_ep_id; size_t current_chunk_size = chunk_size + (i == num_rails - 1 ? remainder : 0); - if (current_chunk_size == 0) break; + if (current_chunk_size == 0) { + break; + } // Allocate request nixlLibfabricReq *req = rails_[rail_id]->allocateDataRequest(op_type, xfer_id); if (!req) { diff --git a/src/utils/libfabric/libfabric_rail_manager.h b/src/utils/libfabric/libfabric_rail_manager.h index 10ce041cc7..b5308e52ea 100644 --- a/src/utils/libfabric/libfabric_rail_manager.h +++ b/src/utils/libfabric/libfabric_rail_manager.h @@ -219,7 +219,8 @@ class nixlLibfabricRailManager { size_t &submitted_count_out, int desc_idx, int desc_count, - size_t base_offset); + size_t base_offset, + int device_id = -1); /** Reserve a base offset for a transfer to ensure stable rail assignment * across all descriptors in the transfer. Call once per postXfer. */