From 4183c020a12c5165c01fd14e6ef4780326458c1e Mon Sep 17 00:00:00 2001 From: Keita Watanabe Date: Thu, 27 Aug 2026 22:43:21 +0000 Subject: [PATCH] comm: lift the GIN scale-out ceiling by making the rail barrier a counting barrier The rail barrier gave every peer a dedicated indexed signal, so its demand grew with the team while the per-context EFA signal budget (21 at the shipped 11 contexts) did not; init refused past 22 NVLink domains. The rail team now uses a counting barrier: every peer adds 1 to one reserved signal id and the waiter expects kNumRanks - 1 increments, costing one id at any team size. kNumReservedBarrierSignals takes that id off the bottom of every context's id space and data_signal_id() shifts data-path ids past it, so data channels cannot collide with the barrier. The world barrier is unchanged (its call sites need per-peer release semantics). Static/host asserts forbid a scale-up and a scale-out GIN barrier being live concurrently, and a per-context serviceability static_assert replaces the cross-context total check. --- csrc/kernels/backend/nccl.cu | 16 ++- deep_ep/include/deep_ep/common/comm.cuh | 128 +++++++++++++----- .../deep_ep/common/gin_resource_alloc.cuh | 84 +++++++++--- 3 files changed, 171 insertions(+), 57 deletions(-) diff --git a/csrc/kernels/backend/nccl.cu b/csrc/kernels/backend/nccl.cu index 2cc4df204..df8d03c81 100644 --- a/csrc/kernels/backend/nccl.cu +++ b/csrc/kernels/backend/nccl.cu @@ -137,9 +137,14 @@ NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm, gin_config.gin_indexed_signals_cnt = 0; } - EP_HOST_ASSERT(gin_config.gin_indexed_signals_cnt >= (num_rdma_ranks - 1) and - "GIN indexed-signal budget cannot give each peer rail team a dedicated " - "signal; reduce num_allocated_qps to raise the per-context signal count"); + // The rail barrier is a counting barrier (`gin_barrier_wo_local_sync`, comm.cuh): + // it costs `kNumReservedBarrierSignals` slots per context regardless of the team + // size. Single-domain runs take the NVLink barrier and consume none. + const int barrier_signal_slots = + scaleout_active ? elastic::gin_alloc::kNumReservedBarrierSignals : 0; + EP_HOST_ASSERT(gin_config.gin_indexed_signals_cnt >= barrier_signal_slots and + "GIN indexed-signal budget cannot host the barrier's counting signal; " + "reduce num_allocated_qps to raise the per-context signal count"); this->num_allocated_qps = gin_config.gin_context_cnt; @@ -204,6 +209,11 @@ NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm, } is_scaleup_nvlink = num_scaleup_ranks == num_nvl_ranks; + EP_HOST_ASSERT((is_scaleup_nvlink or num_scaleup_ranks <= 1 or num_scaleout_ranks <= 1) and + "A GIN scale-up barrier and a GIN scale-out barrier would share the " + "reserved barrier signal id; allocate a second reserved id before " + "allowing this combination"); + // Create symmetric memory // num_bytes = GPU + CPU, derive GPU portion this->symmetric_memory = symmetric::alloc( diff --git a/deep_ep/include/deep_ep/common/comm.cuh b/deep_ep/include/deep_ep/common/comm.cuh index 5cbb9e08a..fc2bc6534 100644 --- a/deep_ep/include/deep_ep/common/comm.cuh +++ b/deep_ep/include/deep_ep/common/comm.cuh @@ -15,6 +15,7 @@ #include #include +#include #include #include #include @@ -117,8 +118,10 @@ __device__ __forceinline__ std::pair get_qp_mod template __device__ __forceinline__ int get_qp_signal_id( const int& sm_idx, const int& channel_in_sm_idx) { - return channel_to_signal_id( - sm_idx, channel_in_sm_idx); + // Shift past the barrier's reserved ids; `channel_to_signal_id` itself stays 0-based. + return elastic::gin_alloc::data_signal_id( + channel_to_signal_id( + sm_idx, channel_in_sm_idx)); } // Per-part indexed-signal id: kNumParts contiguous ids under the channel's base id, one @@ -128,8 +131,11 @@ template __device__ __forceinline__ int get_per_part_signal_id( const int& sm_idx, const int& channel_in_sm_idx, const int& part_idx) { - return get_qp_signal_id( - sm_idx, channel_in_sm_idx) * kNumParts + part_idx; + // The reservation is added after the per-part multiply: going through + // `get_qp_signal_id` would scale the offset by `kNumParts`. + return elastic::gin_alloc::data_signal_id( + channel_to_signal_id( + sm_idx, channel_in_sm_idx) * kNumParts + part_idx); } template @@ -208,36 +214,82 @@ __forceinline__ __device__ void gin_barrier_wo_local_sync( ncclTeamWorld(nccl_dev_comm) : ncclTeamRail(nccl_dev_comm); const ncclGin gin(nccl_dev_comm, 0, NCCL_GIN_RESOURCE_SHARING_CTA); - // Compact signal indexing: (kNumRanks - 1) signal slots per rank. Sender rank_idx - // writes to every peer i at the slot that identifies *itself* in the peer's - // enumeration: - // sig = (rank_idx < i) ? rank_idx : (rank_idx - 1) - // So on receiver R, each of the (kNumRanks - 1) slots gets exactly +1 from a - // distinct sender, and the wait side just iterates all slots looking for one - // increment per slot. - for (int i = thread_idx; i < kNumRanks; i += kNumThreads) { - if (i == rank_idx) continue; - const auto sig = static_cast((rank_idx < i) ? rank_idx : (rank_idx - 1)); - gin.signal(team, i, ncclGin_SignalInc{sig}); - } - - for (int i = thread_idx; i < kNumRanks - 1; i += kNumThreads) { - const auto signal_idx = static_cast(i); - const auto shadow_ptr = gin.getSignalShadowPtr(signal_idx); - const auto target = ++(*shadow_ptr); - - // TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check. - timeout_while([=](const bool& is_last_check) { - const auto signal = gin.readSignal(signal_idx, 64, cuda::memory_order_acquire); - if (signal >= target) - return true; - - if (is_last_check) { - printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, thread: %d, " - "signal: %lu, target: %lu\n", kTag, scaleout_rank_idx, scaleup_rank_idx, thread_idx, signal, target); - } - return false; - }); + // The two team instantiations run different protocols. World keeps the per-peer + // barrier: its call sites pass `kFlushStores = true` and rely on per-peer arrival + // for release semantics. Rail is synchronisation-only (`kFlushStores = false` at + // every call site) and uses a counting barrier on one reserved signal id, so its + // per-context signal cost is constant in the team size. + if constexpr (std::is_same_v) { + // Compact signal indexing: (kNumRanks - 1) signal slots per rank. Sender rank_idx + // writes to every peer i at the slot that identifies *itself* in the peer's + // enumeration: + // sig = (rank_idx < i) ? rank_idx : (rank_idx - 1) + // So on receiver R, each of the (kNumRanks - 1) slots gets exactly +1 from a + // distinct sender, and the wait side just iterates all slots looking for one + // increment per slot. + for (int i = thread_idx; i < kNumRanks; i += kNumThreads) { + if (i == rank_idx) continue; + const auto sig = static_cast((rank_idx < i) ? rank_idx : (rank_idx - 1)); + gin.signal(team, i, ncclGin_SignalInc{sig}); + } + + for (int i = thread_idx; i < kNumRanks - 1; i += kNumThreads) { + const auto signal_idx = static_cast(i); + const auto shadow_ptr = gin.getSignalShadowPtr(signal_idx); + const auto target = ++(*shadow_ptr); + + // TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check. + timeout_while([=](const bool& is_last_check) { + const auto signal = gin.readSignal(signal_idx, 64, cuda::memory_order_acquire); + if (signal >= target) + return true; + + if (is_last_check) { + printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, thread: %d, " + "signal: %lu, target: %lu\n", kTag, scaleout_rank_idx, scaleup_rank_idx, thread_idx, signal, target); + } + return false; + }); + } + } else { + // Counting barrier: every sender adds 1 to the same reserved signal id on + // every peer, and the single waiter advances its shadow by the expected + // (kNumRanks - 1) increments. `SignalAdd{.., 1}` matches the accumulation + // pattern of the unordered data path (`hybrid_combine_unordered.cuh`). + constexpr auto kBarrierSignal = + static_cast(elastic::gin_alloc::kBarrierSignalId); + for (int i = thread_idx; i < kNumRanks; i += kNumThreads) { + if (i == rank_idx) continue; + gin.signal(team, i, ncclGin_SignalAdd{kBarrierSignal, 1ull}); + } + + // First sync: every thread has issued its sends before the single waiter + // starts polling. Second sync: no thread leaves the barrier before the wait + // completes. + __syncthreads(); + + if (thread_idx == 0) { + const auto shadow_ptr = gin.getSignalShadowPtr(kBarrierSignal); + const auto target = (*shadow_ptr += static_cast(kNumRanks - 1)); + + // TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check. + timeout_while([=](const bool& is_last_check) { + const auto signal = gin.readSignal(kBarrierSignal, 64, cuda::memory_order_acquire); + if (signal >= target) + return true; + + if (is_last_check) { + // A counting slot cannot identify the stalled peer; report the shortfall. + printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, " + "signal_id: %d, observed: %lu, target: %lu, missing: %lu of %d\n", + kTag, scaleout_rank_idx, scaleup_rank_idx, + static_cast(kBarrierSignal), signal, target, + target - signal, kNumRanks - 1); + } + return false; + }); + } + __syncthreads(); } } } @@ -291,6 +343,14 @@ __forceinline__ __device__ void gpu_barrier(const handle::NCCLGin& gin, EP_STATIC_ASSERT(not kFlushStores, "No data to be flushed"); } + // World's per-peer slots and Rail's counting slot overlap in the same + // (context, signal) space, so the two GIN barriers must never be live concurrently. + // `NCCLSymmetricMemoryContext` enforces the same condition at init. + EP_STATIC_ASSERT(kIsScaleupNVLink or kNumScaleupRanks <= 1 or kNumScaleoutRanks <= 1, + "A GIN scale-up barrier and a GIN scale-out barrier would share the " + "reserved barrier signal id; allocate a second reserved id before " + "allowing this combination"); + do_scaleout &= kNumScaleoutRanks > 1; do_scaleup &= kNumScaleupRanks > 1; if (do_scaleup and do_scaleout) { diff --git a/deep_ep/include/deep_ep/common/gin_resource_alloc.cuh b/deep_ep/include/deep_ep/common/gin_resource_alloc.cuh index 7b2cf7ebf..79b1881a5 100644 --- a/deep_ep/include/deep_ep/common/gin_resource_alloc.cuh +++ b/deep_ep/include/deep_ep/common/gin_resource_alloc.cuh @@ -64,20 +64,32 @@ __forceinline__ __device__ __host__ constexpr GinResourceConfig make_gin_resourc return GinResourceConfig{gin_indexed_signals_for(gin_context_cnt), gin_context_cnt}; } -// Each ScaleOut warp (== channel) needs its own dedicated indexed signal id, so the total -// signal budget (ctx * signals/ctx) must cover the worst-case warp count for EVERY legal -// context count, not just the default. The tightest points are ctx = 13 and ctx = 17, both at -// 221 against the 220-warp ceiling -- one signal of slack. Do not raise `kMaxSM` / -// `kMaxWarpsPerSM`, widen the context range, or lower `kTotalQPBudget` without re-checking. -__forceinline__ __host__ constexpr bool all_gin_context_counts_cover_warps() { - for (int ctx = kMinGinContextCnt; ctx <= kMaxGinContextCnt; ++ ctx) - if (ctx * gin_indexed_signals_for(ctx) < kMaxScaleoutWarps) - return false; - return true; +// Indexed-signal ids reserved for the rail counting barrier (`gin_barrier_wo_local_sync`, +// comm.cuh), taken off the bottom of every context's id space so no data channel can +// produce the barrier's id. Plain `int`, not an NCCL signal type: this header stays +// NCCL-free and host-compilable so the invariants below can be `static_assert`s. +static constexpr int kNumReservedBarrierSignals = 1; +static constexpr int kBarrierSignalId = 0; +static_assert(kBarrierSignalId >= 0 and kBarrierSignalId < kNumReservedBarrierSignals, + "the barrier's signal id must lie inside the reserved range"); + +// The single place the reservation offset is applied; both id derivations in `comm.cuh` +// route through it. +__forceinline__ __device__ __host__ constexpr int data_signal_id(int raw_offset) { + return kNumReservedBarrierSignals + raw_offset; +} + +// Signals the data path may use: the provisioned per-context count minus the barrier +// reservation. +__forceinline__ __device__ __host__ constexpr int usable_signal_budget(const GinResourceConfig& cfg) { + return cfg.gin_indexed_signals_cnt > kNumReservedBarrierSignals + ? cfg.gin_indexed_signals_cnt - kNumReservedBarrierSignals : 0; } -static_assert(all_gin_context_counts_cover_warps(), - "GIN layout cannot give each ScaleOut warp a dedicated signal id " - "for every legal context count"); + +// Every legal context count must remain serviceable at the worst-case launch (`kMaxSM` +// SMs x `kMaxWarpsPerSM` ScaleOut warps): `compute_part_allocation` must return at least +// one channel per SM and at least one part. Defined below, after the math it checks. +__forceinline__ __host__ constexpr bool all_gin_context_counts_are_serviceable(); // Preferred (and workspace-sizing) maximum for per-part signalling. static constexpr int kMaxParts = 4; @@ -120,11 +132,13 @@ __forceinline__ __device__ __host__ constexpr int channels_per_context( } // Per-part signal allocation: pick the largest num_parts (up to kMaxParts) that fits -// channels_per_context(...) * num_parts <= gin_indexed_signals_cnt +// channels_per_context(...) * num_parts <= usable_signal_budget(cfg) // at the requested channels/SM, then reduce channels_per_sm until the budget holds. -__forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_allocation( +// `_raw` is the diagnostics-free math so `static_assert`s can evaluate it; +// `compute_part_allocation` below adds the host-side warning and check. +__forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_allocation_raw( const GinResourceConfig& cfg, int num_sms, int num_available_qps, int num_channels_per_sm) { - const int gin_signals = cfg.gin_indexed_signals_cnt; + const int gin_signals = usable_signal_budget(cfg); const int channels_per_ctx = channels_per_context(num_sms, num_available_qps, num_channels_per_sm); const int budget_parts = gin_signals / (channels_per_ctx > 1 ? channels_per_ctx : 1); GinPartAllocation alloc{}; @@ -135,20 +149,50 @@ __forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_all static_cast(channels_per_context(num_sms, num_available_qps, alloc.num_channels_per_sm)) * alloc.num_parts > gin_signals) --alloc.num_channels_per_sm; + return alloc; +} + +// True when the allocation fits the usable budget (the reduction loop converged). +__forceinline__ __device__ __host__ constexpr bool part_allocation_fits( + const GinResourceConfig& cfg, int num_sms, int num_available_qps, const GinPartAllocation& alloc) { + return static_cast(channels_per_context(num_sms, num_available_qps, + alloc.num_channels_per_sm)) * alloc.num_parts + <= usable_signal_budget(cfg); +} + +__forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_allocation( + const GinResourceConfig& cfg, int num_sms, int num_available_qps, int num_channels_per_sm) { + const GinPartAllocation alloc = + compute_part_allocation_raw(cfg, num_sms, num_available_qps, num_channels_per_sm); #ifndef __CUDA_ARCH__ if (alloc.num_channels_per_sm < num_channels_per_sm) printf("[WARN] DeepEP GIN signal budget reduced the number of channels per SM " "from %d to %d\n", num_channels_per_sm, alloc.num_channels_per_sm); -#endif -#ifndef __CUDA_ARCH__ - EP_HOST_ASSERT(static_cast(channels_per_context(num_sms, num_available_qps, - alloc.num_channels_per_sm)) * alloc.num_parts <= gin_signals and + EP_HOST_ASSERT(part_allocation_fits(cfg, num_sms, num_available_qps, alloc) and "GIN signal budget cannot host even 1 part-signal per channel " "at 1 channel/SM. Reduce --num-sms or num_allocated_qps."); #endif return alloc; } +// The invariant declared above, now that the math it checks is in scope. +__forceinline__ __host__ constexpr bool all_gin_context_counts_are_serviceable() { + for (int ctx = kMinGinContextCnt; ctx <= kMaxGinContextCnt; ++ ctx) { + // The notify warp owns QP 0, so only `ctx - 1` contexts carry data channels. + const int avail = ctx > 1 ? ctx - 1 : 1; + const auto cfg = make_gin_resources(ctx); + const auto alloc = compute_part_allocation_raw(cfg, kMaxSM, avail, kMaxWarpsPerSM); + if (alloc.num_channels_per_sm < 1 or alloc.num_parts < 1) + return false; + if (not part_allocation_fits(cfg, kMaxSM, avail, alloc)) + return false; + } + return true; +} +static_assert(all_gin_context_counts_are_serviceable(), + "some legal GIN context count cannot service the worst-case launch " + "(kMaxSM x kMaxWarpsPerSM) once the barrier reservation is taken out"); + // Kernel-side entry points: derive the per-channel part count (and verify the launched // channel count) as compile-time constants from the provisioned indexed-signal budget. __forceinline__ __device__ __host__ constexpr int constexpr_num_parts(