Add CXI transport support for EP

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,18 +109,23 @@ 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
+
 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) $(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

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:

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=<local_gpu_count> --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=<local_gpu_count> --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))

ep/include/cxi_transport.hpp

@@ -0,0 +1,76 @@
+#pragma once
+
+#include "transport.hpp"
+
+#ifdef USE_LIBFABRIC_CXI
+#include <rdma/fabric.h>
+#include <rdma/fi_cm.h>
+#include <rdma/fi_domain.h>
+#include <rdma/fi_endpoint.h>
+#include <rdma/fi_eq.h>
+#endif
+
+#include <list>
+#include <stdexcept>
+#include <vector>
+
+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<fi_addr_t> peer_addrs_;
+  std::vector<RemoteConnInfo> peer_infos_;
+  std::vector<uint64_t> atomic_operands_;
+  std::vector<uint8_t> atomic_operand_used_;
+  std::vector<uint64_t> 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<size_t>(-1);
+  };
+  std::list<OpContext> op_contexts_;
+#endif
+};

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

ep/include/proxy.hpp

@@ -2,13 +2,15 @@
 #define PROXY_HPP
 
 #include "common.hpp"
+#include "cxi_transport.hpp"
 #include "proxy_ctx.hpp"
 #include "rdma.hpp"
 #include "ring_buffer.cuh"
 #include "util/gpu_rt.h"
 #include <algorithm>
 #include <atomic>
 #include <chrono>
+#include <memory>
 #include <mutex>
 #include <thread>
 #include <unordered_map>
@@ -104,11 +106,19 @@ class Proxy {
   void init_common();
   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);
   void post_gpu_commands_mixed(std::vector<uint64_t> const& wrs_to_post,
                                std::vector<TransferCmd> const& cmds_to_post);
+  void post_cxi_commands(std::vector<uint64_t> const& wrs_to_post,
+                         std::vector<TransferCmd> 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 +151,11 @@ class Proxy {
   std::vector<std::unique_ptr<ProxyCtx>> ctxs_for_all_ranks_;
   std::vector<RDMAConnectionInfo> local_infos_, remote_infos_;
   std::vector<ProxyCtx*> ctx_by_tag_;
-  void* atomic_buffer_ptr_;
+  void* atomic_buffer_ptr_ = nullptr;
+  bool use_cxi_transport_ = false;
+  CxiTransport* cxi_transport_ = nullptr;
+  std::vector<std::unique_ptr<CxiTransport>> cxi_transports_by_rank_;
+  size_t cxi_outstanding_ops_ = 0;
   std::vector<TransferCmd> postponed_atomics_;
   std::vector<uint64_t> postponed_wr_ids_;
 

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

ep/include/transport.hpp

@@ -0,0 +1,42 @@
+#pragma once
+
+#include "rdma.hpp"
+#include <cstddef>
+#include <cstdint>
+
+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;
+};

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
+}  // namespace uccl

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_;

ep/setup.py

@@ -171,6 +171,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
@@ -436,6 +454,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")