From 593e16f7e3af283b2a4073ecab2bca33393828c1 Mon Sep 17 00:00:00 2001 From: Yueyang Pan Date: Tue, 16 Jun 2026 13:30:09 +0200 Subject: [PATCH 01/13] [ep]: add cxi transport support - add a libfabric CXI transport backend and wire it into the EP proxy runtime - add generic build switches for USE_LIBFABRIC_CXI and NUM_MAX_NVL_PEERS - relax internode benchmark assumptions so validation works on non-8-GPU nodes --- ep/Makefile | 30 +- ep/bench/buffer.py | 2 +- ep/bench/test_internode.py | 21 +- ep/include/cxi_transport.hpp | 76 +++++ ep/include/ep_configs.cuh | 2 + ep/include/proxy.hpp | 13 +- ep/include/rdma.hpp | 9 + ep/include/transport.hpp | 42 +++ ep/include/uccl_proxy.hpp | 2 +- ep/setup.py | 19 ++ ep/src/cxi_transport.cpp | 486 ++++++++++++++++++++++++++++++++ ep/src/internode.cu | 39 +-- ep/src/proxy.cpp | 528 +++++++++++++++++++++++++++++++++-- ep/src/uccl_ep.cc | 13 +- ep/src/uccl_proxy.cpp | 63 +++-- 15 files changed, 1268 insertions(+), 77 deletions(-) create mode 100644 ep/include/cxi_transport.hpp create mode 100644 ep/include/transport.hpp create mode 100644 ep/src/cxi_transport.cpp diff --git a/ep/Makefile b/ep/Makefile index 45b6753b4..d7854dedf 100644 --- a/ep/Makefile +++ b/ep/Makefile @@ -23,6 +23,21 @@ else DMABUF_CFLAGS := endif +# Enable libfabric/CXI transport skeleton with USE_LIBFABRIC_CXI=1. +USE_LIBFABRIC_CXI ?= 0 +ifeq ($(USE_LIBFABRIC_CXI),1) + $(info libfabric CXI transport enabled) + CXI_CFLAGS := -DUSE_LIBFABRIC_CXI + CXI_LDFLAGS := -lfabric + ifdef LIBFABRIC_HOME + CXI_CFLAGS += -I$(LIBFABRIC_HOME)/include + CXI_LDFLAGS += -L$(LIBFABRIC_HOME)/lib + endif +else + CXI_CFLAGS := + CXI_LDFLAGS := +endif + ARCH := $(shell uname -m) GPU_NAME ?= $(shell nvidia-smi --query-gpu=name --format=csv,noheader | head -n1) CPU_IS_ARM64 := 0 @@ -94,13 +109,18 @@ ifeq ($(PER_EXPERT_BATCHING),1) override NVCCFLAGS += -DPER_EXPERT_BATCHING endif +ifdef NUM_MAX_NVL_PEERS + override CXXFLAGS += -DNUM_MAX_NVL_PEERS=$(NUM_MAX_NVL_PEERS) + override NVCCFLAGS += -DNUM_MAX_NVL_PEERS=$(NUM_MAX_NVL_PEERS) +endif + -CXXFLAGS += $(EFA_CFLAGS) $(GH_CFLAGS) $(DMABUF_CFLAGS) $(NORMAL_CFLAGS) -NVCCFLAGS += $(EFA_CFLAGS) $(GH_CFLAGS) $(DMABUF_CFLAGS) $(NORMAL_CFLAGS) -LDFLAGS += $(EFA_LDFLAGS) -INCLUDES += $(EFA_CFLAGS) $(GH_CFLAGS) $(DMABUF_CFLAGS) $(NORMAL_CFLAGS) +CXXFLAGS += $(EFA_CFLAGS) $(GH_CFLAGS) $(DMABUF_CFLAGS) $(CXI_CFLAGS) $(NORMAL_CFLAGS) +NVCCFLAGS += $(EFA_CFLAGS) $(GH_CFLAGS) $(DMABUF_CFLAGS) $(CXI_CFLAGS) $(NORMAL_CFLAGS) +LDFLAGS += $(EFA_LDFLAGS) $(CXI_LDFLAGS) +INCLUDES += $(EFA_CFLAGS) $(GH_CFLAGS) $(DMABUF_CFLAGS) $(CXI_CFLAGS) $(NORMAL_CFLAGS) -SRC_CPP := src/proxy.cpp src/rdma.cpp src/common.cpp src/uccl_proxy.cpp src/uccl_bench.cpp src/fifo.cpp +SRC_CPP := src/proxy.cpp src/rdma.cpp src/common.cpp src/uccl_proxy.cpp src/uccl_bench.cpp src/fifo.cpp src/cxi_transport.cpp SRC_CU := src/bench_kernel.cu src/internode_ll.cu src/internode.cu src/layout.cu src/intranode.cu src/ep_runtime.cu # Non-binding .cc translation units linked into the shared library. SRC_CC_LIB := src/adaptive_sleeper.cc diff --git a/ep/bench/buffer.py b/ep/bench/buffer.py index 0aa7d3e1f..b73f6fb8e 100644 --- a/ep/bench/buffer.py +++ b/ep/bench/buffer.py @@ -116,7 +116,7 @@ def __init__( # cudaMalloc/cudaMallocHost-style allocation instead of a possibly # segmented caching-allocator mapping. scratch_dlpack, rdma_buffer_is_host_allocated = ep.get_rdma_buffer( - num_rdma_bytes, device_index + num_rdma_bytes, device_index, bool(low_latency_mode) ) self.scratch = torch.utils.dlpack.from_dlpack(scratch_dlpack) else: diff --git a/ep/bench/test_internode.py b/ep/bench/test_internode.py index 4808f2bb8..c918d0913 100644 --- a/ep/bench/test_internode.py +++ b/ep/bench/test_internode.py @@ -6,13 +6,13 @@ make clean && make -j install On first node: -torchrun --nnodes=2 --nproc_per_node=8 --node_rank=0 \ +torchrun --nnodes=2 --nproc_per_node= --node_rank=0 \ --master_addr=10.1.227.34 --master_port=12355 \ bench/test_internode.py --num-tokens=4096 \ --hidden=7168 --num-topk=8 --num-experts=256 --test-ll-compatibility On second node: -torchrun --nnodes=2 --nproc_per_node=8 --node_rank=1 \ +torchrun --nnodes=2 --nproc_per_node= --node_rank=1 \ --master_addr=10.1.227.34 --master_port=12355 \ bench/test_internode.py --num-tokens=4096 \ --hidden=7168 --num-topk=8 --num-experts=256 --test-ll-compatibility @@ -97,7 +97,7 @@ def test_main( args.num_experts, ) - assert num_experts % num_ranks == 0 and num_local_ranks == 8 + assert num_experts % num_ranks == 0 if local_rank == 0: print( f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk_groups={num_topk_groups}, num_topk={num_topk}", @@ -567,7 +567,7 @@ def test_loop( explicitly_destroy=True, ) - assert num_local_ranks == 8 and num_ranks > 8 + assert num_nodes > 1 and num_ranks > num_local_ranks for seed in range(int(1e9)): if local_rank == 0: @@ -627,8 +627,8 @@ def test_loop( parser.add_argument( "--num-processes", type=int, - default=8, - help="Number of processes to spawn (default: 8)", + default=None, + help="Number of local processes. Defaults to LOCAL_WORLD_SIZE from torchrun.", ) parser.add_argument( "--num-tokens", type=int, default=4096, help="Number of tokens (default: 4096)" @@ -668,9 +668,12 @@ def test_loop( if args.num_topk_groups is None: args.num_topk_groups = min(num_nodes, 4) - num_processes = args.num_processes - if num_processes != 8: - raise ValueError("Only --num-processes=8 is supported for this test.") + num_processes = args.num_processes or local_world_size + if num_processes != local_world_size: + raise ValueError( + f"--num-processes must match LOCAL_WORLD_SIZE={local_world_size}, " + f"got {num_processes}" + ) # NOTE: modified from deep_ep local_rank = int(os.environ["LOCAL_RANK"]) num_local_ranks = int(os.environ.get("LOCAL_WORLD_SIZE", 1)) diff --git a/ep/include/cxi_transport.hpp b/ep/include/cxi_transport.hpp new file mode 100644 index 000000000..e6fcc3975 --- /dev/null +++ b/ep/include/cxi_transport.hpp @@ -0,0 +1,76 @@ +#pragma once + +#include "transport.hpp" + +#ifdef USE_LIBFABRIC_CXI +#include +#include +#include +#include +#include +#endif + +#include +#include +#include + +class CxiTransport final : public EpTransport { + public: + CxiTransport() = default; + ~CxiTransport() override; + + void init(ProxyCtx& ctx) override; + void register_main_buffer(void* ptr, size_t len, int cuda_device) override; + void register_atomic_buffer(void* ptr, size_t len, int cuda_device) override; + LocalConnInfo local_info() const override; + void connect_peer(int peer, RemoteConnInfo const& remote) override; + + void post_write(int dst_rank, uint64_t wr_id, uint64_t local_offset, + uint64_t remote_offset, uint32_t bytes, + bool low_latency) override; + + void post_atomic_add(int dst_rank, uint64_t wr_id, + uint64_t remote_atomic_offset, int64_t value, + bool fence) override; + void post_barrier_atomic_add(int dst_rank, uint64_t wr_id, size_t slot, + uint64_t value); + uint64_t load_barrier_word(size_t slot) const; + + int poll(TransportCompletion* out, int max) override; + void destroy() override; + + private: + [[noreturn]] static void unavailable(); + +#ifdef USE_LIBFABRIC_CXI + ProxyCtx* ctx_ = nullptr; + void* main_buffer_ = nullptr; + void* atomic_buffer_ = nullptr; + size_t main_buffer_len_ = 0; + size_t atomic_buffer_len_ = 0; + fid_fabric* fabric_ = nullptr; + fid_domain* domain_ = nullptr; + fid_ep* ep_ = nullptr; + fid_cq* cq_ = nullptr; + fid_av* av_ = nullptr; + fid_mr* main_mr_ = nullptr; + fid_mr* atomic_mr_ = nullptr; + fid_mr* atomic_operand_mr_ = nullptr; + fid_mr* barrier_mr_ = nullptr; + LocalConnInfo local_info_{}; + std::vector peer_addrs_; + std::vector peer_infos_; + std::vector atomic_operands_; + std::vector atomic_operand_used_; + std::vector barrier_words_; + + struct OpContext { + fi_context context = {}; + uint64_t wr_id = 0; + bool is_write = false; + bool is_atomic = false; + size_t atomic_operand_slot = static_cast(-1); + }; + std::list op_contexts_; +#endif +}; diff --git a/ep/include/ep_configs.cuh b/ep/include/ep_configs.cuh index e1618329a..7c1d9f92f 100644 --- a/ep/include/ep_configs.cuh +++ b/ep/include/ep_configs.cuh @@ -1,6 +1,8 @@ #pragma once +#ifndef NUM_MAX_NVL_PEERS #define NUM_MAX_NVL_PEERS 8 +#endif #define NUM_MAX_RDMA_PEERS 20 #define NUM_WORKSPACE_BYTES (32 * 1024 * 1024) #define NUM_MAX_LOCAL_EXPERTS 1024 diff --git a/ep/include/proxy.hpp b/ep/include/proxy.hpp index 8f978c9ad..2f75aeb02 100644 --- a/ep/include/proxy.hpp +++ b/ep/include/proxy.hpp @@ -2,6 +2,7 @@ #define PROXY_HPP #include "common.hpp" +#include "cxi_transport.hpp" #include "proxy_ctx.hpp" #include "rdma.hpp" #include "ring_buffer.cuh" @@ -9,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -104,11 +106,17 @@ class Proxy { void init_common(); void init_sender(); void init_remote(); + bool use_cxi_transport() const; void notify_gpu_completion(uint64_t& my_tail); void post_gpu_command(uint64_t& my_tail, size_t& seen); void post_gpu_commands_mixed(std::vector const& wrs_to_post, std::vector const& cmds_to_post); + void post_cxi_commands(std::vector const& wrs_to_post, + std::vector const& cmds_to_post); + void poll_cxi_completions(); + CxiTransport* cxi_transport_for_rank(int rank) const; + uint64_t load_cxi_barrier_word_sum(size_t slot) const; void post_barrier_msg(int dst_rank, bool ack, uint64_t seq); void send_barrier(uint64_t wr); void barrier_check(); @@ -141,7 +149,10 @@ class Proxy { std::vector> ctxs_for_all_ranks_; std::vector local_infos_, remote_infos_; std::vector ctx_by_tag_; - void* atomic_buffer_ptr_; + void* atomic_buffer_ptr_ = nullptr; + CxiTransport* cxi_transport_ = nullptr; + std::vector> cxi_transports_by_rank_; + size_t cxi_outstanding_ops_ = 0; std::vector postponed_atomics_; std::vector postponed_wr_ids_; diff --git a/ep/include/rdma.hpp b/ep/include/rdma.hpp index cea471c76..704520b6e 100644 --- a/ep/include/rdma.hpp +++ b/ep/include/rdma.hpp @@ -44,6 +44,15 @@ struct RDMAConnectionInfo { uint32_t data_qp_num[kChannelPerProxy]; // #endif +#ifdef USE_LIBFABRIC_CXI + static constexpr uint32_t kMaxCxiEndpointName = 256; + uint32_t cxi_ep_name_len = 0; + uint8_t cxi_ep_name[kMaxCxiEndpointName] = {}; + uint64_t cxi_main_mr_key = 0; + uint64_t cxi_atomic_mr_key = 0; + uint64_t cxi_barrier_mr_key = 0; +#endif + #ifdef USE_DMABUF // Chunked MR info — exchanged when the GPU buffer is split across // multiple MRs (with IOMMU DMA-BUF 2 GiB limit). num_mr_chunks == 0 means diff --git a/ep/include/transport.hpp b/ep/include/transport.hpp new file mode 100644 index 000000000..5598a7b14 --- /dev/null +++ b/ep/include/transport.hpp @@ -0,0 +1,42 @@ +#pragma once + +#include "rdma.hpp" +#include +#include + +enum class TransportKind { Verbs, Cxi }; + +using LocalConnInfo = RDMAConnectionInfo; +using RemoteConnInfo = RDMAConnectionInfo; + +struct TransportCompletion { + uint64_t wr_id = 0; + uint64_t data = 0; + bool is_write = false; + bool is_atomic = false; + bool is_recv_data = false; + int status = 0; +}; + +class EpTransport { + public: + virtual ~EpTransport() = default; + + virtual void init(ProxyCtx& ctx) = 0; + virtual void register_main_buffer(void* ptr, size_t len, int cuda_device) = 0; + virtual void register_atomic_buffer(void* ptr, size_t len, + int cuda_device) = 0; + virtual LocalConnInfo local_info() const = 0; + virtual void connect_peer(int peer, RemoteConnInfo const& remote) = 0; + + virtual void post_write(int dst_rank, uint64_t wr_id, uint64_t local_offset, + uint64_t remote_offset, uint32_t bytes, + bool low_latency) = 0; + + virtual void post_atomic_add(int dst_rank, uint64_t wr_id, + uint64_t remote_atomic_offset, int64_t value, + bool fence) = 0; + + virtual int poll(TransportCompletion* out, int max) = 0; + virtual void destroy() = 0; +}; diff --git a/ep/include/uccl_proxy.hpp b/ep/include/uccl_proxy.hpp index 5b33e855c..2dcd12c3a 100644 --- a/ep/include/uccl_proxy.hpp +++ b/ep/include/uccl_proxy.hpp @@ -98,7 +98,7 @@ class UcclProxy { int local_rank_; int device_index_; int nic_local_rank_; - void* atomic_buffer_ptr_; + void* atomic_buffer_ptr_ = nullptr; bool atomic_buffer_is_host_allocated_ = false; // true => cudaFreeHost, false => cudaFree int node_idx_; diff --git a/ep/setup.py b/ep/setup.py index 59f7e61ea..28a74057c 100644 --- a/ep/setup.py +++ b/ep/setup.py @@ -163,6 +163,24 @@ def run(self): nvcc_dlink = [] extra_link_args = [] use_dmabuf = False + use_libfabric_cxi = int(os.getenv("USE_LIBFABRIC_CXI", "0")) + num_max_nvl_peers = os.getenv("NUM_MAX_NVL_PEERS") + if num_max_nvl_peers: + nvl_peers_flag = f"-DNUM_MAX_NVL_PEERS={int(num_max_nvl_peers)}" + print(f"Building with NUM_MAX_NVL_PEERS={int(num_max_nvl_peers)}") + cxx_flags.append(nvl_peers_flag) + nvcc_flags.append(nvl_peers_flag) + + if use_libfabric_cxi: + print("Building with libfabric CXI transport support (USE_LIBFABRIC_CXI)") + cxx_flags.append("-DUSE_LIBFABRIC_CXI") + nvcc_flags.append("-DUSE_LIBFABRIC_CXI") + libraries.append("fabric") + + libfabric_home = os.getenv("LIBFABRIC_HOME") + if libfabric_home: + include_dirs.append(Path(libfabric_home) / "include") + library_dirs.append(Path(libfabric_home) / "lib") if torch.version.cuda: # Add CUDA library directory to library_dirs @@ -428,6 +446,7 @@ def run(self): print(f" > EFA Support: {'Yes' if has_efa else 'No'}") print(f" > GH200 Support: {'Yes' if has_gh200 else 'No'}") print(f" > DMA-BUF Support: {'Yes' if use_dmabuf else 'No'}") + print(f" > libfabric CXI Support: {'Yes' if use_libfabric_cxi else 'No'}") print(f" > Device Arch: {device_arch}") print(f" > Sources: {len(sources)} files") print(f" > Headers (tracked): {len(header_files)} files") diff --git a/ep/src/cxi_transport.cpp b/ep/src/cxi_transport.cpp new file mode 100644 index 000000000..79cdce55c --- /dev/null +++ b/ep/src/cxi_transport.cpp @@ -0,0 +1,486 @@ +#include "cxi_transport.hpp" +#include "util/gpu_rt.h" +#include +#include +#include +#include +#include +#include + +CxiTransport::~CxiTransport() { destroy(); } + +void CxiTransport::unavailable() { + throw std::runtime_error( + "CxiTransport requires building with USE_LIBFABRIC_CXI=1"); +} + +#ifndef USE_LIBFABRIC_CXI + +void CxiTransport::init(ProxyCtx&) { unavailable(); } +void CxiTransport::register_main_buffer(void*, size_t, int) { unavailable(); } +void CxiTransport::register_atomic_buffer(void*, size_t, int) { unavailable(); } +LocalConnInfo CxiTransport::local_info() const { unavailable(); } +void CxiTransport::connect_peer(int, RemoteConnInfo const&) { unavailable(); } +void CxiTransport::post_write(int, uint64_t, uint64_t, uint64_t, uint32_t, + bool) { + unavailable(); +} +void CxiTransport::post_atomic_add(int, uint64_t, uint64_t, int64_t, bool) { + unavailable(); +} +void CxiTransport::post_barrier_atomic_add(int, uint64_t, size_t, uint64_t) { + unavailable(); +} +uint64_t CxiTransport::load_barrier_word(size_t) const { unavailable(); } +int CxiTransport::poll(TransportCompletion*, int) { unavailable(); } +void CxiTransport::destroy() {} + +#else + +namespace { + +void check_fi(int rc, char const* what) { + if (rc < 0) { + throw std::runtime_error(std::string(what) + + " failed: " + fi_strerror(-rc)); + } +} + +void check_gpu(gpuError_t rc, char const* what) { + if (rc != gpuSuccess) { + throw std::runtime_error(std::string(what) + + " failed: " + gpuGetErrorString(rc)); + } +} + +fid_mr* register_cuda_mr(fid_domain* domain, fid_ep* ep, void* ptr, size_t len, + int cuda_device, char const* label) { + int old_device = 0; + check_gpu(gpuGetDevice(&old_device), "gpuGetDevice"); + check_gpu(gpuSetDevice(cuda_device), "gpuSetDevice"); + + iovec iov = {}; + iov.iov_base = ptr; + iov.iov_len = len; + + fi_mr_attr attr = {}; + attr.mr_iov = &iov; + attr.iov_count = 1; + attr.access = + FI_SEND | FI_RECV | FI_WRITE | FI_READ | FI_REMOTE_WRITE | FI_REMOTE_READ; + attr.iface = FI_HMEM_CUDA; + attr.device.cuda = static_cast(cuda_device); + + fid_mr* mr = nullptr; + int rc = fi_mr_regattr(domain, &attr, 0, &mr); + gpuError_t restore_rc = gpuSetDevice(old_device); + check_fi(rc, label); + check_gpu(restore_rc, "gpuSetDevice(restore)"); + + check_fi(fi_mr_bind(mr, &ep->fid, 0), "fi_mr_bind"); + check_fi(fi_control(&mr->fid, FI_ENABLE, nullptr), "fi_control(FI_ENABLE)"); + return mr; +} + +fid_mr* register_host_mr(fid_domain* domain, fid_ep* ep, void* ptr, size_t len, + char const* label) { + iovec iov = {}; + iov.iov_base = ptr; + iov.iov_len = len; + + fi_mr_attr attr = {}; + attr.mr_iov = &iov; + attr.iov_count = 1; + attr.access = + FI_SEND | FI_RECV | FI_WRITE | FI_READ | FI_REMOTE_WRITE | FI_REMOTE_READ; + + fid_mr* mr = nullptr; + check_fi(fi_mr_regattr(domain, &attr, 0, &mr), label); + check_fi(fi_mr_bind(mr, &ep->fid, 0), "fi_mr_bind"); + check_fi(fi_control(&mr->fid, FI_ENABLE, nullptr), "fi_control(FI_ENABLE)"); + return mr; +} + +fi_info* select_cxi_info(fi_info* infos, int local_rank) { + char target[32]; + std::snprintf(target, sizeof(target), "cxi%d", local_rank); + + for (fi_info* cur = infos; cur; cur = cur->next) { + char const* name = (cur->domain_attr && cur->domain_attr->name) + ? cur->domain_attr->name + : ""; + if (std::strcmp(name, target) == 0) { + std::fprintf(stderr, "Selected CXI domain %s for local_rank %d\n", name, + local_rank); + return cur; + } + } + + std::string available; + for (fi_info* cur = infos; cur; cur = cur->next) { + char const* name = (cur->domain_attr && cur->domain_attr->name) + ? cur->domain_attr->name + : ""; + if (!available.empty()) available += " "; + available += name; + } + throw std::runtime_error("Could not find required CXI domain " + + std::string(target) + " for local_rank " + + std::to_string(local_rank) + + "; available domains: " + available); +} + +} // namespace + +void CxiTransport::init(ProxyCtx& ctx) { + ctx_ = &ctx; + + fi_info* hints = fi_allocinfo(); + if (!hints) throw std::runtime_error("fi_allocinfo failed"); + + hints->caps = FI_RMA | FI_WRITE | FI_ATOMIC | FI_FENCE | FI_REMOTE_WRITE | + FI_REMOTE_READ | FI_HMEM | FI_LOCAL_COMM | FI_REMOTE_COMM; + hints->mode = FI_CONTEXT; + hints->ep_attr->type = FI_EP_RDM; + hints->fabric_attr->prov_name = strdup("cxi"); + hints->domain_attr->mr_mode = + FI_MR_ENDPOINT | FI_MR_ALLOCATED | FI_MR_PROV_KEY; + hints->domain_attr->threading = FI_THREAD_DOMAIN; + + fi_info* info = nullptr; + int rc = fi_getinfo(FI_VERSION(1, 15), nullptr, nullptr, 0, hints, &info); + fi_freeinfo(hints); + check_fi(rc, "fi_getinfo(cxi)"); + + try { + fi_info* selected = select_cxi_info(info, ctx.local_rank); + + check_fi(fi_fabric(selected->fabric_attr, &fabric_, nullptr), "fi_fabric"); + check_fi(fi_domain(fabric_, selected, &domain_, nullptr), "fi_domain"); + + fi_cq_attr cq_attr = {}; + cq_attr.format = FI_CQ_FORMAT_DATA; + cq_attr.size = 131072; + check_fi(fi_cq_open(domain_, &cq_attr, &cq_, nullptr), "fi_cq_open"); + + fi_av_attr av_attr = {}; + av_attr.type = FI_AV_MAP; + check_fi(fi_av_open(domain_, &av_attr, &av_, nullptr), "fi_av_open"); + + check_fi(fi_endpoint(domain_, selected, &ep_, nullptr), "fi_endpoint"); + check_fi(fi_ep_bind(ep_, &cq_->fid, FI_TRANSMIT | FI_RECV), + "fi_ep_bind(cq)"); + check_fi(fi_ep_bind(ep_, &av_->fid, 0), "fi_ep_bind(av)"); + check_fi(fi_enable(ep_), "fi_enable"); + } catch (...) { + fi_freeinfo(info); + destroy(); + throw; + } + + fi_freeinfo(info); + + size_t ep_name_len = sizeof(local_info_.cxi_ep_name); + check_fi(fi_getname(&ep_->fid, local_info_.cxi_ep_name, &ep_name_len), + "fi_getname"); + local_info_.cxi_ep_name_len = static_cast(ep_name_len); + + atomic_operands_.assign(ProxyCtx::kMaxAtomicOps, 0); + atomic_operand_used_.assign(ProxyCtx::kMaxAtomicOps, 0); + atomic_operand_mr_ = + register_host_mr(domain_, ep_, atomic_operands_.data(), + atomic_operands_.size() * sizeof(atomic_operands_[0]), + "fi_mr_regattr(atomic operands)"); + barrier_words_.assign(8, 0); + barrier_mr_ = + register_host_mr(domain_, ep_, barrier_words_.data(), + barrier_words_.size() * sizeof(barrier_words_[0]), + "fi_mr_regattr(barrier words)"); + local_info_.cxi_barrier_mr_key = fi_mr_key(barrier_mr_); + + std::fprintf(stderr, "[CXI] libfabric/CXI endpoint initialized\n"); +} + +void CxiTransport::register_main_buffer(void* ptr, size_t len, + int cuda_device) { + if (!domain_ || !ep_) { + throw std::runtime_error("CxiTransport::init must run before MR setup"); + } + main_mr_ = register_cuda_mr(domain_, ep_, ptr, len, cuda_device, + "fi_mr_regattr(main)"); + main_buffer_ = ptr; + main_buffer_len_ = len; + local_info_.addr = reinterpret_cast(ptr); + local_info_.len = len; + local_info_.cxi_main_mr_key = fi_mr_key(main_mr_); + local_info_.rkey = static_cast(local_info_.cxi_main_mr_key); +} + +void CxiTransport::register_atomic_buffer(void* ptr, size_t len, + int cuda_device) { + if (!domain_ || !ep_) { + throw std::runtime_error("CxiTransport::init must run before MR setup"); + } + atomic_mr_ = register_cuda_mr(domain_, ep_, ptr, len, cuda_device, + "fi_mr_regattr(atomic)"); + atomic_buffer_ = ptr; + atomic_buffer_len_ = len; + local_info_.atomic_buffer_addr = reinterpret_cast(ptr); + local_info_.atomic_buffer_len = len; + local_info_.cxi_atomic_mr_key = fi_mr_key(atomic_mr_); + local_info_.atomic_buffer_rkey = + static_cast(local_info_.cxi_atomic_mr_key); +} + +LocalConnInfo CxiTransport::local_info() const { return local_info_; } + +void CxiTransport::connect_peer(int peer, RemoteConnInfo const& remote) { + if (!av_) throw std::runtime_error("CxiTransport::init must run first"); + if (remote.cxi_ep_name_len == 0 || + remote.cxi_ep_name_len > RemoteConnInfo::kMaxCxiEndpointName) { + throw std::runtime_error("invalid CXI endpoint address in remote info"); + } + + if (peer >= static_cast(peer_addrs_.size())) { + peer_addrs_.resize(peer + 1, FI_ADDR_UNSPEC); + peer_infos_.resize(peer + 1); + } + + fi_addr_t addr = FI_ADDR_UNSPEC; + int rc = fi_av_insert(av_, const_cast(remote.cxi_ep_name), 1, &addr, + 0, nullptr); + if (rc != 1) { + if (rc < 0) check_fi(rc, "fi_av_insert"); + throw std::runtime_error("fi_av_insert inserted no CXI peer address"); + } + peer_addrs_[peer] = addr; + peer_infos_[peer] = remote; +} + +void CxiTransport::post_write(int dst_rank, uint64_t wr_id, + uint64_t local_offset, uint64_t remote_offset, + uint32_t bytes, bool) { + if (dst_rank >= static_cast(peer_addrs_.size()) || + peer_addrs_[dst_rank] == FI_ADDR_UNSPEC) { + throw std::runtime_error("CXI peer is not connected"); + } + if (!main_mr_ || !main_buffer_) { + throw std::runtime_error("CXI main buffer MR is not registered"); + } + if (local_offset + bytes > main_buffer_len_) { + throw std::runtime_error("CXI write local range exceeds main buffer"); + } + + op_contexts_.push_back({}); + OpContext& op = op_contexts_.back(); + op.wr_id = wr_id; + op.is_write = true; + + auto* local = static_cast(main_buffer_) + local_offset; + ssize_t rc = fi_write(ep_, local, bytes, fi_mr_desc(main_mr_), + peer_addrs_[dst_rank], remote_offset, + peer_infos_[dst_rank].cxi_main_mr_key, &op.context); + check_fi(static_cast(rc), "fi_write"); +} + +void CxiTransport::post_atomic_add(int dst_rank, uint64_t wr_id, + uint64_t remote_atomic_offset, int64_t value, + bool fence) { + if (dst_rank >= static_cast(peer_addrs_.size()) || + peer_addrs_[dst_rank] == FI_ADDR_UNSPEC) { + throw std::runtime_error("CXI peer is not connected"); + } + if (!atomic_mr_) { + throw std::runtime_error("CXI atomic buffer MR is not registered"); + } + + op_contexts_.push_back({}); + OpContext& op = op_contexts_.back(); + op.wr_id = wr_id; + op.is_atomic = true; + + size_t slot = static_cast(-1); + for (size_t i = 0; i < atomic_operand_used_.size(); ++i) { + if (!atomic_operand_used_[i]) { + slot = i; + atomic_operand_used_[i] = 1; + break; + } + } + if (slot == static_cast(-1)) { + op_contexts_.pop_back(); + throw std::runtime_error("no free CXI atomic operand slot"); + } + op.atomic_operand_slot = slot; + atomic_operands_[slot] = static_cast(value); + + fi_ioc msg_iov = {}; + msg_iov.addr = &atomic_operands_[slot]; + msg_iov.count = 1; + void* desc = fi_mr_desc(atomic_operand_mr_); + + fi_rma_ioc rma_iov = {}; + rma_iov.addr = remote_atomic_offset; + rma_iov.count = 1; + rma_iov.key = peer_infos_[dst_rank].cxi_atomic_mr_key; + + fi_msg_atomic msg = {}; + msg.msg_iov = &msg_iov; + msg.desc = &desc; + msg.iov_count = 1; + msg.addr = peer_addrs_[dst_rank]; + msg.rma_iov = &rma_iov; + msg.rma_iov_count = 1; + msg.datatype = FI_UINT64; + msg.op = FI_SUM; + msg.context = &op.context; + + ssize_t rc = fi_atomicmsg(ep_, &msg, fence ? FI_FENCE : 0); + check_fi(static_cast(rc), "fi_atomicmsg"); +} + +void CxiTransport::post_barrier_atomic_add(int dst_rank, uint64_t wr_id, + size_t slot, uint64_t value) { + if (dst_rank >= static_cast(peer_addrs_.size()) || + peer_addrs_[dst_rank] == FI_ADDR_UNSPEC) { + throw std::runtime_error("CXI peer is not connected"); + } + if (peer_infos_[dst_rank].cxi_barrier_mr_key == 0) { + throw std::runtime_error("CXI barrier MR is not registered for peer"); + } + + op_contexts_.push_back({}); + OpContext& op = op_contexts_.back(); + op.wr_id = wr_id; + op.is_atomic = true; + + size_t operand_slot = static_cast(-1); + for (size_t i = 0; i < atomic_operand_used_.size(); ++i) { + if (!atomic_operand_used_[i]) { + operand_slot = i; + atomic_operand_used_[i] = 1; + break; + } + } + if (operand_slot == static_cast(-1)) { + op_contexts_.pop_back(); + throw std::runtime_error("no free CXI atomic operand slot"); + } + op.atomic_operand_slot = operand_slot; + atomic_operands_[operand_slot] = value; + + fi_ioc msg_iov = {}; + msg_iov.addr = &atomic_operands_[operand_slot]; + msg_iov.count = 1; + void* desc = fi_mr_desc(atomic_operand_mr_); + + fi_rma_ioc rma_iov = {}; + rma_iov.addr = slot * sizeof(uint64_t); + rma_iov.count = 1; + rma_iov.key = peer_infos_[dst_rank].cxi_barrier_mr_key; + + fi_msg_atomic msg = {}; + msg.msg_iov = &msg_iov; + msg.desc = &desc; + msg.iov_count = 1; + msg.addr = peer_addrs_[dst_rank]; + msg.rma_iov = &rma_iov; + msg.rma_iov_count = 1; + msg.datatype = FI_UINT64; + msg.op = FI_SUM; + msg.context = &op.context; + + ssize_t rc = fi_atomicmsg(ep_, &msg, FI_FENCE); + check_fi(static_cast(rc), "fi_atomicmsg(barrier)"); +} + +uint64_t CxiTransport::load_barrier_word(size_t slot) const { + if (slot >= barrier_words_.size()) { + throw std::runtime_error("CXI barrier slot out of range"); + } + return __atomic_load_n(&barrier_words_[slot], __ATOMIC_ACQUIRE); +} + +int CxiTransport::poll(TransportCompletion* out, int max) { + if (!cq_ || max <= 0) return 0; + + std::vector entries(max); + ssize_t rc = fi_cq_read(cq_, entries.data(), max); + if (rc == -FI_EAGAIN) return 0; + if (rc < 0) { + fi_cq_err_entry err = {}; + ssize_t err_rc = fi_cq_readerr(cq_, &err, 0); + if (err_rc >= 0) { + throw std::runtime_error( + std::string("CXI CQ error: ") + + fi_cq_strerror(cq_, err.prov_errno, err.err_data, nullptr, 0)); + } + check_fi(static_cast(rc), "fi_cq_read"); + } + + for (ssize_t i = 0; i < rc; ++i) { + auto* op = static_cast(entries[i].op_context); + out[i].wr_id = op ? op->wr_id : 0; + out[i].data = entries[i].data; + out[i].is_write = op ? op->is_write : ((entries[i].flags & FI_WRITE) != 0); + out[i].is_atomic = + op ? op->is_atomic : ((entries[i].flags & FI_ATOMIC) != 0); + out[i].is_recv_data = (entries[i].flags & FI_REMOTE_CQ_DATA) != 0; + out[i].status = 0; + + if (op) { + if (op->atomic_operand_slot != static_cast(-1) && + op->atomic_operand_slot < atomic_operand_used_.size()) { + atomic_operand_used_[op->atomic_operand_slot] = 0; + } + for (auto it = op_contexts_.begin(); it != op_contexts_.end(); ++it) { + if (&*it == op) { + op_contexts_.erase(it); + break; + } + } + } + } + return static_cast(rc); +} + +void CxiTransport::destroy() { + if (barrier_mr_) { + fi_close(&barrier_mr_->fid); + barrier_mr_ = nullptr; + } + if (atomic_operand_mr_) { + fi_close(&atomic_operand_mr_->fid); + atomic_operand_mr_ = nullptr; + } + if (atomic_mr_) { + fi_close(&atomic_mr_->fid); + atomic_mr_ = nullptr; + } + if (main_mr_) { + fi_close(&main_mr_->fid); + main_mr_ = nullptr; + } + if (ep_) { + fi_close(&ep_->fid); + ep_ = nullptr; + } + if (av_) { + fi_close(&av_->fid); + av_ = nullptr; + } + if (cq_) { + fi_close(&cq_->fid); + cq_ = nullptr; + } + if (domain_) { + fi_close(&domain_->fid); + domain_ = nullptr; + } + if (fabric_) { + fi_close(&fabric_->fid); + fabric_ = nullptr; + } +} + +#endif diff --git a/ep/src/internode.cu b/ep/src/internode.cu index e37968036..38f6b160f 100644 --- a/ep/src/internode.cu +++ b/ep/src/internode.cu @@ -16,7 +16,7 @@ namespace internode { struct SourceMeta { int src_rdma_rank, is_token_in_nvl_rank_bits; - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8, + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 8, "Invalid number of maximum NVL peers"); __forceinline__ SourceMeta() = default; @@ -39,6 +39,16 @@ struct SourceMeta { EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`"); +__device__ __forceinline__ uint64_t +pack_is_token_in_nvl_ranks(bool const* values) { + uint64_t packed = 0; +#pragma unroll + for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) { + packed |= static_cast(values[i]) << (i * 8); + } + return packed; +} + int get_source_meta_bytes() { return sizeof(SourceMeta); } __host__ __device__ __forceinline__ int get_num_bytes_per_token( @@ -366,17 +376,15 @@ __global__ void notify_dispatch( int total_count = 0, per_nvl_rank_count[NUM_MAX_NVL_PEERS] = {0}; for (int64_t i = token_start_idx + lane_id; i < token_end_idx; i += WARP_SIZE) { - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(uint64_t), - "Invalid number of NVL peers"); - auto is_token_in_rank_uint64 = *reinterpret_cast( - is_token_in_rank + i * num_ranks + - dst_rdma_rank * NUM_MAX_NVL_PEERS); - auto is_token_in_rank_values = - reinterpret_cast(&is_token_in_rank_uint64); + auto is_token_in_rank_values = is_token_in_rank + i * num_ranks + + dst_rdma_rank * NUM_MAX_NVL_PEERS; + bool is_token_in_any_nvl_rank = false; #pragma unroll - for (int j = 0; j < NUM_MAX_NVL_PEERS; ++j) + for (int j = 0; j < NUM_MAX_NVL_PEERS; ++j) { per_nvl_rank_count[j] += is_token_in_rank_values[j]; - total_count += (is_token_in_rank_uint64 != 0); + is_token_in_any_nvl_rank |= is_token_in_rank_values[j]; + } + total_count += is_token_in_any_nvl_rank; } // Warp reduce @@ -563,8 +571,7 @@ __global__ void __launch_bounds__( EP_DEVICE_ASSERT(num_topk <= WARP_SIZE); // RDMA symmetric layout - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(uint64_t), - "Invalid number of NVL peers"); + EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 8, "Invalid number of NVL peers"); auto hidden_bytes = hidden_int4 * sizeof(int4); auto scale_bytes = num_scales * sizeof(float); auto num_bytes_per_token = @@ -757,9 +764,9 @@ __global__ void __launch_bounds__( // Read RDMA rank existence uint64_t is_token_in_rank_uint64 = 0; if (lane_id < kNumRDMARanks) { - is_token_in_rank_uint64 = *(reinterpret_cast( + is_token_in_rank_uint64 = pack_is_token_in_nvl_ranks( is_token_in_rank + token_idx * num_ranks + - lane_id * NUM_MAX_NVL_PEERS)); + lane_id * NUM_MAX_NVL_PEERS); } // Acquire sequential lock @@ -806,9 +813,9 @@ __global__ void __launch_bounds__( // Read RDMA rank existence uint64_t is_token_in_rank_uint64 = 0; if (lane_id < kNumRDMARanks) { - is_token_in_rank_uint64 = __ldg(reinterpret_cast( + is_token_in_rank_uint64 = pack_is_token_in_nvl_ranks( is_token_in_rank + token_idx * num_ranks + - lane_id * NUM_MAX_NVL_PEERS)); + lane_id * NUM_MAX_NVL_PEERS); global_rdma_tail_idx += (is_token_in_rank_uint64 != 0); } diff --git a/ep/src/proxy.cpp b/ep/src/proxy.cpp index ec587c775..be0ebb79d 100644 --- a/ep/src/proxy.cpp +++ b/ep/src/proxy.cpp @@ -7,6 +7,8 @@ #include // for htonl, ntohl #include #include +#include +#include #include #include #include @@ -14,6 +16,47 @@ #include #include +static bool proxy_trace_enabled() { + static int const enabled = [] { + char const* env = std::getenv("UCCL_PROXY_TRACE"); + return env && env[0] != '\0' && env[0] != '0'; + }(); + return enabled != 0; +} + +static char const* cmd_name(CmdType cmd) { + switch (get_base_cmd(cmd)) { + case CmdType::WRITE: + return "WRITE"; + case CmdType::ATOMIC: + return "ATOMIC"; + case CmdType::QUIET: + return "QUIET"; + case CmdType::BARRIER: + return "BARRIER"; + case CmdType::EMPTY: + return "EMPTY"; + default: + return "UNKNOWN"; + } +} + +static int ranks_per_node_for_proxy(Proxy::Config const& cfg) { + if (cfg.num_nodes > 0 && cfg.num_ranks > 0) { + return std::max(1, cfg.num_ranks / cfg.num_nodes); + } + return MAX_NUM_GPUS; +} + +static bool is_normal_mode_remote_peer(Proxy::Config const& cfg, int my_rank, + int peer) { + if (!cfg.use_normal_mode) return true; + int const ranks_per_node = ranks_per_node_for_proxy(cfg); + return std::abs(peer - my_rank) % ranks_per_node == 0; +} + +static constexpr uint64_t kCxiInternalWrId = UINT64_MAX; + #ifndef USE_SUBSET_BARRIER static std::string shm_name_for_barrier(std::string const& ip, bool use_normal_mode, int thread_idx) { @@ -110,6 +153,11 @@ double Proxy::avg_wr_latency_us() const { uint64_t Proxy::completed_wr() const { return completion_count_; } +bool Proxy::use_cxi_transport() const { + char const* transport = std::getenv("UCCL_EP_TRANSPORT"); + return transport && std::string(transport) == "cxi"; +} + void Proxy::pin_thread_to_cpu_wrapper() { if (cfg_.pin_thread) { // TODO(MaoZiming): improves pinning. @@ -190,6 +238,159 @@ void Proxy::set_bench_d2h_channel_addrs(std::vector const& addrs) { void Proxy::init_common() { int const my_rank = cfg_.rank; + ctx_.local_rank = cfg_.local_rank; + ctx_.thread_idx = cfg_.thread_idx; + + if (use_cxi_transport()) { + if (ctxs_for_all_ranks_.empty()) { + fprintf(stderr, + "Error: peers metadata not set before init_common (peers_.size() " + "=%zu)\n", + peers_.size()); + std::abort(); + } + if (!atomic_buffer_ptr_) { + throw std::runtime_error( + "UCCL_EP_TRANSPORT=cxi requires atomic_buffer_ptr to be set"); + } + + int num_ranks = ctxs_for_all_ranks_.size(); + local_infos_.assign(num_ranks, RDMAConnectionInfo{}); + remote_infos_.assign(num_ranks, RDMAConnectionInfo{}); + cxi_transports_by_rank_.clear(); + cxi_transports_by_rank_.resize(num_ranks); + cxi_transport_ = nullptr; + for (int peer = 0; peer < num_ranks; ++peer) { + if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || + !is_normal_mode_remote_peer(cfg_, my_rank, peer)) + continue; + + auto transport = std::make_unique(); + transport->init(ctx_); + transport->register_main_buffer(cfg_.gpu_buffer, cfg_.total_size, + cfg_.local_rank); + transport->register_atomic_buffer(atomic_buffer_ptr_, kAtomicBufferSize, + cfg_.local_rank); + local_infos_[peer] = transport->local_info(); + if (!cxi_transport_) cxi_transport_ = transport.get(); + cxi_transports_by_rank_[peer] = std::move(transport); + } + + std::thread receiver_thread([this, num_ranks, my_rank]() { + for (int peer = 0; peer < num_ranks; ++peer) { + if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || + !is_normal_mode_remote_peer(cfg_, my_rank, peer)) + continue; + int actual_peer; + recv_connection_info_as_server(my_rank, &actual_peer, listen_fd_, + remote_infos_.data()); + } + }); + + for (int peer = 0; peer < num_ranks; ++peer) { + if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || + !is_normal_mode_remote_peer(cfg_, my_rank, peer)) + continue; + char const* peer_ip = peers_[peer].ip.c_str(); + int const peer_listen_port = peers_[peer].listen_ports[cfg_.thread_idx]; + send_connection_info_as_client(my_rank, peer, peer_ip, peer_listen_port, + &local_infos_[peer]); + } + + receiver_thread.join(); + + ctx_by_tag_.clear(); + ctx_by_tag_.resize(ctxs_for_all_ranks_.size() + 1, nullptr); + for (int peer = 0; peer < num_ranks; ++peer) { + auto& c = *ctxs_for_all_ranks_[peer]; + c.tag = static_cast(peer + 1); + if (c.tag >= ctx_by_tag_.size()) ctx_by_tag_.resize(c.tag + 1, nullptr); + ctx_by_tag_[c.tag] = &c; + if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || + !is_normal_mode_remote_peer(cfg_, my_rank, peer)) + continue; + + c.remote_addr = remote_infos_[peer].addr; + c.remote_len = remote_infos_[peer].len; + c.remote_rkey = remote_infos_[peer].rkey; + c.remote_atomic_buffer_addr = remote_infos_[peer].atomic_buffer_addr; + c.remote_atomic_buffer_len = remote_infos_[peer].atomic_buffer_len; + c.remote_atomic_buffer_rkey = remote_infos_[peer].atomic_buffer_rkey; + CxiTransport* transport = cxi_transport_for_rank(peer); + if (!transport) { + throw std::runtime_error("CXI peer transport is not initialized"); + } + transport->connect_peer(peer, remote_infos_[peer]); + if (proxy_trace_enabled()) { +#ifdef USE_LIBFABRIC_CXI + fprintf(stderr, + "[proxy-trace] rank=%d thread=%d connected CXI peer=%d " + "main_key=0x%llx atomic_key=0x%llx\n", + cfg_.rank, cfg_.thread_idx, peer, + (unsigned long long)remote_infos_[peer].cxi_main_mr_key, + (unsigned long long)remote_infos_[peer].cxi_atomic_mr_key); +#else + fprintf(stderr, + "[proxy-trace] rank=%d thread=%d connected CXI peer=%d\n", + cfg_.rank, cfg_.thread_idx, peer); +#endif + } + } + + if (cfg_.use_normal_mode) { + std::string const my_ip = peers_[cfg_.rank].ip; + std::vector local_ranks; + local_ranks.reserve(ctxs_for_all_ranks_.size()); + int leader_rank = cfg_.rank; + for (int r = 0; r < static_cast(peers_.size()); ++r) { + if (peers_[r].ip == my_ip) { + local_ranks.push_back(r); + if (r < leader_rank) leader_rank = r; + } + } + ctx_.num_local_ranks = static_cast(local_ranks.size()); + ctx_.node_leader_rank = leader_rank; + ctx_.local_rank = cfg_.local_rank; + ctx_.thread_idx = cfg_.thread_idx; + + if (ctx_.num_local_ranks > UCCL_MAX_LOCAL_RANKS) { + fprintf(stderr, "num_local_ranks=%d exceeds UCCL_MAX_LOCAL_RANKS=%d\n", + ctx_.num_local_ranks, static_cast(UCCL_MAX_LOCAL_RANKS)); + std::abort(); + } +#ifndef USE_SUBSET_BARRIER + std::string const shm_name = + shm_name_for_barrier(my_ip, cfg_.use_normal_mode, cfg_.thread_idx); + ctx_.lb = map_local_barrier_shm(shm_name, &ctx_.lb_owner); + if (!ctx_.lb) { + fprintf(stderr, "Failed to map local barrier shm: %s\n", + shm_name.c_str()); + std::abort(); + } + if (ctx_.lb_owner) { + ctx_.lb->full_mask = (ctx_.num_local_ranks >= 64) + ? ~0ULL + : ((1ULL << ctx_.num_local_ranks) - 1ULL); + for (int i = 0; i < ctx_.num_local_ranks; ++i) { + ctx_.lb->arrive_seq[i].store(0, std::memory_order_relaxed); + ctx_.lb->release_seq[i].store(0, std::memory_order_relaxed); + } + } else { + while (ctx_.lb->full_mask == 0ULL) cpu_relax(); + } +#endif + } + + fprintf(stderr, + "[Proxy] UCCL_EP_TRANSPORT=cxi initialized for rank %d thread %d\n", + cfg_.rank, cfg_.thread_idx); +#ifdef USE_MSCCLPP_FIFO_BACKEND + fifo_seq_.assign(cfg_.d2h_queues.size(), 0); + fifo_pending_.assign(cfg_.d2h_queues.size(), + std::deque>{}); +#endif + return; + } per_thread_rdma_init(ctx_, cfg_.gpu_buffer, cfg_.total_size, my_rank, cfg_.thread_idx, cfg_.device_index, cfg_.nic_local_rank); @@ -275,8 +476,7 @@ void Proxy::init_common() { for (int p = 0; p < num_ranks; ++p) { if (p == my_rank) continue; if (peers_[p].ip == peers_[my_rank].ip) continue; - if (cfg_.use_normal_mode && std::abs(p - my_rank) % MAX_NUM_GPUS != 0) - continue; + if (!is_normal_mode_remote_peer(cfg_, my_rank, p)) continue; ++num_active_peers; } int const ack_depth = @@ -316,8 +516,7 @@ void Proxy::init_common() { if (peer == my_rank) continue; // Skip rdma connection for intra-node. if (peers_[peer].ip == peers_[my_rank].ip) continue; - if (cfg_.use_normal_mode && std::abs(peer - my_rank) % MAX_NUM_GPUS != 0) - continue; + if (!is_normal_mode_remote_peer(cfg_, my_rank, peer)) continue; #ifdef EFA // Alias the shared SRD QPs from ctx_; dst_ah/dst_qpn (set later in // modify_qp_to_rtr) routes per WR via ibv_wr_set_ud_addr. @@ -344,8 +543,7 @@ void Proxy::init_common() { for (int peer = 0; peer < num_ranks; ++peer) { // Skip rdma connection for intra-node. if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - (cfg_.use_normal_mode && - std::abs(peer - my_rank) % MAX_NUM_GPUS != 0)) + !is_normal_mode_remote_peer(cfg_, my_rank, peer)) continue; int actual_peer; recv_connection_info_as_server(my_rank, &actual_peer, listen_fd_, @@ -356,7 +554,7 @@ void Proxy::init_common() { // Then send our info to all peers for (int peer = 0; peer < num_ranks; ++peer) { if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - (cfg_.use_normal_mode && std::abs(peer - my_rank) % MAX_NUM_GPUS != 0)) + !is_normal_mode_remote_peer(cfg_, my_rank, peer)) continue; char const* peer_ip = peers_[peer].ip.c_str(); int const peer_listen_port = peers_[peer].listen_ports[cfg_.thread_idx]; @@ -370,7 +568,7 @@ void Proxy::init_common() { // Verify remote info correctness for (int peer = 0; peer < num_ranks; ++peer) { if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - (cfg_.use_normal_mode && std::abs(peer - my_rank) % MAX_NUM_GPUS != 0)) + !is_normal_mode_remote_peer(cfg_, my_rank, peer)) continue; if (remote_infos_[peer].addr != peers_[peer].ptr) { fprintf(stderr, @@ -387,8 +585,7 @@ void Proxy::init_common() { if (peer == my_rank) continue; // Skip rdma connection for intra-node. if (peers_[peer].ip == peers_[my_rank].ip) continue; - if (cfg_.use_normal_mode && std::abs(peer - my_rank) % MAX_NUM_GPUS != 0) - continue; + if (!is_normal_mode_remote_peer(cfg_, my_rank, peer)) continue; auto& c = *ctxs_for_all_ranks_[peer]; // qp is different from each rank. @@ -537,6 +734,10 @@ void Proxy::init_remote() { } void Proxy::run_sender() { + if (use_cxi_transport()) { + throw std::runtime_error( + "CXI transport currently requires Proxy::run_dual"); + } printf("CPU sender thread %d started\n", cfg_.thread_idx); init_sender(); size_t seen = 0; @@ -549,6 +750,10 @@ void Proxy::run_sender() { } void Proxy::run_remote() { + if (use_cxi_transport()) { + throw std::runtime_error( + "CXI transport currently requires Proxy::run_dual"); + } printf("Remote CPU thread %d started\n", cfg_.thread_idx); init_remote(); std::set pending_atomic_updates; @@ -571,10 +776,10 @@ void Proxy::run_dual() { for (int peer = 0; peer < (int)ctxs_for_all_ranks_.size(); ++peer) { if (peer == cfg_.rank) continue; if (peers_[peer].ip == peers_[cfg_.rank].ip) continue; - if (cfg_.use_normal_mode && std::abs(peer - cfg_.rank) % MAX_NUM_GPUS != 0) - continue; + if (!is_normal_mode_remote_peer(cfg_, cfg_.rank, peer)) continue; auto& ctx_ptr = ctxs_for_all_ranks_[peer]; if (!ctx_ptr) continue; + if (use_cxi_transport()) continue; #ifndef EFA // EFA: posted once on the shared recv_ack_qp in init_common. local_post_ack_buf(*ctx_ptr, kSenderAckQueueDepth); @@ -592,10 +797,14 @@ void Proxy::run_dual() { while (ctx_.progress_run.load(std::memory_order_acquire)) { adaptive_sleeper_.maybe_sleep(ctx_); - poll_cq_dual(ctx_, acked_wrs_, cfg_.thread_idx, ring, ctx_by_tag_, - atomic_buffer_ptr_, cfg_.num_ranks, cfg_.num_experts, - pending_atomic_updates, cfg_.rank, cfg_.num_nodes, - adaptive_sleeper_, cfg_.use_normal_mode); + if (use_cxi_transport()) { + poll_cxi_completions(); + } else { + poll_cq_dual(ctx_, acked_wrs_, cfg_.thread_idx, ring, ctx_by_tag_, + atomic_buffer_ptr_, cfg_.num_ranks, cfg_.num_experts, + pending_atomic_updates, cfg_.rank, cfg_.num_nodes, + adaptive_sleeper_, cfg_.use_normal_mode); + } notify_gpu_completion(my_tail); post_gpu_command(my_tail, seen); #ifdef USE_RECEIVER_BARRIER @@ -742,10 +951,12 @@ void Proxy::post_gpu_command(uint64_t& my_tail, size_t& seen) { assert(!ctx_.barrier_inflight); assert(ctx_.barrier_wr == -1); ctx_.barrier_inflight = true; + if (use_cxi_transport()) ctx_.barrier_wr = unique_wr_id; } else if (get_base_cmd(cmd.cmd_type) == CmdType::QUIET) { assert(!ctx_.quiet_inflight); assert(ctx_.quiet_wr == -1); ctx_.quiet_inflight = true; + if (use_cxi_transport()) ctx_.quiet_wr = unique_wr_id; } break; } else { @@ -820,6 +1031,23 @@ void Proxy::post_gpu_command(uint64_t& my_tail, size_t& seen) { // Process all collected commands in batch if (!wrs_to_post.empty()) { + if (proxy_trace_enabled()) { + static std::atomic trace_batches{0}; + uint64_t const batch_id = + trace_batches.fetch_add(1, std::memory_order_relaxed); + if (batch_id < 64 || (batch_id % 1024) == 0) { + fprintf(stderr, + "[proxy-trace] rank=%d thread=%d batch=%lu cmds=%zu first=%s " + "dst=%d bytes=%u req_lptr=0x%lx req_rptr=0x%lx wr=0x%lx\n", + cfg_.rank, cfg_.thread_idx, batch_id, cmds_to_post.size(), + cmd_name(cmds_to_post[0].cmd_type), + static_cast(cmds_to_post[0].dst_rank), + cmds_to_post[0].bytes, + static_cast(cmds_to_post[0].req_lptr), + static_cast(cmds_to_post[0].req_rptr), + static_cast(wrs_to_post[0])); + } + } #ifdef MEASURE_PER_OP_LATENCY auto start = std::chrono::high_resolution_clock::now(); #endif @@ -940,9 +1168,130 @@ void Proxy::run_local() { cfg_.thread_idx, total_seen, cfg_.d2h_queues.size()); } +CxiTransport* Proxy::cxi_transport_for_rank(int rank) const { + if (rank < 0 || rank >= static_cast(cxi_transports_by_rank_.size())) { + return nullptr; + } + return cxi_transports_by_rank_[rank].get(); +} + +uint64_t Proxy::load_cxi_barrier_word_sum(size_t slot) const { + uint64_t total = 0; + for (auto const& transport : cxi_transports_by_rank_) { + if (transport) total += transport->load_barrier_word(slot); + } + return total; +} + +void Proxy::poll_cxi_completions() { + if (cxi_outstanding_ops_ == 0) return; + TransportCompletion completions[kMaxOutstandingSends]; + int const max_poll = static_cast(std::min( + cxi_outstanding_ops_, static_cast(kMaxOutstandingSends))); + int ne = 0; + for (auto& transport : cxi_transports_by_rank_) { + if (!transport || ne >= max_poll) continue; + int const got = transport->poll(completions + ne, max_poll - ne); + if (got > 0) ne += got; + } + for (int i = 0; i < ne; ++i) { + if (completions[i].status != 0) { + throw std::runtime_error("CXI completion returned non-zero status"); + } +#ifdef USE_LIBFABRIC_CXI + if (completions[i].wr_id != kCxiInternalWrId) { + acked_wrs_.insert(completions[i].wr_id); + } +#else + if (completions[i].wr_id != 0) { + acked_wrs_.insert(completions[i].wr_id); + } +#endif + } + if (static_cast(ne) > cxi_outstanding_ops_) { + throw std::runtime_error("CXI completion count exceeded outstanding ops"); + } + cxi_outstanding_ops_ -= static_cast(ne); + if (ne > 0) adaptive_sleeper_.update_timer(); +} + +void Proxy::post_cxi_commands(std::vector const& wrs_to_post, + std::vector const& cmds_to_post) { + if (!cxi_transport_) { + throw std::runtime_error("CXI transport has not been initialized"); + } + if (wrs_to_post.size() != cmds_to_post.size()) { + throw std::runtime_error("CXI command/wr_id size mismatch"); + } + + for (size_t i = 0; i < cmds_to_post.size(); ++i) { + auto const& cmd = cmds_to_post[i]; + int const dst_rank = static_cast(cmd.dst_rank); + CxiTransport* transport = cxi_transport_for_rank(dst_rank); + switch (get_base_cmd(cmd.cmd_type)) { + case CmdType::WRITE: { + if (!transport) { + throw std::runtime_error( + "CXI destination transport is not initialized"); + } + bool const low_latency = !cfg_.use_normal_mode; + uint64_t const local_offset = + decode_write_offset(cmd.req_lptr, low_latency); + uint64_t const remote_offset = + decode_write_offset(cmd.req_rptr, low_latency); + bool const has_signaling_atomic = + cfg_.use_normal_mode && cmd.atomic_offset > 0 && cmd.atomic_val > 0; + transport->post_write( + dst_rank, has_signaling_atomic ? kCxiInternalWrId : wrs_to_post[i], + local_offset, remote_offset, cmd.bytes, low_latency); + ++cxi_outstanding_ops_; + + if (has_signaling_atomic) { + transport->post_atomic_add( + dst_rank, wrs_to_post[i], cmd.atomic_offset, + static_cast(cmd.atomic_val), /*fence=*/true); + ++cxi_outstanding_ops_; + } + break; + } + case CmdType::ATOMIC: { + if (!transport) { + throw std::runtime_error( + "CXI destination transport is not initialized"); + } + int v = static_cast(cmd.value); + if (v == kLargeAtomicValue) v = kMaxSendAtomicValue; + if (get_is_combine(cmd.cmd_type)) v = 1; + transport->post_atomic_add( + dst_rank, wrs_to_post[i], cmd.req_rptr, + static_cast(static_cast(v)), /*fence=*/true); + ++cxi_outstanding_ops_; + break; + } + case CmdType::QUIET: + while (cxi_outstanding_ops_ > 0) { + poll_cxi_completions(); + if (cxi_outstanding_ops_ > 0) cpu_relax(); + } + acked_wrs_.insert(wrs_to_post[i]); + break; + case CmdType::BARRIER: + send_barrier(wrs_to_post[i]); + break; + default: + throw std::runtime_error("Unknown CXI command type"); + } + } +} + void Proxy::post_gpu_commands_mixed( std::vector const& wrs_to_post, std::vector const& cmds_to_post) { + if (use_cxi_transport()) { + post_cxi_commands(wrs_to_post, cmds_to_post); + return; + } + // Separate atomic operations from regular RDMA writes std::vector rdma_wrs, atomic_wrs, quiet_wrs, barrier_wrs; std::vector rdma_cmds, atomic_cmds, quiet_cmds, barrier_cmds; @@ -1035,6 +1384,18 @@ void Proxy::post_gpu_commands_mixed( 0) { return; } + if (proxy_trace_enabled()) { + static std::atomic trace_mixed_batches{0}; + uint64_t const batch_id = + trace_mixed_batches.fetch_add(1, std::memory_order_relaxed); + if (batch_id < 64 || (batch_id % 1024) == 0) { + fprintf(stderr, + "[proxy-trace] rank=%d thread=%d mixed=%lu rdma=%zu atomic=%zu " + "quiet=%zu barrier=%zu\n", + cfg_.rank, cfg_.thread_idx, batch_id, rdma_wrs.size(), + atomic_wrs.size(), quiet_wrs.size(), barrier_wrs.size()); + } + } // Handle regular RDMA writes if (!rdma_wrs.empty()) { post_rdma_async_batched(ctx_, cfg_.gpu_buffer, rdma_wrs.size(), rdma_wrs, @@ -1119,6 +1480,50 @@ void Proxy::quiet(std::vector wrs, std::vector cmds) { } void Proxy::destroy(bool free_gpu_buffer) { + if (cxi_transport_) { + for (auto& transport : cxi_transports_by_rank_) { + if (transport) transport->destroy(); + } + cxi_transports_by_rank_.clear(); + cxi_transport_ = nullptr; + + if (free_gpu_buffer && cfg_.gpu_buffer) { + cudaError_t e; + if (cfg_.free_buffer_with_cuda_free_host) { + e = cudaFreeHost(cfg_.gpu_buffer); + if (e != cudaSuccess) + fprintf(stderr, "[destroy] cudaFreeHost failed: %s\n", + cudaGetErrorString(e)); + } else { + e = cudaFree(cfg_.gpu_buffer); + if (e != cudaSuccess) + fprintf(stderr, "[destroy] cudaFree failed: %s\n", + cudaGetErrorString(e)); + } + if (e == cudaSuccess) cfg_.gpu_buffer = nullptr; + } + + acked_wrs_.clear(); + wr_id_to_start_time_.clear(); + ctxs_for_all_ranks_.clear(); + ctx_by_tag_.clear(); + local_infos_.clear(); + remote_infos_.clear(); +#ifndef USE_SUBSET_BARRIER + if (ctx_.lb) { + std::string const my_ip = (cfg_.rank < static_cast(peers_.size())) + ? peers_[cfg_.rank].ip + : ""; + std::string const shm_name = + shm_name_for_barrier(my_ip, cfg_.use_normal_mode, cfg_.thread_idx); + unmap_local_barrier_shm(shm_name, ctx_.lb, ctx_.lb_owner); + ctx_.lb = nullptr; + ctx_.lb_owner = false; + } +#endif + return; + } + for (auto& ctx_ptr : ctxs_for_all_ranks_) { if (!ctx_ptr) continue; if (ctx_ptr->qps_are_shared) continue; // owned by ctx_ (EFA) @@ -1374,8 +1779,11 @@ void Proxy::send_barrier(uint64_t wr) { assert(!ctx_.barrier_inflight && "only one barrier at a time"); ctx_.barrier_inflight = true; #endif - assert(ctx_.barrier_wr == -1 && "barrier_wr should be 0"); - ctx_.barrier_wr = wr; + if (ctx_.barrier_wr == -1) { + ctx_.barrier_wr = wr; + } else { + assert(use_cxi_transport() && static_cast(ctx_.barrier_wr) == wr); + } ctx_.barrier_seq = (ctx_.barrier_seq + 1) & BarrierImm::kSeqMask; if (cfg_.rank == ctx_.node_leader_rank) { @@ -1472,6 +1880,90 @@ void Proxy::barrier_check() { void Proxy::barrier_check() { if (!ctx_.barrier_inflight) return; + if (use_cxi_transport()) { + constexpr size_t kArrivalSlot = 0; + constexpr size_t kReleaseSlot = 1; + uint64_t const seq = ctx_.barrier_seq; + auto* lb = ctx_.lb; + + if (cfg_.rank == ctx_.node_leader_rank) { + bool all_local_arrived = true; + for (int lr = 0; lr < ctx_.num_local_ranks; ++lr) { + uint32_t seen = lb->arrive_seq[lr].load(std::memory_order_acquire); + if (seen != static_cast(seq)) { + all_local_arrived = false; + break; + } + } + if (!all_local_arrived) return; + + static thread_local uint64_t last_sent_seq = 0; + if (last_sent_seq != seq) { + last_sent_seq = seq; + if (cfg_.rank != 0) { + CxiTransport* transport = cxi_transport_for_rank(0); + if (!transport) { + throw std::runtime_error("CXI barrier transport to rank 0 missing"); + } + transport->post_barrier_atomic_add( + /*dst_rank=*/0, kCxiInternalWrId, kArrivalSlot, 1); + ++cxi_outstanding_ops_; + } + } + + if (cfg_.rank == 0) { + uint64_t const expected_arrivals = + seq * static_cast(std::max(0, cfg_.num_nodes - 1)); + if (load_cxi_barrier_word_sum(kArrivalSlot) < expected_arrivals) { + return; + } + + int const ranks_per_node = ranks_per_node_for_proxy(cfg_); + for (int r = ranks_per_node; r < static_cast(peers_.size()); + r += ranks_per_node) { + CxiTransport* transport = cxi_transport_for_rank(r); + if (!transport) { + throw std::runtime_error("CXI barrier release transport missing"); + } + transport->post_barrier_atomic_add(r, kCxiInternalWrId, kReleaseSlot, + 1); + ++cxi_outstanding_ops_; + } + for (int lr = 0; lr < ctx_.num_local_ranks; ++lr) { + lb->release_seq[lr].store(seq, std::memory_order_release); + } + + acked_wrs_.insert(ctx_.barrier_wr); +#ifndef USE_MSCCLPP_FIFO_BACKEND + ctx_.barrier_inflight = false; + ctx_.barrier_wr = -1; +#endif + return; + } + + if (load_cxi_barrier_word_sum(kReleaseSlot) < seq) return; + for (int lr = 0; lr < ctx_.num_local_ranks; ++lr) { + lb->release_seq[lr].store(seq, std::memory_order_release); + } + acked_wrs_.insert(ctx_.barrier_wr); +#ifndef USE_MSCCLPP_FIFO_BACKEND + ctx_.barrier_inflight = false; + ctx_.barrier_wr = -1; +#endif + return; + } + + if (lb->release_seq[ctx_.local_rank].load(std::memory_order_acquire) == + seq) { + acked_wrs_.insert(ctx_.barrier_wr); +#ifndef USE_MSCCLPP_FIFO_BACKEND + ctx_.barrier_inflight = false; + ctx_.barrier_wr = -1; +#endif + } + return; + } + auto* lb = ctx_.lb; uint64_t const seq = ctx_.barrier_seq; diff --git a/ep/src/uccl_ep.cc b/ep/src/uccl_ep.cc index c4eceed87..ba5f68bd8 100644 --- a/ep/src/uccl_ep.cc +++ b/ep/src/uccl_ep.cc @@ -42,6 +42,7 @@ std::map> g_proxies_by_dev; std::map>& proxies_by_dev() { return g_proxies_by_dev; } + } // namespace uccl #define NUM_MAX_LOCAL_EXPERTS 1024 @@ -275,7 +276,8 @@ nb::object make_rdma_buffer_dlpack_capsule(void* ptr, std::size_t bytes, } std::tuple allocate_rdma_buffer_dlpack( - std::size_t num_rdma_bytes, int device_index) { + std::size_t num_rdma_bytes, int device_index, + bool force_device_alloc = false) { std::size_t const alloc_bytes = std::max(num_rdma_bytes, 1); bool is_host_allocated = false; void* ptr = nullptr; @@ -290,7 +292,7 @@ std::tuple allocate_rdma_buffer_dlpack( CUDA_CHECK(cudaMemset(ptr, 0, alloc_bytes)); #else bool const use_host_alloc = - num_rdma_bytes > 0 && has_any_nic() && + !force_device_alloc && num_rdma_bytes > 0 && has_any_nic() && !can_register_gpu_memory_for_rdma(device_index, num_rdma_bytes); if (!use_host_alloc) { CUDA_CHECK(cudaMalloc(&ptr, alloc_bytes)); @@ -1915,10 +1917,13 @@ NB_MODULE(ep, m) { m.def( "get_rdma_buffer", - [](std::size_t num_rdma_bytes, int device_index) { - return allocate_rdma_buffer_dlpack(num_rdma_bytes, device_index); + [](std::size_t num_rdma_bytes, int device_index, + bool force_device_alloc) { + return allocate_rdma_buffer_dlpack(num_rdma_bytes, device_index, + force_device_alloc); }, nb::arg("num_rdma_bytes"), nb::arg("device_index"), + nb::arg("force_device_alloc") = false, R"doc( Allocate the RDMA scratch buffer outside PyTorch's CUDA allocator and return it as a DLPack capsule plus an `is_host_allocated` flag. diff --git a/ep/src/uccl_proxy.cpp b/ep/src/uccl_proxy.cpp index 57153b28a..b570d943e 100644 --- a/ep/src/uccl_proxy.cpp +++ b/ep/src/uccl_proxy.cpp @@ -5,6 +5,7 @@ #include "ring_buffer.cuh" #include #include +#include #include #include #include @@ -64,9 +65,14 @@ UcclProxy::UcclProxy(int thread_idx, uintptr_t gpu_buffer_addr, node_idx_ = node_idx; if (thread_idx == 0) { -#ifdef USE_GRACE_HOPPER - cudaMallocManaged(&atomic_buffer_ptr_, kAtomicBufferSize); +#ifdef USE_LIBFABRIC_CXI + cudaMalloc(&atomic_buffer_ptr_, kAtomicBufferSize); atomic_buffer_is_host_allocated_ = false; +#elif defined(USE_GRACE_HOPPER) + cudaMallocManaged(&atomic_buffer_ptr_, kAtomicBufferSize); + atomic_buffer_is_host_allocated_ = + false; // Uses unified memory when we are in Grace hopper (so Alps + // clariden will use this) #elif defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) hipExtMallocWithFlags(&atomic_buffer_ptr_, kAtomicBufferSize, hipDeviceMallocUncached); @@ -101,7 +107,7 @@ UcclProxy::~UcclProxy() { } if (thread_idx_ == 0 && atomic_buffer_ptr_) { -#if defined(USE_GRACE_HOPPER) +#if defined(USE_LIBFABRIC_CXI) || defined(USE_GRACE_HOPPER) cudaFree(atomic_buffer_ptr_); #elif defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) hipFree(atomic_buffer_ptr_); @@ -163,26 +169,39 @@ void UcclProxy::start(Mode m) { running_.store(true, std::memory_order_release); thread_ = std::thread([this]() { - if (is_intranode_) { - std::printf("UcclProxy: no peer IP set, running in local mode\n"); - proxy_->run_local(); - return; - } - switch (mode_) { - case Mode::Sender: - proxy_->run_sender(); - break; - case Mode::Remote: - proxy_->run_remote(); - break; - case Mode::Local: + try { + if (is_intranode_) { + std::printf("UcclProxy: no peer IP set, running in local mode\n"); proxy_->run_local(); - break; - case Mode::Dual: - proxy_->run_dual(); - break; - default: - break; + return; + } + switch (mode_) { + case Mode::Sender: + proxy_->run_sender(); + break; + case Mode::Remote: + proxy_->run_remote(); + break; + case Mode::Local: + proxy_->run_local(); + break; + case Mode::Dual: + proxy_->run_dual(); + break; + default: + break; + } + } catch (std::exception const& e) { + std::fprintf(stderr, "UcclProxy thread %d failed with exception: %s\n", + thread_idx_, e.what()); + std::fflush(stderr); + std::terminate(); + } catch (...) { + std::fprintf(stderr, + "UcclProxy thread %d failed with unknown exception\n", + thread_idx_); + std::fflush(stderr); + std::terminate(); } }); } From 7bdfab28387fd0f5690c6a0c858e7e4b383cc032 Mon Sep 17 00:00:00 2001 From: Yueyang Pan Date: Tue, 16 Jun 2026 19:05:34 +0200 Subject: [PATCH 02/13] [bugfix]: drain all proxy queues during sync - remove premature loop exits when posting quiet and barrier commands - ensure every proxy thread is drained before RDMA buffer reuse --- ep/include/uccl_ibgda.cuh | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ep/include/uccl_ibgda.cuh b/ep/include/uccl_ibgda.cuh index 1b56abc9e..f774eddff 100644 --- a/ep/include/uccl_ibgda.cuh +++ b/ep/include/uccl_ibgda.cuh @@ -361,7 +361,6 @@ __device__ static __forceinline__ void nvshmemi_ibgda_quiet( } } #endif - break; } // Then wait for all QUIET commands to complete @@ -410,7 +409,6 @@ __forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx( } } #endif - break; } // Then wait for each proxy’s barrier to complete @@ -421,4 +419,4 @@ __forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx( } } -} // namespace uccl \ No newline at end of file +} // namespace uccl From 098f7ed8c0d982be303db4c0dfa3cf45bf27565a Mon Sep 17 00:00:00 2001 From: Yueyang Pan Date: Tue, 16 Jun 2026 19:05:34 +0200 Subject: [PATCH 03/13] [bugfix]: drain all proxy queues during sync - remove premature loop exits when posting quiet and barrier commands - ensure every proxy thread is drained before RDMA buffer reuse --- ep/include/uccl_ibgda.cuh | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ep/include/uccl_ibgda.cuh b/ep/include/uccl_ibgda.cuh index 1b56abc9e..f774eddff 100644 --- a/ep/include/uccl_ibgda.cuh +++ b/ep/include/uccl_ibgda.cuh @@ -361,7 +361,6 @@ __device__ static __forceinline__ void nvshmemi_ibgda_quiet( } } #endif - break; } // Then wait for all QUIET commands to complete @@ -410,7 +409,6 @@ __forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx( } } #endif - break; } // Then wait for each proxy’s barrier to complete @@ -421,4 +419,4 @@ __forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx( } } -} // namespace uccl \ No newline at end of file +} // namespace uccl From 179fd96bf2f94a4e684b7d786bad5deeab04c02d Mon Sep 17 00:00:00 2001 From: Yueyang Pan Date: Fri, 19 Jun 2026 23:26:20 +0200 Subject: [PATCH 04/13] [fix]: guard CXI libfabric atomic include --- ep/src/cxi_transport.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ep/src/cxi_transport.cpp b/ep/src/cxi_transport.cpp index 79cdce55c..3e928881f 100644 --- a/ep/src/cxi_transport.cpp +++ b/ep/src/cxi_transport.cpp @@ -1,6 +1,10 @@ #include "cxi_transport.hpp" #include "util/gpu_rt.h" + +#ifdef USE_LIBFABRIC_CXI #include +#endif + #include #include #include From fbbaf3d29bb375e59e963b4d6395d4b735faba11 Mon Sep 17 00:00:00 2001 From: Yueyang Pan Date: Sun, 21 Jun 2026 19:48:12 +0200 Subject: [PATCH 05/13] [refactor]: remove proxy trace logging --- ep/src/proxy.cpp | 68 ------------------------------------------------ 1 file changed, 68 deletions(-) diff --git a/ep/src/proxy.cpp b/ep/src/proxy.cpp index be0ebb79d..2bfa7dcf8 100644 --- a/ep/src/proxy.cpp +++ b/ep/src/proxy.cpp @@ -16,31 +16,6 @@ #include #include -static bool proxy_trace_enabled() { - static int const enabled = [] { - char const* env = std::getenv("UCCL_PROXY_TRACE"); - return env && env[0] != '\0' && env[0] != '0'; - }(); - return enabled != 0; -} - -static char const* cmd_name(CmdType cmd) { - switch (get_base_cmd(cmd)) { - case CmdType::WRITE: - return "WRITE"; - case CmdType::ATOMIC: - return "ATOMIC"; - case CmdType::QUIET: - return "QUIET"; - case CmdType::BARRIER: - return "BARRIER"; - case CmdType::EMPTY: - return "EMPTY"; - default: - return "UNKNOWN"; - } -} - static int ranks_per_node_for_proxy(Proxy::Config const& cfg) { if (cfg.num_nodes > 0 && cfg.num_ranks > 0) { return std::max(1, cfg.num_ranks / cfg.num_nodes); @@ -321,20 +296,6 @@ void Proxy::init_common() { throw std::runtime_error("CXI peer transport is not initialized"); } transport->connect_peer(peer, remote_infos_[peer]); - if (proxy_trace_enabled()) { -#ifdef USE_LIBFABRIC_CXI - fprintf(stderr, - "[proxy-trace] rank=%d thread=%d connected CXI peer=%d " - "main_key=0x%llx atomic_key=0x%llx\n", - cfg_.rank, cfg_.thread_idx, peer, - (unsigned long long)remote_infos_[peer].cxi_main_mr_key, - (unsigned long long)remote_infos_[peer].cxi_atomic_mr_key); -#else - fprintf(stderr, - "[proxy-trace] rank=%d thread=%d connected CXI peer=%d\n", - cfg_.rank, cfg_.thread_idx, peer); -#endif - } } if (cfg_.use_normal_mode) { @@ -1031,23 +992,6 @@ void Proxy::post_gpu_command(uint64_t& my_tail, size_t& seen) { // Process all collected commands in batch if (!wrs_to_post.empty()) { - if (proxy_trace_enabled()) { - static std::atomic trace_batches{0}; - uint64_t const batch_id = - trace_batches.fetch_add(1, std::memory_order_relaxed); - if (batch_id < 64 || (batch_id % 1024) == 0) { - fprintf(stderr, - "[proxy-trace] rank=%d thread=%d batch=%lu cmds=%zu first=%s " - "dst=%d bytes=%u req_lptr=0x%lx req_rptr=0x%lx wr=0x%lx\n", - cfg_.rank, cfg_.thread_idx, batch_id, cmds_to_post.size(), - cmd_name(cmds_to_post[0].cmd_type), - static_cast(cmds_to_post[0].dst_rank), - cmds_to_post[0].bytes, - static_cast(cmds_to_post[0].req_lptr), - static_cast(cmds_to_post[0].req_rptr), - static_cast(wrs_to_post[0])); - } - } #ifdef MEASURE_PER_OP_LATENCY auto start = std::chrono::high_resolution_clock::now(); #endif @@ -1384,18 +1328,6 @@ void Proxy::post_gpu_commands_mixed( 0) { return; } - if (proxy_trace_enabled()) { - static std::atomic trace_mixed_batches{0}; - uint64_t const batch_id = - trace_mixed_batches.fetch_add(1, std::memory_order_relaxed); - if (batch_id < 64 || (batch_id % 1024) == 0) { - fprintf(stderr, - "[proxy-trace] rank=%d thread=%d mixed=%lu rdma=%zu atomic=%zu " - "quiet=%zu barrier=%zu\n", - cfg_.rank, cfg_.thread_idx, batch_id, rdma_wrs.size(), - atomic_wrs.size(), quiet_wrs.size(), barrier_wrs.size()); - } - } // Handle regular RDMA writes if (!rdma_wrs.empty()) { post_rdma_async_batched(ctx_, cfg_.gpu_buffer, rdma_wrs.size(), rdma_wrs, From b0a4d68512484a8c48f850550ab69d4a601a8f5f Mon Sep 17 00:00:00 2001 From: Yueyang Pan Date: Sun, 21 Jun 2026 19:49:31 +0200 Subject: [PATCH 06/13] [refactor]: share proxy connection exchange --- ep/include/proxy.hpp | 2 + ep/src/proxy.cpp | 104 +++++++++++++++---------------------------- 2 files changed, 37 insertions(+), 69 deletions(-) diff --git a/ep/include/proxy.hpp b/ep/include/proxy.hpp index 2f75aeb02..be002f6fd 100644 --- a/ep/include/proxy.hpp +++ b/ep/include/proxy.hpp @@ -107,6 +107,8 @@ class Proxy { void init_sender(); void init_remote(); bool use_cxi_transport() const; + bool should_connect_peer(int peer) const; + void exchange_peer_connection_info(int num_ranks); void notify_gpu_completion(uint64_t& my_tail); void post_gpu_command(uint64_t& my_tail, size_t& seen); diff --git a/ep/src/proxy.cpp b/ep/src/proxy.cpp index 2bfa7dcf8..fa7454bd7 100644 --- a/ep/src/proxy.cpp +++ b/ep/src/proxy.cpp @@ -133,6 +133,33 @@ bool Proxy::use_cxi_transport() const { return transport && std::string(transport) == "cxi"; } +bool Proxy::should_connect_peer(int peer) const { + return peer != cfg_.rank && peers_[peer].ip != peers_[cfg_.rank].ip && + is_normal_mode_remote_peer(cfg_, cfg_.rank, peer); +} + +void Proxy::exchange_peer_connection_info(int num_ranks) { + int const my_rank = cfg_.rank; + std::thread receiver_thread([this, num_ranks, my_rank]() { + for (int peer = 0; peer < num_ranks; ++peer) { + if (!should_connect_peer(peer)) continue; + int actual_peer; + recv_connection_info_as_server(my_rank, &actual_peer, listen_fd_, + remote_infos_.data()); + } + }); + + for (int peer = 0; peer < num_ranks; ++peer) { + if (!should_connect_peer(peer)) continue; + char const* peer_ip = peers_[peer].ip.c_str(); + int const peer_listen_port = peers_[peer].listen_ports[cfg_.thread_idx]; + send_connection_info_as_client(my_rank, peer, peer_ip, peer_listen_port, + &local_infos_[peer]); + } + + receiver_thread.join(); +} + void Proxy::pin_thread_to_cpu_wrapper() { if (cfg_.pin_thread) { // TODO(MaoZiming): improves pinning. @@ -236,9 +263,7 @@ void Proxy::init_common() { cxi_transports_by_rank_.resize(num_ranks); cxi_transport_ = nullptr; for (int peer = 0; peer < num_ranks; ++peer) { - if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - !is_normal_mode_remote_peer(cfg_, my_rank, peer)) - continue; + if (!should_connect_peer(peer)) continue; auto transport = std::make_unique(); transport->init(ctx_); @@ -251,28 +276,7 @@ void Proxy::init_common() { cxi_transports_by_rank_[peer] = std::move(transport); } - std::thread receiver_thread([this, num_ranks, my_rank]() { - for (int peer = 0; peer < num_ranks; ++peer) { - if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - !is_normal_mode_remote_peer(cfg_, my_rank, peer)) - continue; - int actual_peer; - recv_connection_info_as_server(my_rank, &actual_peer, listen_fd_, - remote_infos_.data()); - } - }); - - for (int peer = 0; peer < num_ranks; ++peer) { - if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - !is_normal_mode_remote_peer(cfg_, my_rank, peer)) - continue; - char const* peer_ip = peers_[peer].ip.c_str(); - int const peer_listen_port = peers_[peer].listen_ports[cfg_.thread_idx]; - send_connection_info_as_client(my_rank, peer, peer_ip, peer_listen_port, - &local_infos_[peer]); - } - - receiver_thread.join(); + exchange_peer_connection_info(num_ranks); ctx_by_tag_.clear(); ctx_by_tag_.resize(ctxs_for_all_ranks_.size() + 1, nullptr); @@ -281,9 +285,7 @@ void Proxy::init_common() { c.tag = static_cast(peer + 1); if (c.tag >= ctx_by_tag_.size()) ctx_by_tag_.resize(c.tag + 1, nullptr); ctx_by_tag_[c.tag] = &c; - if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - !is_normal_mode_remote_peer(cfg_, my_rank, peer)) - continue; + if (!should_connect_peer(peer)) continue; c.remote_addr = remote_infos_[peer].addr; c.remote_len = remote_infos_[peer].len; @@ -435,10 +437,7 @@ void Proxy::init_common() { // Pre-post recv WRs once on the shared recv_ack_qp, sized for all peers. int num_active_peers = 0; for (int p = 0; p < num_ranks; ++p) { - if (p == my_rank) continue; - if (peers_[p].ip == peers_[my_rank].ip) continue; - if (!is_normal_mode_remote_peer(cfg_, my_rank, p)) continue; - ++num_active_peers; + if (should_connect_peer(p)) ++num_active_peers; } int const ack_depth = std::min(static_cast(kMaxOutstandingRecvs), @@ -474,10 +473,7 @@ void Proxy::init_common() { c.atomic_old_values_buf = ctx_.atomic_old_values_buf; c.atomic_old_values_mr = ctx_.atomic_old_values_mr; - if (peer == my_rank) continue; - // Skip rdma connection for intra-node. - if (peers_[peer].ip == peers_[my_rank].ip) continue; - if (!is_normal_mode_remote_peer(cfg_, my_rank, peer)) continue; + if (!should_connect_peer(peer)) continue; #ifdef EFA // Alias the shared SRD QPs from ctx_; dst_ah/dst_qpn (set later in // modify_qp_to_rtr) routes per WR via ibv_wr_set_ud_addr. @@ -499,38 +495,11 @@ void Proxy::init_common() { usleep(50 * 1000); - // Out-of-band exchange info per pair: start receiver thread first - std::thread receiver_thread([this, num_ranks, my_rank]() { - for (int peer = 0; peer < num_ranks; ++peer) { - // Skip rdma connection for intra-node. - if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - !is_normal_mode_remote_peer(cfg_, my_rank, peer)) - continue; - int actual_peer; - recv_connection_info_as_server(my_rank, &actual_peer, listen_fd_, - remote_infos_.data()); - } - }); - - // Then send our info to all peers - for (int peer = 0; peer < num_ranks; ++peer) { - if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - !is_normal_mode_remote_peer(cfg_, my_rank, peer)) - continue; - char const* peer_ip = peers_[peer].ip.c_str(); - int const peer_listen_port = peers_[peer].listen_ports[cfg_.thread_idx]; - send_connection_info_as_client(my_rank, peer, peer_ip, peer_listen_port, - &local_infos_[peer]); - } - - // Wait for receiver thread to finish - receiver_thread.join(); + exchange_peer_connection_info(num_ranks); // Verify remote info correctness for (int peer = 0; peer < num_ranks; ++peer) { - if (peer == my_rank || peers_[peer].ip == peers_[my_rank].ip || - !is_normal_mode_remote_peer(cfg_, my_rank, peer)) - continue; + if (!should_connect_peer(peer)) continue; if (remote_infos_[peer].addr != peers_[peer].ptr) { fprintf(stderr, "Rank %d thread %d: Warning: remote addr mismatch for peer %d: " @@ -543,10 +512,7 @@ void Proxy::init_common() { // Bring each per-peer QP to RTR/RTS for (int peer = 0; peer < num_ranks; ++peer) { - if (peer == my_rank) continue; - // Skip rdma connection for intra-node. - if (peers_[peer].ip == peers_[my_rank].ip) continue; - if (!is_normal_mode_remote_peer(cfg_, my_rank, peer)) continue; + if (!should_connect_peer(peer)) continue; auto& c = *ctxs_for_all_ranks_[peer]; // qp is different from each rank. From f46f191b2c4e7e45f1831c3ce5aa1d18a6abd8ba Mon Sep 17 00:00:00 2001 From: Yueyang Pan Date: Sun, 21 Jun 2026 19:50:05 +0200 Subject: [PATCH 07/13] [refactor]: cache CXI transport selection --- ep/include/proxy.hpp | 1 + ep/src/proxy.cpp | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/ep/include/proxy.hpp b/ep/include/proxy.hpp index be002f6fd..4fcb55c8c 100644 --- a/ep/include/proxy.hpp +++ b/ep/include/proxy.hpp @@ -152,6 +152,7 @@ class Proxy { std::vector local_infos_, remote_infos_; std::vector ctx_by_tag_; void* atomic_buffer_ptr_ = nullptr; + bool use_cxi_transport_ = false; CxiTransport* cxi_transport_ = nullptr; std::vector> cxi_transports_by_rank_; size_t cxi_outstanding_ops_ = 0; diff --git a/ep/src/proxy.cpp b/ep/src/proxy.cpp index fa7454bd7..73bbe08e8 100644 --- a/ep/src/proxy.cpp +++ b/ep/src/proxy.cpp @@ -104,6 +104,9 @@ void unmap_local_barrier_shm(std::string const& name, LocalBarrier* lb, #endif Proxy::Proxy(Config const& cfg) : cfg_(cfg) { + char const* transport = std::getenv("UCCL_EP_TRANSPORT"); + use_cxi_transport_ = transport && std::string(transport) == "cxi"; + // Unset (-1) device/NIC ranks fall back to local_rank. if (cfg_.device_index < 0) cfg_.device_index = cfg_.local_rank; if (cfg_.nic_local_rank < 0) cfg_.nic_local_rank = cfg_.local_rank; @@ -128,10 +131,7 @@ double Proxy::avg_wr_latency_us() const { uint64_t Proxy::completed_wr() const { return completion_count_; } -bool Proxy::use_cxi_transport() const { - char const* transport = std::getenv("UCCL_EP_TRANSPORT"); - return transport && std::string(transport) == "cxi"; -} +bool Proxy::use_cxi_transport() const { return use_cxi_transport_; } bool Proxy::should_connect_peer(int peer) const { return peer != cfg_.rank && peers_[peer].ip != peers_[cfg_.rank].ip && From 7227ea48676c7f377fc57ac55743beada15454b7 Mon Sep 17 00:00:00 2001 From: Yueyang Pan Date: Sun, 21 Jun 2026 19:50:44 +0200 Subject: [PATCH 08/13] [fix]: keep CXI barrier sequence monotonic --- ep/src/proxy.cpp | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/ep/src/proxy.cpp b/ep/src/proxy.cpp index 73bbe08e8..50f043217 100644 --- a/ep/src/proxy.cpp +++ b/ep/src/proxy.cpp @@ -1682,7 +1682,14 @@ void Proxy::send_barrier(uint64_t wr) { } else { assert(use_cxi_transport() && static_cast(ctx_.barrier_wr) == wr); } - ctx_.barrier_seq = (ctx_.barrier_seq + 1) & BarrierImm::kSeqMask; + if (use_cxi_transport()) { + ++ctx_.barrier_seq; + if (ctx_.barrier_seq == 0) { + throw std::runtime_error("CXI barrier sequence wrapped"); + } + } else { + ctx_.barrier_seq = (ctx_.barrier_seq + 1) & BarrierImm::kSeqMask; + } if (cfg_.rank == ctx_.node_leader_rank) { if (ctx_.barrier_arrived.size() != static_cast(cfg_.num_nodes)) { From 6edc474dd4d8aaa2f22b0678f9f2c429b1c02d44 Mon Sep 17 00:00:00 2001 From: PanJason Date: Sun, 21 Jun 2026 20:37:16 +0200 Subject: [PATCH 09/13] [ep]: make MAX_NUM_GPUS configurable - keep the default MAX_NUM_GPUS value at 8 - allow setup.py and Makefile builds to override MAX_NUM_GPUS - preserve NUM_MAX_NVL_PEERS as a separate EP topology knob --- ep/Makefile | 5 +++++ ep/include/common.hpp | 2 ++ ep/setup.py | 8 ++++++++ 3 files changed, 15 insertions(+) diff --git a/ep/Makefile b/ep/Makefile index d7854dedf..bb0f8b511 100644 --- a/ep/Makefile +++ b/ep/Makefile @@ -114,6 +114,11 @@ ifdef NUM_MAX_NVL_PEERS override NVCCFLAGS += -DNUM_MAX_NVL_PEERS=$(NUM_MAX_NVL_PEERS) endif +ifdef MAX_NUM_GPUS + override CXXFLAGS += -DMAX_NUM_GPUS=$(MAX_NUM_GPUS) + override NVCCFLAGS += -DMAX_NUM_GPUS=$(MAX_NUM_GPUS) +endif + CXXFLAGS += $(EFA_CFLAGS) $(GH_CFLAGS) $(DMABUF_CFLAGS) $(CXI_CFLAGS) $(NORMAL_CFLAGS) NVCCFLAGS += $(EFA_CFLAGS) $(GH_CFLAGS) $(DMABUF_CFLAGS) $(CXI_CFLAGS) $(NORMAL_CFLAGS) diff --git a/ep/include/common.hpp b/ep/include/common.hpp index f3aee7b9d..72f686128 100644 --- a/ep/include/common.hpp +++ b/ep/include/common.hpp @@ -85,7 +85,9 @@ static constexpr int kMaxMRChunks = 128; #define kReorderingBufferSize 16 // Right now only 4 bits. #define kRemoteBufferSize (kBatchSize * kNumProxyThs * kObjectSize * 100) #define MAIN_THREAD_CPU_IDX 31 +#ifndef MAX_NUM_GPUS #define MAX_NUM_GPUS 8 +#endif #define RECEIVER_BATCH_SIZE 16 #define kAtomicWrTag 0xa70a000000000000ULL #define kAtomicMask 0x0000FFFFFFFFFFFFULL diff --git a/ep/setup.py b/ep/setup.py index 28a74057c..ae22ba00b 100644 --- a/ep/setup.py +++ b/ep/setup.py @@ -153,6 +153,14 @@ def run(self): cxx_flags.extend(abi_flags) nvcc_flags.extend(abi_flags) + max_num_gpus = os.getenv("MAX_NUM_GPUS") + if max_num_gpus: + max_num_gpus = int(max_num_gpus) + max_num_gpus_flag = f"-DMAX_NUM_GPUS={max_num_gpus}" + print(f"Building with MAX_NUM_GPUS={max_num_gpus}") + cxx_flags.append(max_num_gpus_flag) + nvcc_flags.append(max_num_gpus_flag) + # Collect header files for dependency tracking header_files = [] for inc_dir in include_dirs: From 2adc58621442c11719ef46eb5129e3e46d4dd71a Mon Sep 17 00:00:00 2001 From: PanJason Date: Mon, 22 Jun 2026 17:56:01 +0200 Subject: [PATCH 10/13] [bugfix]: keep sync barrier single-queued - restore the sync loop exit so BARRIER remains posted through one representative queue - keep QUIET handling separate from BARRIER semantics for CXI timeout debugging --- ep/include/uccl_ibgda.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/ep/include/uccl_ibgda.cuh b/ep/include/uccl_ibgda.cuh index f774eddff..a79e71988 100644 --- a/ep/include/uccl_ibgda.cuh +++ b/ep/include/uccl_ibgda.cuh @@ -252,6 +252,7 @@ __device__ __forceinline__ void nvshmemi_ibgda_amo_nonfetch_add( } } #endif + break; } } From d43db3177b10dd6797d2e4712a6f16046a8ff15f Mon Sep 17 00:00:00 2001 From: PanJason Date: Mon, 22 Jun 2026 17:56:16 +0200 Subject: [PATCH 11/13] [bugfix]: preserve fifo control wr ids - widen quiet and barrier wr id storage so FIFO ring bits are not truncated - keep the existing negative sentinel while preserving high 32-bit queue indices --- ep/include/proxy_ctx.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ep/include/proxy_ctx.hpp b/ep/include/proxy_ctx.hpp index 85f12d4ef..ebb2eb6ef 100644 --- a/ep/include/proxy_ctx.hpp +++ b/ep/include/proxy_ctx.hpp @@ -201,10 +201,10 @@ struct ProxyCtx { // Async-barrier state (single inflight assumed) bool barrier_inflight = false; uint64_t barrier_seq = 0; - int barrier_wr = -1; + int64_t barrier_wr = -1; bool quiet_inflight = false; - int quiet_wr = -1; + int64_t quiet_wr = -1; // Rank-0 bookkeeping std::vector barrier_arrived; // size = num_ranks; 1 if arrival seen From a5420121c64bd660f2e7b47c81e4aa4d2eb72ddf Mon Sep 17 00:00:00 2001 From: PanJason Date: Mon, 22 Jun 2026 17:57:02 +0200 Subject: [PATCH 12/13] [bugfix]: quiet all CXI proxy queues - post QUIET to every CXI D2H queue so sibling proxy queues are drained - track the posted queue index and wait on the exact queue for completion --- ep/include/uccl_ibgda.cuh | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/ep/include/uccl_ibgda.cuh b/ep/include/uccl_ibgda.cuh index a79e71988..760b867e1 100644 --- a/ep/include/uccl_ibgda.cuh +++ b/ep/include/uccl_ibgda.cuh @@ -326,16 +326,23 @@ __device__ static __forceinline__ void nvshmemi_ibgda_quiet( EP_DEVICE_ASSERT( num_d2h_channel_addrs % kChannelPerProxy == 0 && "num_d2h_channel_addrs must be multiple of kChannelPerProxy"); - /* NOTE(MaoZiming): This is sent to all proxy threads. Since each proxy - * thread manages kChannelPerProxy ring buffers, we just need to post a quiet - * command to one out of the kChannelPerProxy ring buffer per cpu thread. */ EP_DEVICE_ASSERT(num_d2h_channel_addrs % kChannelPerProxy == 0); EP_DEVICE_ASSERT(num_d2h_channel_addrs / kChannelPerProxy == kNumProxyThs); - // First, atomically commit QUIET to one ring per proxy - uint64_t slots[kNumProxyThs]; + // First, atomically commit QUIET to all CXI rings. Each proxy thread owns + // multiple D2H rings, so posting to one ring per proxy is not enough to fence + // sibling rings before RDMA buffer reuse. +#if defined(USE_LIBFABRIC_CXI) + constexpr int kQuietStride = 1; + constexpr int kMaxQuietPosts = kNumProxyThs * kChannelPerProxy; +#else + constexpr int kQuietStride = kChannelPerProxy; + constexpr int kMaxQuietPosts = kNumProxyThs; +#endif + uint64_t slots[kMaxQuietPosts]; + int posted_d2h_channel_idxs[kMaxQuietPosts]; int num_posted = 0; for (int d2h_channel_idx = 0; d2h_channel_idx < num_d2h_channel_addrs; - d2h_channel_idx += kChannelPerProxy) { + d2h_channel_idx += kQuietStride) { auto* h = reinterpret_cast( static_cast(d2h_channel_addrs[d2h_channel_idx])); #ifdef USE_MSCCLPP_FIFO_BACKEND @@ -344,7 +351,9 @@ __device__ static __forceinline__ void nvshmemi_ibgda_quiet( TransferCmd cmd{}; cmd.cmd_type = CmdType::QUIET; h->atomic_set_and_commit(cmd, &slot); - slots[num_posted++] = slot; + slots[num_posted] = slot; + posted_d2h_channel_idxs[num_posted] = d2h_channel_idx; + ++num_posted; } #else while (true) { @@ -357,6 +366,7 @@ __device__ static __forceinline__ void nvshmemi_ibgda_quiet( cmd.cmd_type = CmdType::QUIET; h->atomic_set_and_commit(cmd, &slot); slots[num_posted] = slot; + posted_d2h_channel_idxs[num_posted] = d2h_channel_idx; ++num_posted; break; } @@ -367,7 +377,7 @@ __device__ static __forceinline__ void nvshmemi_ibgda_quiet( // Then wait for all QUIET commands to complete for (int i = 0; i < num_posted; ++i) { auto* h = reinterpret_cast( - static_cast(d2h_channel_addrs[i * kChannelPerProxy])); + static_cast(d2h_channel_addrs[posted_d2h_channel_idxs[i]])); wait_until_cmd_consumed(h, slots[i], nvl_rank, CmdType::QUIET); } } From 6c1035b34bb1e62610a6e45e8d7fe83c6f1a44e5 Mon Sep 17 00:00:00 2001 From: Fergus Finn Date: Mon, 22 Jun 2026 18:15:13 +0100 Subject: [PATCH 13/13] fix(merge): drop stray break in amo_nonfetch_add #997 refactored the per-channel for-loop into a direct index computation (removing its trailing break); #998 touched adjacent lines so the 3-way merge kept #998's body, resurrecting a break now outside any loop. Invalid C++ (uccl_ibgda.cuh: 'break may only be used within a loop or switch'). Semantic resolution of the #997+#998 compose that the clean textual merge missed. --- ep/include/uccl_ibgda.cuh | 1 - 1 file changed, 1 deletion(-) diff --git a/ep/include/uccl_ibgda.cuh b/ep/include/uccl_ibgda.cuh index 760b867e1..b2909b039 100644 --- a/ep/include/uccl_ibgda.cuh +++ b/ep/include/uccl_ibgda.cuh @@ -252,7 +252,6 @@ __device__ __forceinline__ void nvshmemi_ibgda_amo_nonfetch_add( } } #endif - break; } }