comm: lift the GIN scale-out ceiling by making the rail barrier a counting barrier

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(

deep_ep/include/deep_ep/common/comm.cuh

@@ -15,6 +15,7 @@
 #include <nccl.h>
 #include <nccl_device.h>
 
+#include <deep_ep/common/gin_resource_alloc.cuh>
 #include <deep_ep/common/handle.cuh>
 #include <deep_ep/common/ptx.cuh>
 #include <deep_ep/common/layout.cuh>
@@ -117,8 +118,10 @@ __device__ __forceinline__ std::pair<int, ncclGinResourceSharingMode> get_qp_mod
 template <int kNumSMs, int kNumQPs, int kNumChannelsPerSM, bool kWithNotifyWarps = false>
 __device__ __forceinline__ int get_qp_signal_id(
     const int& sm_idx, const int& channel_in_sm_idx) {
-    return channel_to_signal_id<kNumSMs, kNumQPs, kNumChannelsPerSM, kWithNotifyWarps>(
-        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<kNumSMs, kNumQPs, kNumChannelsPerSM, kWithNotifyWarps>(
+            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 <int kNumSMs, int kNumQPs, int kNumChannelsPerSM, int kNumParts,
           bool kWithNotifyWarps = false>
 __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<kNumSMs, kNumQPs, kNumChannelsPerSM, kWithNotifyWarps>(
-               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<kNumSMs, kNumQPs, kNumChannelsPerSM, kWithNotifyWarps>(
+            sm_idx, channel_in_sm_idx) * kNumParts + part_idx);
 }
 
 template <int kNumRanks, int kNumSMs, int kNumThreads, int64_t kNumTimeoutCycles, int kTag = kDeviceBarrierTag>
@@ -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<ncclGinSignal_t>((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<ncclGinSignal_t>(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<kNumTimeoutCycles>([=](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<team_t, ncclTeamTagWorld>) {
+            // 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<ncclGinSignal_t>((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<ncclGinSignal_t>(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<kNumTimeoutCycles>([=](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<ncclGinSignal_t>(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<uint64_t>(kNumRanks - 1));
+
+                // TODO(NCCL): Using the official NCCL wait signal API, after they added timeout check.
+                timeout_while<kNumTimeoutCycles>([=](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<int>(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) {

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<long long>(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<long long>(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<long long>(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(