[Do not merge] Experiment: iteration-flag combine gate (EP_COMBINE_GATE)

csrc/elastic/buffer.hpp

@@ -43,6 +43,8 @@ class ElasticBuffer {
 
     // Workspace
     // NOTES: for all workspace, we must keep them as zeros
+    // (exception: the combine gate flags are monotonic counters — zeroed once at
+    // construction, never reset; see `get_combine_gate_flag_ptr`)
     void *workspace;
     void *host_workspace, *mapped_host_workspace;
     std::shared_ptr<layout::WorkspaceLayout> workspace_layout_wo_expert;
@@ -58,6 +60,8 @@ class ElasticBuffer {
 
     mutable int dispatch_iteration = 0;
 
+    mutable uint32_t combine_iteration = 0;
+
     // Whether to prefer overlapping communication with compute (use more SMs and channels if false)
     bool prefer_overlap_with_compute;
 
@@ -1409,6 +1413,7 @@ class ElasticBuffer {
 
         // Push data into remote buffers
         // NOTES: we don't use `num_hidden_bytes` due to enable later quantization possibility
+        ++ combine_iteration;
         const auto reduce_buffer = launch_combine(
             x.data_ptr(),
             topk_weights.has_value() ? topk_weights->data_ptr() : nullptr,
@@ -1420,6 +1425,7 @@ class ElasticBuffer {
             nccl_context->dev_comm, nccl_context->window,
             buffer, workspace,
             num_reduced_tokens, num_combined_tokens,
+            combine_iteration,
             num_max_tokens_per_rank,
             hidden, num_experts, num_topk,
             num_qps, num_gpu_timeout_cycles,

csrc/kernels/backend/nccl.cu

@@ -130,9 +130,34 @@ 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 instantiation of `gin_barrier_wo_local_sync` is now a counting barrier:
+        // every peer adds to a single signal id and the waiter advances its shadow by the
+        // expected arrival count, so the rail barrier costs ONE indexed-signal slot whatever
+        // the team size. Single-domain runs take the NVLink barrier and consume none.
+        //
+        // The change is scoped to Rail on purpose. This arm is the only one with the
+        // ceiling, and the Rail barrier is the only one that never carries release
+        // semantics -- every rail call site passes `kFlushStores = false`. The World
+        // instantiation, reached only from the direct / ordered arm below (which asks for
+        // `num_ranks + 2 * 2` signals and therefore has no ceiling), keeps the per-peer
+        // barrier: `dispatch.cuh` and `combine.cuh` use it with `kFlushStores = true` to
+        // "ensure data arrival", and an anonymous counter cannot establish that N distinct
+        // peers arrived -- a peer a round ahead can supply two of the increments.
+        //
+        // This is what removes the scale-out ceiling. The previous check scaled with the
+        // team size against a per-context budget fixed at (kTotalQPBudget - 2c)/c, and so
+        // refused to initialize past 22 NVLink domains at the default context count --
+        // measured on p6-b200: 22 domains complete, 23 refuse. The budget's only remaining
+        // TEAM-SIZE-DEPENDENT consumer is the data path, whose requirement
+        // (ceil(channels/qp) * num_parts) does not grow with the team and is enforced by
+        // `compute_part_allocation`. The barrier still consumes a fixed
+        // `kNumReservedBarrierSignals` on every context -- that reservation is what keeps the
+        // data path from ever producing the barrier's id.
+        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");
 
         if (scaleout_active)
             this->num_allocated_qps = gin_config.gin_context_cnt;
@@ -198,6 +223,18 @@ NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm,
     }
     is_scaleup_nvlink = num_scaleup_ranks == num_nvl_ranks;
 
+    // The two device barriers overlap in the same (context, signal) space -- World's
+    // per-peer slots start at id 0 and Rail's counting slot IS id 0 -- so a GIN scale-up
+    // barrier and a GIN scale-out barrier must never be live concurrently. See
+    // `gpu_barrier` in `comm.cuh` for why, and for the same condition as a static assert.
+    // The two branches above make this unreachable, but the kernels are JIT-generated from
+    // exactly these values, so without a host gate a violation would surface as a compile
+    // exception on first launch rather than here at init.
+    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(

csrc/kernels/elastic/combine.hpp

@@ -32,6 +32,9 @@ class CombineRuntime final : public jit::LaunchRuntime<CombineRuntime> {
         // Use the ordered (upstream) hybrid kernel instead of the unordered one.
         // Resolved once from `EP_HYBRID_KERNEL`; only affects the hybrid path.
         bool use_ordered_kernel;
+        // Gate the unordered combine's return on iteration flags instead of
+        // counting signals. Resolved once from `EP_COMBINE_GATE`.
+        bool use_flag_gate;
         int num_scaleup_warps, num_forward_warps;
         int num_scaleout_ranks, num_scaleup_ranks;
         int hidden;
@@ -56,6 +59,7 @@ class CombineRuntime final : public jit::LaunchRuntime<CombineRuntime> {
         int scaleout_rank_idx, scaleup_rank_idx;
         int num_reduced_tokens;
         int num_combined_tokens;
+        uint32_t combine_iteration;
 
         jit::LaunchArgs launch_args;
     };
@@ -77,9 +81,13 @@ class CombineRuntime final : public jit::LaunchRuntime<CombineRuntime> {
                                     args.num_qps, args.num_timeout_cycles);
         } else {
             header_name = args.use_ordered_kernel ? "hybrid_combine" : "hybrid_combine_unordered";
-            func_name = fmt::format("{}<{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
+            // Only the unordered kernel takes the `kUseFlagGate` template argument
+            const std::string flag_gate_arg =
+                args.use_ordered_kernel ? "" : fmt::format("{}, ", args.use_flag_gate);
+            func_name = fmt::format("{}<{}, {}, {}{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}>",
                                     args.use_ordered_kernel ? "hybrid_combine_impl" : "hybrid_unordered_combine_impl",
                                     args.use_expanded_layout, args.allow_multiple_reduction,
+                                    flag_gate_arg,
                                     args.launch_args.grid_dim.first,
                                     args.num_scaleup_warps, args.num_forward_warps,
                                     args.num_scaleout_ranks, args.num_scaleup_ranks,
@@ -133,7 +141,8 @@ static void __instantiate_kernel() {{
                                                      args.buffer, args.workspace,
                                                      args.scaleout_rank_idx, args.scaleup_rank_idx,
                                                      args.num_reduced_tokens,
-                                                     args.num_combined_tokens));
+                                                     args.num_combined_tokens,
+                                                     args.combine_iteration));
         }
     }
 };
@@ -153,6 +162,7 @@ static void* launch_combine(void* x,
                             const jit::NoRefPtr& nccl_dev_comm, const ncclWindow_t& nccl_window,
                             void* buffer, void* workspace,
                             const int& num_reduced_tokens, const int& num_combined_tokens,
+                            const uint32_t& combine_iteration,
                             const int& num_max_tokens_per_rank,
                             const int& hidden,
                             const int& num_experts, const int& num_topk,
@@ -170,6 +180,9 @@ static void* launch_combine(void* x,
 
     // Decide warps
     const bool use_ordered_kernel = use_ordered_hybrid_kernel();
+    const bool use_flag_gate = use_flag_combine_gate();
+    EP_HOST_ASSERT((not (use_flag_gate and use_ordered_kernel)) and
+                   "EP_COMBINE_GATE=flag requires the unordered hybrid kernel");
     int num_scaleup_warps = 0, num_forward_warps = 0;
     if (num_scaleout_ranks > 1) {
         EP_HOST_ASSERT(num_channels % num_sms == 0 and
@@ -209,6 +222,7 @@ static void* launch_combine(void* x,
         .use_expanded_layout = use_expanded_layout,
         .allow_multiple_reduction = allow_multiple_reduction,
         .use_ordered_kernel = use_ordered_kernel,
+        .use_flag_gate = use_flag_gate,
         .num_scaleup_warps = num_scaleup_warps, .num_forward_warps = num_forward_warps,
         .num_scaleout_ranks = num_scaleout_ranks, .num_scaleup_ranks = num_scaleup_ranks,
         .hidden = hidden,
@@ -228,6 +242,7 @@ static void* launch_combine(void* x,
         .scaleout_rank_idx = scaleout_rank_idx, .scaleup_rank_idx = scaleup_rank_idx,
         .num_reduced_tokens = num_reduced_tokens,
         .num_combined_tokens = num_combined_tokens,
+        .combine_iteration = combine_iteration,
         // NOTES: make cluster dim 2 to overlap with clustered computation kernels
         .launch_args = jit::LaunchArgs(num_sms, num_threads, num_smem_bytes, 2 - (num_sms % 2), true)
     };

csrc/kernels/elastic/kernel_select.hpp

@@ -34,6 +34,33 @@ namespace deep_ep::elastic {
 //
 // The value is read once per process. It changes the JIT-generated source (header and
 // kernel names differ), so the JIT cache distinguishes the two variants automatically.
+// `EP_COMBINE_GATE` selects how the unordered combine's return gate synchronizes:
+//
+//   "signal" (default) — counting GIN indexed signals (the existing path).
+//
+//   "flag"             — a 4-byte iteration number written by putValue into a
+//                        symmetric workspace slot; the receiver polls it. The
+//                        sender flushes its data puts before writing the flag.
+//
+// Read once per process; JIT-cached per variant. Only the unordered kernel has
+// this gate, so "flag" rejects EP_HYBRID_KERNEL=ordered.
+static bool use_flag_combine_gate() {
+    static const bool flag = [] {
+        const auto value = get_env<std::string>("EP_COMBINE_GATE");
+        bool result = false;
+        if (value.empty() or value == "signal") {
+            result = false;
+        } else if (value == "flag") {
+            result = true;
+        } else {
+            EP_HOST_ASSERT(false and "EP_COMBINE_GATE must be `signal` or `flag`");
+        }
+        printf("DeepEP combine gate selection: %s\n", result ? "flag" : "signal");
+        return result;
+    }();
+    return flag;
+}
+
 static bool use_ordered_hybrid_kernel() {
     static const bool ordered = [] {
         const auto value = get_env<std::string>("EP_HYBRID_KERNEL");