diff --git a/ep/bench/buffer.py b/ep/bench/buffer.py index b73f6fb8e..b4d75cb18 100644 --- a/ep/bench/buffer.py +++ b/ep/bench/buffer.py @@ -781,7 +781,7 @@ def get_combine_config(num_ranks: int) -> Config: config_map = { 2: Config(Buffer.num_sms, 10, 256, 6, 128), 4: Config(Buffer.num_sms, 9, 256, 6, 128), - 8: Config(Buffer.num_sms, 4, 256, 6, 128), + 8: Config(Buffer.num_sms, 4, 256, 8, 128), 16: Config(Buffer.num_sms, 4, 288, 12, 512 if Buffer._is_efa() else 128), 24: Config(Buffer.num_sms, 1, 288, 8, 128), 32: Config(Buffer.num_sms, 1, 288, 8, 512 if Buffer._is_efa() else 128), diff --git a/ep/include/proxy.hpp b/ep/include/proxy.hpp index 2f75aeb02..d4b2255fa 100644 --- a/ep/include/proxy.hpp +++ b/ep/include/proxy.hpp @@ -150,6 +150,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/include/uccl_ibgda.cuh b/ep/include/uccl_ibgda.cuh index f774eddff..5ddee6162 100644 --- a/ep/include/uccl_ibgda.cuh +++ b/ep/include/uccl_ibgda.cuh @@ -325,16 +325,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]; + // CXI work is sharded across every D2H ring. A QUIET on one + // ring per proxy only fences commands that were already posted to the + // transport from that ring, not commands still queued on sibling rings. +#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 @@ -343,7 +350,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) { @@ -356,6 +365,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; } @@ -366,7 +376,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); } } @@ -381,7 +391,8 @@ __forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx( uint64_t slots[kNumProxyThs]; int num_posted = 0; - // First, post one BARRIER command per proxy + // First, post one BARRIER command per proxy thread. QUIET handles per-ring + // write fencing; the barrier drains each proxy thread's rendezvous. for (int d2h_channel_idx = 0; d2h_channel_idx < num_d2h_channel_addrs; d2h_channel_idx += kChannelPerProxy) { auto* h = reinterpret_cast( @@ -411,7 +422,7 @@ __forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx( #endif } - // Then wait for each proxy’s barrier to complete + // Then wait for each proxy thread's barrier to complete. for (int i = 0; i < num_posted; ++i) { auto* h = reinterpret_cast( static_cast(d2h_channel_addrs[i * kChannelPerProxy])); diff --git a/ep/src/proxy.cpp b/ep/src/proxy.cpp index be0ebb79d..8ba0b87e5 100644 --- a/ep/src/proxy.cpp +++ b/ep/src/proxy.cpp @@ -128,7 +128,13 @@ void unmap_local_barrier_shm(std::string const& name, LocalBarrier* lb, } #endif +static bool env_requests_cxi_transport() { + char const* transport = std::getenv("UCCL_EP_TRANSPORT"); + return transport && std::string(transport) == "cxi"; +} + Proxy::Proxy(Config const& cfg) : cfg_(cfg) { + use_cxi_transport_ = env_requests_cxi_transport(); // 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; @@ -153,10 +159,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_; } void Proxy::pin_thread_to_cpu_wrapper() { if (cfg_.pin_thread) { @@ -853,13 +856,16 @@ void Proxy::notify_gpu_completion(uint64_t& my_tail) { std::memory_order_release); } - if (ctx_.quiet_wr != -1 && front_wr == (uint64_t)ctx_.quiet_wr) { + uint32_t const front_seq = static_cast(front_wr); + if (ctx_.quiet_wr != -1 && + front_seq == static_cast(ctx_.quiet_wr)) { ctx_.quiet_inflight = false; ctx_.quiet_wr = -1; fifo->pop(); } - if (ctx_.barrier_wr != -1 && front_wr == (uint64_t)ctx_.barrier_wr) { + if (ctx_.barrier_wr != -1 && + front_seq == static_cast(ctx_.barrier_wr)) { ctx_.barrier_inflight = false; ctx_.barrier_wr = -1; fifo->pop(); @@ -1782,7 +1788,8 @@ void Proxy::send_barrier(uint64_t wr) { if (ctx_.barrier_wr == -1) { ctx_.barrier_wr = wr; } else { - assert(use_cxi_transport() && static_cast(ctx_.barrier_wr) == wr); + assert(use_cxi_transport() && static_cast(ctx_.barrier_wr) == + static_cast(wr)); } ctx_.barrier_seq = (ctx_.barrier_seq + 1) & BarrierImm::kSeqMask;