Skip to content

fix(ep): drain CXI FIFO rings before quiet #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
fergusfinn wants to merge 2 commits into from
Closed

fix(ep): drain CXI FIFO rings before quiet #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
05a7d4d
fix(ep): drain CXI FIFO rings before quiet
fergusfinn Jun 17, 2026
f248138
fix(ep): keep barriers per proxy thread
fergusfinn Jun 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ep/bench/buffer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
1 change: 1 addition & 0 deletions ep/include/proxy.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,7 @@ class Proxy {
std::vector<RDMAConnectionInfo> local_infos_, remote_infos_;
std::vector<ProxyCtx*> ctx_by_tag_;
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;
Expand Down
31 changes: 21 additions & 10 deletions ep/include/uccl_ibgda.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Comment thread
fergusfinn marked this conversation as resolved.
// 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<d2hq::D2HHandle*>(
static_cast<uintptr_t>(d2h_channel_addrs[d2h_channel_idx]));
#ifdef USE_MSCCLPP_FIFO_BACKEND
Expand All @@ -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) {
Expand All @@ -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;
}
Expand All @@ -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<d2hq::D2HHandle*>(
static_cast<uintptr_t>(d2h_channel_addrs[i * kChannelPerProxy]));
static_cast<uintptr_t>(d2h_channel_addrs[posted_d2h_channel_idxs[i]]));
wait_until_cmd_consumed(h, slots[i], nvl_rank, CmdType::QUIET);
}
}
Expand All @@ -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<d2hq::D2HHandle*>(
Expand Down Expand Up @@ -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<d2hq::D2HHandle*>(
static_cast<uintptr_t>(d2h_channel_addrs[i * kChannelPerProxy]));
Expand Down
21 changes: 14 additions & 7 deletions ep/src/proxy.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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) {
Expand Down Expand Up @@ -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<uint32_t>(front_wr);
if (ctx_.quiet_wr != -1 &&
front_seq == static_cast<uint32_t>(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<uint32_t>(ctx_.barrier_wr)) {
ctx_.barrier_inflight = false;
ctx_.barrier_wr = -1;
fifo->pop();
Expand Down Expand Up @@ -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<uint64_t>(ctx_.barrier_wr) == wr);
assert(use_cxi_transport() && static_cast<uint32_t>(ctx_.barrier_wr) ==
static_cast<uint32_t>(wr));
}
ctx_.barrier_seq = (ctx_.barrier_seq + 1) & BarrierImm::kSeqMask;

Expand Down