From e0f110a3771293fe493bc44a84e626b1192c1383 Mon Sep 17 00:00:00 2001 From: Xuan Jiang Date: Wed, 29 Jul 2026 04:48:04 +0000 Subject: [PATCH 1/3] envs: probe RDMA link rate via sysfs, survive probe failure get_rdma_gbs() only knew how to ask ibstat for a CA named EP_NIC_NAME (default mlx5_0). EFA devices expose no umad CA, so on EFA hosts the probe returned 0 and get_theoretical_num_sms() divided by it -- any multi-node run without an explicit --num-sms crashed with ZeroDivisionError. Read /sys/class/infiniband//ports/*/rate first, which works for every verbs provider, and fall back to ibstat for setups whose rate only shows there. When EP_NIC_NAME is unset and the default device is absent, pick the fastest device under /sys/class/infiniband instead of failing; an explicitly named NIC still fails loudly rather than guessing. --- deep_ep/utils/envs.py | 55 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/deep_ep/utils/envs.py b/deep_ep/utils/envs.py index f6e34d988..af4183fcb 100644 --- a/deep_ep/utils/envs.py +++ b/deep_ep/utils/envs.py @@ -1,3 +1,14 @@ +# MIT License +# +# Copyright (c) 2025 DeepSeek +# Changes and additions copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + import functools import inspect import os @@ -242,17 +253,61 @@ def check_fast_rdma_atomic_support(nic_name: str = _DEFAULT_NIC_NAME) -> bool: return False +def _get_sysfs_rdma_gbs(nic_name: str) -> float: + """ + Read one RDMA device's link rate from sysfs (`/sys/class/infiniband//ports/*/rate`). + Works for any verbs provider (mlx5, EFA's `rdmap*`, ...) without external tools. + + Arguments: + nic_name: the NIC device name. + + Returns: + gbs: the device's link rate in GB/s (0 if the device or its rate is unavailable). + """ + rate = 0 + ports_dir = os.path.join('/sys/class/infiniband', nic_name, 'ports') + try: + for port in os.listdir(ports_dir): + with open(os.path.join(ports_dir, port, 'rate')) as f: + match = re.match(r'\s*(\d+)\s*Gb/sec', f.read()) + if match: + rate = max(rate, int(match.group(1))) + except OSError: + pass + return rate / 8 + + @functools.lru_cache() def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float: """ Get the RDMA bandwidth in GB/s, cached. + Probes sysfs first, which covers any verbs provider; `ibstat` is kept as a fallback but + cannot see providers without a umad interface (e.g. EFA). When `EP_NIC_NAME` is not set + and the default device does not exist (EFA hosts have no `mlx5_0`), the fastest device + under `/sys/class/infiniband` is used instead. + Arguments: nic_name: the NIC device name. Returns: gbs: the RDMA bandwidth in GB/s (0 if detection fails). """ + gbs = _get_sysfs_rdma_gbs(nic_name) + if gbs > 0: + return gbs + + # The un-overridden default may simply not exist on this fabric; an explicitly named NIC + # must not fall back silently + if 'EP_NIC_NAME' not in os.environ: + try: + devices = sorted(os.listdir('/sys/class/infiniband')) + except OSError: + devices = [] + gbs = max((_get_sysfs_rdma_gbs(device) for device in devices), default=0) + if gbs > 0: + return gbs + # noinspection PyBroadException try: result = subprocess.run(['ibstat'], capture_output=True, text=True, check=True) From 02efc268a37802fc00812ede8f5ad7f535ceea0e Mon Sep 17 00:00:00 2001 From: Aviv Benchorin Date: Fri, 21 Aug 2026 17:54:06 +0000 Subject: [PATCH 2/3] tests: bounded pressure loops and teardown barrier for test_ep Add --pressure-iterations to bound the pressure-test loop (upstream's --do-pressure-test runs int(1e9) seeds, i.e. until killed; 0 keeps that behavior), with argument validation. Add a barrier before dist.destroy_process_group() so a fast rank cannot tear down the TCPStore while slower ranks are still in destroy. Signed-off-by: Xuan Jiang --- tests/elastic/test_ep.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/tests/elastic/test_ep.py b/tests/elastic/test_ep.py index 7344af5f2..456e20340 100644 --- a/tests/elastic/test_ep.py +++ b/tests/elastic/test_ep.py @@ -1,3 +1,14 @@ +# MIT License +# +# Copyright (c) 2025 DeepSeek +# Changes and additions copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + import argparse import os import torch @@ -545,7 +556,12 @@ def construct_elastic_buffer(): test_dispatch_combine(buffer, args) # Pressure tests - for seed in range(int(1e9) if args.do_pressure_test else 0): + if args.do_pressure_test: + pressure_iteration_count = args.pressure_iterations if args.pressure_iterations != 0 else int(1e9) + else: + pressure_iteration_count = 0 + + for seed in range(pressure_iteration_count): if not args.reuse_elastic_buffer: # Recreate elastic buffer buffer.destroy() @@ -558,6 +574,7 @@ def construct_elastic_buffer(): # Destroy the runtime and communication group buffer.destroy() + dist.barrier() dist.destroy_process_group() @@ -591,6 +608,12 @@ def construct_elastic_buffer(): parser.add_argument('--skip-check', action='store_true', help='Whether to skip correctness checks') parser.add_argument('--skip-perf-test', action='store_true', help='Whether to skip performance tests') parser.add_argument('--do-pressure-test', action='store_true', help='Whether to do pressure test') + parser.add_argument( + '--pressure-iterations', + type=int, + default=0, + help='Number of pressure-loop seeds; 0 represents the default unbounded value of 1e9 seeds', + ) parser.add_argument('--reuse-elastic-buffer', action='store_true', help='Whether to reuse elastic buffer for each test') parser.add_argument('--test-first-only', action='store_true', help='Only test the first case') parser.add_argument('--unbalanced-ratio', type=float, default=1.0, help='The MoE unbalanced ratio') @@ -599,6 +622,10 @@ def construct_elastic_buffer(): parser.add_argument('--dump-profile-traces', type=str, default='', help='Dump profiling trace JSONs') parser.add_argument('--ignore-local-traffic', action='store_true', help='Whether to ignore local traffic during bandwidth calculation') args = parser.parse_args() + if args.pressure_iterations < 0: + parser.error("--pressure-iterations must be non-negative") + if args.pressure_iterations and not args.do_pressure_test: + parser.error("--pressure-iterations requires --do-pressure-test") # Create dump trace directories if args.dump_profile_traces: From 2542d9641f2ec280213e875feb04be7862dda57c Mon Sep 17 00:00:00 2001 From: Keita Watanabe Date: Sun, 23 Aug 2026 18:41:17 +0000 Subject: [PATCH 3/3] comm: lift the GIN scale-out ceiling by making the rail barrier a counting barrier On EFA the shipped default refuses to initialize past 22 NVLink domains. Measured on p6-b200.48xlarge: 22 domains complete, 23 refuse at `nccl.cu` with gin_config.gin_indexed_signals_cnt >= (num_rdma_ranks - 1) and "GIN indexed-signal budget cannot give each peer rail team a dedicated signal..." The assert is correct and load-bearing, not redundant: `gin_barrier_wo_local_sync` really did need one indexed-signal slot per peer, so the rail barrier's demand grew with the team while the per-context budget did not. A GIN signal id is not free on EFA -- one id is one `gdaki_sc_endpoint`, a complete QP and CQ -- so the budget is `gin_indexed_signals_for(c) = (kTotalQPBudget - 2c)/c`, which is 21 at the shipped 11 contexts. 21 slots is 22 domains, and the 23rd is refused. Raising the budget is not available: it is a NIC resource, independently measured at 512 completion counters per NIC (256 QPs at two counters each), exactly what `kTotalQPBudget = 256` already assumes. The fix removes the team-size term instead. A barrier is a counting predicate and does not need to distinguish senders, so every peer now adds 1 to the SAME signal id and the waiter advances its shadow by `kNumRanks - 1`. That costs ONE slot whatever the team size. The pattern is the one the unordered data path already relies on (`hybrid_combine_unordered.cuh`'s `num_expected_arrivals` wait), and `SignalAdd{.., 1}` matches that precedent. This is scoped to the RAIL instantiation, deliberately, because the two teams have opposite requirements: * Only rail has a ceiling. The unordered-hybrid arm of `NCCLSymmetricMemoryContext` asks for `gin_indexed_signals_cnt`, the per-context budget above. The direct / ordered arm asks for `num_ranks + 2 * 2`, commented "Customized RDMA barrier needs extra signals" -- the world barrier's per-peer slots are already budgeted there and scale with the team. * Only world is used as a release barrier. Every rail call site passes `kFlushStores = false` (`hybrid_{dispatch,combine}{,_unordered}.cuh`, `barrier.cuh`), so the rail barrier never even issues the QP flush; it is a pure synchronisation point. The hybrid kernels' two "ensure data arrival" barriers pass `do_scaleout = false` and run over NVLink. By contrast `dispatch.cuh` and `combine.cuh` use the world path with `kFlushStores = true` and then read what peers wrote. * A counting barrier cannot carry release. Its counter is anonymous, so a peer one round ahead can supply an increment standing in for a delayed current-round arrival: the count reaches its target without every distinct peer having arrived. "Everyone arrived" survives; "every peer's prior writes are visible to me" does not. Signal strength does not repair this -- strong signals order a sender's own prior puts, they never say which sender incremented. Identity is the missing half and only per-peer slots have it, so this would hold on InfiniBand too. So the world body is left unchanged apart from the four-space re-indent that the new `if constexpr` block forces around it, and can be reviewed as untouched. To check that mechanically, strip comments and whitespace from the block in both revisions and compare -- 746 characters of code, identical. A second, live defect had to be fixed for the counting barrier to be safe at all. The barrier hardcoded signal id 0 and nothing stopped the data path from producing that same id. In `cached_mode` the unordered hybrid kernels run with zero notify warps, which makes `kQPStartIdx = 0` and puts data channels on context 0 -- the barrier's own context -- where 0-based `channel_to_signal_id` yields id 0 for channel 0 part 0. NCCL addresses a shadow by (context, signal) alone, so that was one shared counter with each side inflating the other's arrival count. `kNumReservedBarrierSignals` now takes one id off the bottom of every context's id space and `data_signal_id()` is the single place the offset is applied, so both derivations in `comm.cuh` shift together. Note the offset is added AFTER the per-part multiply; going through `get_qp_signal_id` would scale it by `kNumParts` and burn ids. Two invariants were repaired alongside. `all_gin_context_counts_cover_warps()` compared a cross-context TOTAL against a warp count, which says nothing about the most crowded context; it is replaced by a serviceability check over every legal context count, verified non-vacuous (forcing the reservation to 200 fails the build). `compute_part_allocation` is split into a pure `_raw` and a diagnosing wrapper so the invariant can be a `static_assert` -- a `printf` or a throwing assert reached during constant evaluation makes the expression non-constant. `gpu_barrier` gains a static assert, mirrored by a host assert in `NCCLSymmetricMemoryContext`, that a GIN scale-up and a GIN scale-out barrier are never live concurrently: their id ranges overlap, since world's per-peer slots start at 0 and rail's counting slot IS 0. No instantiation can violate it today; the host assert is what fails early, because these kernels are NVRTC-generated and a device static assert would surface as a JIT exception in production. Cost, measured with `compute_part_allocation` at the shipped context count: the reservation costs one part at 16-17 SMs (3 -> 2) and one channel per SM at 51-52 SMs (4 -> 3, with the host warning). No second id is needed, so the usable per-context budget stays at 20. Verified on 36 x p6-b200.48xlarge (8x B200, 8 EFA/node, efa kmod 3.3.0g), treatment and control built from the same base and each pod printing the md5 of the JIT headers it compiled: 23 domains, stock -> refuses at init, verbatim assert above 23 domains, this change -> completes, correctness checks pass 32 domains, this change -> completes (256 ranks) 22 domains, this change -> completes 2 domains, this change -> completes 3 domains, direct mode -> completes; exercises the world instantiation 3 domains, direct, stock -> completes; control for the row above Barrier cost at 22 domains is unchanged to within run-to-run spread, and the counting barrier is never slower than the per-peer one: 8.2% faster at 2 domains, 4.9% at 8, converging to no difference by 22. The single-counter hotspot does not materialise through 32 domains -- marginal cost 3.38 us/domain over 22->32 against the stock barrier's 3.25 us/domain over 16->22. End-to-end dispatch and combine throughput is unchanged. Remaining limitations and untested paths are tracked separately rather than in this change. --- csrc/kernels/backend/nccl.cu | 43 +++- deep_ep/include/deep_ep/common/comm.cuh | 220 +++++++++++++++--- .../deep_ep/common/gin_resource_alloc.cuh | 130 +++++++++-- 3 files changed, 336 insertions(+), 57 deletions(-) diff --git a/csrc/kernels/backend/nccl.cu b/csrc/kernels/backend/nccl.cu index d86b0a075..c72fb55c8 100644 --- a/csrc/kernels/backend/nccl.cu +++ b/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( diff --git a/deep_ep/include/deep_ep/common/comm.cuh b/deep_ep/include/deep_ep/common/comm.cuh index 443027441..f8fb726ca 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 @@ -111,8 +112,12 @@ __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); + // Shifted past the barrier's reserved ids. `channel_to_signal_id` stays 0-based (pure + // integer math); the reservation is applied here, at the boundary where a raw offset + // becomes a signal id the data path will actually signal. + 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 @@ -122,8 +127,12 @@ 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; + // NOTE: the reservation is added AFTER the multiply, and `channel_to_signal_id` is called + // directly rather than through `get_qp_signal_id`. Going through the latter would scale + // the offset by `kNumParts` and burn `kNumReservedBarrierSignals * kNumParts` ids. + return elastic::gin_alloc::data_signal_id( + channel_to_signal_id( + sm_idx, channel_in_sm_idx) * kNumParts + part_idx); } template @@ -202,36 +211,144 @@ __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 get DIFFERENT barrier protocols, and the split is + // deliberate. Rail gets a counting barrier; World keeps the per-peer barrier + // unchanged. Three facts force it: + // + // 1. Only Rail has a ceiling to remove. `NCCLSymmetricMemoryContext` + // (`csrc/kernels/backend/nccl.cu`) provisions the two paths from different + // arms: the unordered-hybrid arm asks for `gin_indexed_signals_cnt`, a + // per-context budget of `(kTotalQPBudget - 2c)/c` that does NOT grow with the + // team -- that is the budget a per-peer rail barrier overran at 23 domains. + // The direct / ordered arm asks for `num_ranks + 2 * 2`, explicitly commented + // "Customized RDMA barrier needs extra signals". The World barrier's per-peer + // slots are already budgeted there and scale with the team, so that path never + // had the ceiling and gains nothing from being converted. + // + // 2. Only World is used as a RELEASE barrier. Every Rail call site passes + // `kFlushStores = false` (`hybrid_{dispatch,combine}{,_unordered}.cuh`, + // `barrier.cuh`), so the Rail barrier never even issues the QP flush above -- + // it is a pure synchronisation point. The two "ensure data arrival" barriers + // in the hybrid kernels explicitly pass `do_scaleout = false` and run over + // NVLink. By contrast `dispatch.cuh:398` and `combine.cuh:240` take + // `kFlushStores = true` on this World path and then read what peers wrote + // (they trigger the copy-epilogue kernel immediately after). + // + // 3. A counting barrier cannot carry release. Its counter is anonymous, so a peer + // one round ahead can supply an increment that stands in for a delayed + // current-round arrival: the count reaches its target without every distinct + // peer having arrived. "Everyone arrived" survives that (a rank can only be a + // round ahead if it already saw everyone), but "every peer's prior writes are + // visible to me" does not. Signal STRENGTH cannot repair it either -- strong + // signals order a sender's own prior puts, they do not say which sender + // incremented. Identity is the missing half, and only per-peer slots have it. + // + // So: convert the path that has the ceiling and does not need release; leave the + // path that needs release and has no ceiling exactly as it was. + if constexpr (std::is_same_v) { + // UNCHANGED from the base commit apart from this block's four-space re-indent, + // so this path can be reviewed as "not touched" (strip comments and whitespace + // from both revisions and compare: 746 characters of code, identical). 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 signal barrier, for the RAIL team only. Every sender adds 1 to the + // SAME signal id on every peer, and the waiter advances its shadow by the + // number of increments one barrier round delivers, (kNumRanks - 1). A + // synchronisation-only barrier is a counting predicate, so it does not need to + // distinguish senders -- see the three-point argument above for why that is + // true here and false on the World path. + // + // This costs ONE indexed-signal slot instead of (kNumRanks - 1), which is what + // removes the team-size term from the per-context signal budget, and hence the + // scale-out ceiling: measured on p6-b200, 22 NVLink domains complete and 23 + // refuse at the shipped context count. + // + // The pattern is the one the unordered data path already relies on: a single + // signal accumulating increments from many remote senders, polled against a + // shadow advanced by the expected delta (`hybrid_combine_unordered.cuh`, the + // `num_expected_arrivals` wait; sends there use `ncclGin_SignalAdd{.., 1}`). + // `SignalAdd{.., 1}` is used here rather than `SignalInc` to match that + // precedent exactly. + // Reserved id, off-limits to the data path -- see `kNumReservedBarrierSignals`. + 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}); + } + + // Two block-wide syncs bracket the wait, and they do different jobs. + // + // BEFORE: every thread has issued its strided share of the sends above. Without + // this, thread 0 enters the spin below while lanes 1..31 of its own warp still + // have sends pending -- a divergent warp with one tight polling loop starves the + // sending lanes, and at large `kNumRanks` that delays outbound sends against a + // running timeout. + // + // AFTER: a single slot means a single waiter, so the other threads must not run + // ahead of the barrier. The per-peer layout had every thread wait on its own + // slot, which made "all threads have observed completion" implicit; it has to be + // restored explicitly here. + // + // The enclosing `if (sm_idx == 0)` is uniform across the block (`sm_idx` is + // derived from `blockIdx.x` at every call site, including the synthetic + // `sm_idx - 1` in `gpu_barrier`'s hybrid split), so both syncs are reached by + // all of the block's threads. + __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) { + // Report the shortfall: with one counting slot the stalled *peer* is + // no longer identifiable, so print how many of the expected arrivals + // are missing, and the signal id. + 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(); } } } @@ -285,6 +402,41 @@ __forceinline__ __device__ void gpu_barrier(const handle::NCCLGin& gin, EP_STATIC_ASSERT(not kFlushStores, "No data to be flushed"); } + // A GIN scale-up barrier (`ncclTeamTagWorld`) and a GIN scale-out barrier + // (`ncclTeamTagRail`) would both land on the SAME shadow: NCCL addresses shadows by + // (context, signal) only -- `_signalShadows = comm.ginSignalShadows + contextIndex * + // comm.ginSignalCount` -- with no team or tag term. Both use context 0, and their id + // ranges overlap even though the two now run different protocols: World's per-peer + // slots start at 0 and Rail's counting slot IS 0. So if they were ever live + // CONCURRENTLY they would corrupt each other -- World would read Rail's increments as + // peer 0's arrival, and Rail's count would absorb World's. Rail is a subset of World, + // so the two rank sets are not even disjoint. + // + // Three ways that cannot happen, any one of which is sufficient: + // * scale-up runs over NVLink, so it never touches a GIN signal at all; or + // * there is no scale-up team to synchronize (`do_scaleup` is masked off below); or + // * there is no scale-out team to synchronize (`do_scaleout` likewise). + // Note this is about CONCURRENCY: `barrier.cuh`'s sequential path issues Rail and World + // from two separate, globally ordered `gpu_barrier` calls and is safe for that reason, + // which is why the condition below is a disjunction rather than the stricter + // `kIsScaleupNVLink or kNumScaleoutRanks <= 1`. + // + // Today no instantiation can violate it: `NCCLSymmetricMemoryContext` + // (`csrc/kernels/backend/nccl.cu`) sets `num_scaleup_ranks = num_nvl_ranks` in hybrid mode + // -- making `is_scaleup_nvlink` true there by construction -- and `num_scaleout_ranks = 1` + // in direct mode. This assert exists so that stops being an accident. + // + // CAVEAT on the safety net: these kernels are NVRTC-generated from runtime values + // (`csrc/kernels/elastic/*.hpp` format the template arguments into the instantiation), so + // a violation surfaces as a JIT compile exception on first launch, in production -- NOT + // as a build failure. The host-side `EP_HOST_ASSERT` in `NCCLSymmetricMemoryContext` is + // the gate that actually fails early; this one is the backstop for a caller that + // constructs the template arguments some other way. + 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..240d77b06 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,66 @@ __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 device barrier, taken off the bottom of EVERY +// context's id space. This is for the RAIL instantiation of `gin_barrier_wo_local_sync`, +// which is a counting barrier and needs exactly one slot whatever the team size -- but it +// must be an id no data channel can ever produce, which is what this reservation buys. +// +// The World instantiation keeps the per-peer barrier and is NOT covered by this +// reservation. It does not need to be: it is provisioned from the other arm of +// `NCCLSymmetricMemoryContext` (`reqs.ginSignalCount = num_ranks + 2 * 2`), and the +// unordered data path that `data_signal_id` shifts does not run on that arm at all -- +// `get_qp_signal_id` / `get_per_part_signal_id` are called only from +// `hybrid_{dispatch,combine}_unordered.cuh`. +// +// This is a fix, not hardening. In `cached_mode` the unordered hybrid kernels run with zero +// notify warps (`dispatch.hpp`), which makes `kQPStartIdx = 0` (`qp_mapping.cuh`) and puts +// data channels on context 0 -- the barrier's own context. `channel_to_signal_id` is 0-based, +// so channel 0 part 0 produced id 0 there: exactly the barrier's slot. NCCL addresses a +// shadow by (context, signal) alone, so that was a shared counter, and each side inflated +// the other's arrival count. +// +// Only the barrier's own context (QP 0) strictly needs the hole, yet the reservation is +// applied uniformly, so the id derivation does not have to depend on which QP a channel +// landed on. That is NOT free -- measured with `compute_part_allocation` at the shipped +// context count, 16-17 SMs lose one part (3 -> 2) and 51-52 SMs lose a channel per SM +// (4 -> 3, with the host `[WARN]`). Weigh that before changing `kDefaultGinContextCnt`. +// +// ONE id, not two, because the Rail barrier is the only counting barrier and a GIN +// scale-up barrier and a GIN scale-out barrier can never be live at once -- `get_logical_domain_size` (`nccl.cu`) sets `num_scaleup_ranks` to +// `num_nvl_ranks` in hybrid mode (so `kIsScaleupNVLink` is true there by construction) and +// forces `num_scaleout_ranks` to 1 in direct mode. `gpu_barrier` asserts that, and +// `NCCLSymmetricMemoryContext` asserts it again on the host where it fails at init. +// +// NOTE: plain `int`, not an NCCL signal type -- this header is deliberately NCCL-free so it +// stays host-compilable, which is what lets the invariants below 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` +// (per-channel and per-part) route through this, so the offset cannot be dropped from one +// and kept in the other. +__forceinline__ __device__ __host__ constexpr int data_signal_id(int raw_offset) { + return kNumReservedBarrierSignals + raw_offset; } -static_assert(all_gin_context_counts_cover_warps(), - "GIN layout cannot give each ScaleOut warp a dedicated signal id " - "for every legal context count"); + +// For every legal context count the layout 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, without tripping its own host assert. +// +// This replaces an earlier aggregate check, `ctx * signals/ctx >= kMaxScaleoutWarps`, which +// claimed to prove "every ScaleOut warp gets a dedicated id". It never did: the ids are +// per-context, so an aggregate total says nothing about the busiest context, and at ctx = 13 +// the worst-case launch already needed 19 ids against 17 available and was silently relying +// on the tuner to cut channels. The reservation made the old check additionally stale +// (usable is `ctx * (signals - 1)`), which is what surfaced the problem. +// +// Channel reduction at high SM counts is expected and is a graceful degradation -- the tuner +// warns and proceeds. What must never happen is an unserviceable configuration, and that is +// what this asserts. `compute_part_allocation` is defined below; the check sits with it. +__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 +166,21 @@ __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 <= gin_indexed_signals_cnt - kNumReservedBarrierSignals // at the requested channels/SM, then reduce channels_per_sm until the budget holds. -__forceinline__ __device__ __host__ constexpr GinPartAllocation compute_part_allocation( +// +// Split in two: `_raw` is the pure math, free of diagnostics, so it can be evaluated inside a +// `static_assert` (a `printf` or a throwing assert reached during constant evaluation makes +// the expression non-constant). `compute_part_allocation` is the shipping entry point and +// adds the host-side warning and check on top. Keep the math in `_raw` only -- duplicating it +// into the invariant below is exactly the drift this split exists to prevent. +__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; + // The data path may only use ids at or above `kNumReservedBarrierSignals`, so the + // usable budget is that much smaller than the provisioned count. + const int provisioned = cfg.gin_indexed_signals_cnt; + const int gin_signals = provisioned > kNumReservedBarrierSignals + ? provisioned - kNumReservedBarrierSignals : 0; 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 +191,54 @@ __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 actually fits the usable budget -- i.e. the reduction loop above +// converged rather than bottoming out at one channel per SM and still not fitting. +__forceinline__ __device__ __host__ constexpr bool part_allocation_fits( + const GinResourceConfig& cfg, int num_sms, int num_available_qps, const GinPartAllocation& alloc) { + const int provisioned = cfg.gin_indexed_signals_cnt; + const int gin_signals = provisioned > kNumReservedBarrierSignals + ? provisioned - kNumReservedBarrierSignals : 0; + return static_cast(channels_per_context(num_sms, num_available_qps, + alloc.num_channels_per_sm)) * alloc.num_parts + <= gin_signals; +} + +__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(