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

[KeitaW]

Problem

csrc/kernels/backend/nccl.cu asserts at Buffer init that the per-context GIN indexed-signal
budget can give one dedicated signal to every rail peer:

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");

num_rdma_ranks is ncclTeamRail(comm).nRanks, the number of NVLink domains. The supply is
(kTotalQPBudget - 2c)/c for c contexts — a constant in the domain count — while the
requirement grows linearly with it. The two cross:

contexts	signals/ctx	passes to	refuses at
2	126	127 domains	128
11 (shipped default)	21	22 domains	23
17	13	14 domains	15

At 8 GPUs per domain the shipped default stops at 176 ranks.

The ceiling is not a fixed number, though — it is a function of the QP count, and the two are
traded against each other. A num_allocated_qps of 5, 6 or 7 reaches 32 or more domains today,
and gin_resource_alloc.cuh lists {5, 6, 7, 8, 9, 14} as part-neutral alternatives to the
default 11, so that trade need not cost a part.

What it costs is QPs, and that was measured rather than assumed. On 8 × p6-b300 (64 ranks,
256 experts, 4096 tokens, 24 SMs, GDAKI GIN), holding everything constant but
--num-allocated-qps, dropping 11 → 5 costs +1.9 % on dispatch and +30.5 % on combine
(3241 → 4228 µs; combine bandwidth 96.3 → 74.6 GB/s). Eight cells, both orders, two repeats;
the effect holds in both orders, so it is not run-order bias. Dispatch barely notices losing six
of eleven QPs; combine, the direction with the heavier concurrent per-rank traffic, loses nearly
a third.

So the defect is not that 22 is too low. It is that scale-out reach and data-path QP
parallelism are drawn from the same budget, and no setting buys both. This change makes the
rail barrier's cost constant in the domain count, which removes the trade rather than relocating
it.

A GIN signal id is not free on EFA. One id is one gdaki_sc_endpoint, a complete QP and CQ,
which is why the budget is expressed in QPs. Raising the budget is not available either. The
boundary has been measured directly in request space on B200: a devComm's request costs one
endpoint for each context's data path plus one per signal id, and requests totalling 256
endpoints allocate while 257 fail. Three independent context/signal splits land on exactly
256 and allocate, each failing at the next step up. kTotalQPBudget = 256 is therefore the
measured constant at the granularity it is written in — one endpoint carries one QP and one
CQ, and the cells do not separate which of those is the scarce object, only that the cap is
one per endpoint. The scope is one devComm's request rather than the device total; those
measurements are being submitted separately.

A second defect on the same ids

Independent of the ceiling, the barrier's ids were already colliding with the data path.
dispatch.hpp sets num_notify_warps = 0 in cached_mode, which makes kQPStartIdx = 0 and
places data channels on context 0 — the barrier's own context. channel_to_signal_id is
0-based, so channel 0 / part 0 derived id 0: exactly the barrier's slot.

NCCL addresses a signal's shadow by (context, signal) alone:

// nccl_device/impl/gin__funcs.h
gin->_signalShadows = comm.ginSignalShadows + contextIndex * comm.ginSignalCount;

No team term, no tag term — so that is a genuinely shared 64-bit counter, with each side
inflating the other's arrival count.

Why the assert is load-bearing

Two independent consumers draw on the per-context budget, and only one scales with the domain
count:

consumer	indexing	slots	grows with domains
data path	(sm, channel, part), no peer term	ceil(channels/qp) × num_parts	no
rail barrier	per peer	num_rdma_ranks − 1	yes

Deleting the assert would convert a loud init failure into the barrier indexing outside the
provisioned range — the silent mode the header already warns about, where no counts arrive and
dispatch times out with all-zero received counts.

Fix

A barrier is a counting predicate and does not need to distinguish senders. Every peer now adds
1 to the same signal id, and the waiter advances its shadow by kNumRanks - 1 and polls
that one signal. Message count is unchanged; the slot requirement drops from N−1 to one
whatever the team size, which removes the team-size term from the init check.

This is the pattern the unordered data path already relies on — one 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). SignalAdd{.., 1} matches
that precedent exactly, and avoids the documented rule that Inc may not be mixed with other
signal operators without an intervening reset().

Scoped to the rail team, deliberately

The two team instantiations have opposite requirements, so only rail is converted.

Only rail has a ceiling. The unordered-hybrid arm requests gin_indexed_signals_cnt, the
per-context budget above. The direct / ordered arm requests 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,
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 holds on InfiniBand too.

The world body is therefore left unchanged apart from the four-space re-indent the new
if constexpr forces around it. Stripping comments and whitespace from the block in both
revisions gives 746 characters of identical code.

Reserving an id the data path cannot produce

Collapsing to one slot is not enough on its own, because of the collision above.
kNumReservedBarrierSignals 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.

The reservation is applied uniformly across contexts although only QP 0 needs it: the tuner
budgets for the worst context anyway, and a uniform offset keeps the id derivation independent
of which QP a channel landed on. The offset goes in after the per-part multiply — inside
channel_to_signal_id it would be scaled by kNumParts and burn ids, so that function stays
0-based and keeps its host unit test unchanged.

