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 diff --git a/ep/include/uccl_ibgda.cuh b/ep/include/uccl_ibgda.cuh index f774eddff..40b48e005 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]; + // 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 @@ -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); } } @@ -409,6 +419,10 @@ __forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx( } } #endif + // NOTE: the `break` here is intentional (only post to proxy thread 0): + // unlike QUIET, one proxy thread per GPU suffices since one GPU proxy + // thread is enough to form a complete barrier across all ranks. + break; } // Then wait for each proxy’s barrier to complete