One id, not two, because a GIN scale-up and a GIN scale-out barrier can never be live
concurrently — now asserted rather than left implicit, since their id ranges overlap (world's
per-peer slots start at 0 and rail's counting slot is 0):

EP_STATIC_ASSERT(kIsScaleupNVLink or kNumScaleupRanks <= 1 or kNumScaleoutRanks <= 1, ...);

A disjunction of three sufficient conditions. The stricter
kIsScaleupNVLink or kNumScaleoutRanks <= 1 would be wrong: it rejects barrier.cuh's
sequential path, which issues rail and world from two separate, globally ordered calls and is
safe for that reason. The hazard is concurrency, so the condition is about concurrency.

The static assert is a backstop, not the gate. These kernels are NVRTC-generated from
runtime values, so a violation surfaces as a JIT exception on first launch rather than as a
build failure. The same condition is an EP_HOST_ASSERT in NCCLSymmetricMemoryContext, where
both values are known and it fails at Buffer init.

Two invariants repaired alongside

all_gin_context_counts_cover_warps() compared a cross-context total against a warp count.
Ids are per-context, so an aggregate says nothing about the busiest context; at ctx=13 the
worst-case launch already needed 19 ids against 17 and was relying on the tuner to cut channels.
It is replaced by the real contract: every legal context count must remain serviceable at
the worst-case launch. Verified to bite — forcing the reservation to 200 fails the build.

To make that assertable, compute_part_allocation is split into a pure _raw and a diagnosing
wrapper: a printf or a throwing assert reached during constant evaluation makes the expression
non-constant, so the invariant could not otherwise call the shipping math — and duplicating the
math is exactly the drift the split prevents.

Cost

Host-compiled the shipped tuner across 24 (SMs, channels/SM) shapes at the default 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). Usable per-context budget goes 21 → 20; a second id would have
cost roughly three times as many shapes, which is what the concurrency assert buys.

One __syncthreads() is added. One slot means one waiter, and the per-peer layout had every
thread wait on its own slot, which implicitly guaranteed that no thread in the block ran ahead
of the barrier — that has to be restored explicitly.

The timeout diagnostic changes: with one counting slot the stalled peer is no longer
identifiable, so the message prints how many of the expected arrivals are missing plus the
signal id. The previous message never printed the slot index either, so per-peer attribution was
not available before.

Validation

36 × p6-b200.48xlarge, 8× B200 per node, one NVLink domain per node, 8 EFA per node, NCCL
2.31.2, aws-ofi-nccl with the EFA-GDA GIN backend, libfabric 2.6.0. Treatment and control built
from the same base and differing only in the DeepEP source tree; every rank printed the md5 of
the JIT headers it compiled, so each arm is provably the tree it claims.

domains	build	result
23	stock	refuses at init, verbatim assert above
23	this change	completes, correctness checks pass
32	this change	completes (256 ranks)
22	this change	completes
2	this change	completes
3, --allow-hybrid-mode 0	this change	completes; exercises the world instantiation
3, --allow-hybrid-mode 0	stock	completes; control for the row above

Host-side, with no GPU, since these headers are NCCL-free by design: the serviceability
invariant compiles for every legal context count and fails when the reservation is forced to
200; and every id the data path can derive falls in [1, provisioned−1] across all legal
context counts, with and without notify warps.

Performance

Barrier in isolation, mean over runs in both job orders; the 8-rank spread within a run is under
1%. The counting barrier is never slower:

domains	per-peer	counting	delta
2	27.97 µs	25.68 µs	−8.2 %
8	38.16	36.30	−4.9 %
16	60.15	58.15	−3.3 %
22	79.64	78.94	−0.9 %
26	n/a	94.59	—
32	n/a	112.65	—

The saving is a fixed ~2.2 µs per call rather than a scaling one, so it fades as a percentage
while the barrier's absolute cost grows. The single-counter hotspot does not materialise through
32 domains: marginal cost is 3.38 µs/domain over 22→32 against the per-peer barrier's
3.25 µs/domain over 16→22, and the slope falls again over 26→32. End-to-end dispatch and combine
throughput is unchanged.

Note on #3

This change is textually independent of #3 (different files), but on Blackwell no kernel builds
until #3 lands, so the validation above was run with both applied. Remaining limitations and
untested paths are tracked separately rather than in this PR.

[Xuan-1998]

Thanks for the pr! I have two high level comments:

1. Let remove the long comments which give the reason about why the change from the code and put them in the pr description if necessary. In the code we want to have short explanation about what current code is doing if any comments added.
2. Please rebase on latest main to resolve the conflicts.

[KeitaW]

Thanks for the review. Rebased on latest main; the envs and tests commits were already there, so the branch is now just the barrier commit. The code comments are trimmed to short descriptions of what the code does (the rationale was already in the PR description), and usable_signal_budget() is factored out and called from both places as suggested. The comment changes are a separate commit (6cfc91f) to make the delta easy to see.