fix: resolve the GIN context count on the single-node path (kNumQPs == 0 corrupts combine)

[KeitaW]

Summary

On a single node the GIN context count is never resolved, so kNumQPs == 0 reaches the combine kernel and balanced_partition() is called with q == 0. Nothing rejects it, and the run does not fail: it produces a wrong result, reported by tests/elastic/test_ep.py as AssertionError: Diff: nan.

Two nodes are unaffected, which is why this does not show up in multi-node testing. "Single node" here still means eight local ranks, so the device-communicator setup is entered normally; what differs is that ncclTeamRail(comm).nRanks is 1.

The chain

All at main @ 02efc268.

1. deep_ep/buffers/elastic.py:350-358, automatic QP count. When the caller leaves num_allocated_qps at its default of 0 and the unordered hybrid kernels are in use, the Python auto-fill is deliberately skipped, and the comment says why: the count is resolved in C++ from the GIN signal budget. In that configuration 0 is handed to C++ on purpose.

2. csrc/kernels/backend/nccl.cu:115-138. resolve_gin_context_cnt() is the resolver Python delegates to, and it is only called on the scaleout_active branch. The single-node branch copies the sentinel verbatim:

   if (scaleout_active) {
       gin_config = make_gin_resources(resolve_gin_context_cnt());
   } else {
       gin_config.gin_context_cnt = this->num_allocated_qps;   // still 0
       gin_config.gin_indexed_signals_cnt = 0;
   }
   ...
   if (scaleout_active)
       this->num_allocated_qps = gin_config.gin_context_cnt;   // write-back also gated

   The range check that would have rejected 0 (ctx >= kMinGinContextCnt, minimum 2) lives inside that lambda, so it is not on this path.

3. elastic.py:862-879, get_theoretical_num_qps(). min(num_sms * 16 + 1, self.num_allocated_qps) becomes min(1025, 0) = 0, and that 0 becomes the kNumQPs template parameter of the combine kernel. Because it is a template parameter, the branch is selected at JIT compile time.

4. common/comm.cuh:77-99 and common/qp_mapping.cuh:47-57,62-88. get_qp_mode special-cases kNumQPs == 1 only. With 0, kNumAvailableQPs is 0, kNumSMs <= kNumAvailableQPs is false, and control reaches balanced_partition(global_channel_idx, 512, 0), whose const int base = n / q; sits directly under the comment "Callers guarantee q >= 1".

Compiling the shipped qp_mapping.cuh on the host at kNumSMs=64, kNumQPs=0 exits 136 (SIGFPE), which confirms the branch is reachable with the real parameters. In the GPU runs below the same specialization produced corrupted output and no reported fault.

Why the symptom is misleading

The assertion message says NaN, but in an instrumented single-node run (eight ranks, unpatched) the combine output differs from the reference in 469,453,194 of 469,561,344 elements, 99.977%, while only 8,300 elements, 0.0018%, are actually NaN. The inputs are NaN-free, rows containing no NaNs still differ, and sampled bad rows carry thousands of distinct bit patterns.

The clearest signal is the timing: with #QPs: 0/0 the combine reports 622,881,536 bytes in 1.148 us, which the harness prints as 542,580 GB/s. That timing is not compatible with the kernel having completed the stated payload. calc_diff then returns NaN because some leftover bit patterns decode as NaN, which is how a resource-indexing problem ends up presenting as a numerics one.

The change

• nccl.cu: call resolve_gin_context_cnt() on the single-node branch too, and make the write-back unconditional. Python reads the value back through get_num_allocated_qps() and caps the per-launch QP count with it, so a path that does not write back leaves the Python side at 0. gin_indexed_signals_cnt = 0 is unchanged; a single node has no peer rail teams.
• comm.cuh, qp_mapping.cuh: static_assert(kNumQPs >= 1, ...) in get_qp_mode, channel_to_qp and channel_to_signal_id, so an illegal specialization fails to compile instead of miscomputing. I deliberately did not widen the fast path to kNumQPs <= 1: 0 means no QP was allocated, and running it as if one existed would hide a resource-allocation problem rather than surface it.
• elastic.py: assert the delegation invariant immediately after get_num_allocated_qps(), where it is observable. With the C++ change in place this never fires; it exists so a future branch that skips resolution fails at the handoff rather than downstream in a kernel. The existing assert num_qps <= self.num_allocated_qps guards pass vacuously when both sides are 0.

Verification

p6-b200.48xlarge, EFA kernel driver 3.3.0g, CUDA 13.1.2, NCCL v2.31.2-1, aws-ofi-nccl v1.21.1, EFA installer 1.50.0, torch 2.11.0+cu130. Predictions were written before the runs.

The two builds are the same container image differing only in the DeepEP tree installed into it: base 02efc268 versus this branch. They ran on different but matched nodes of the same instance type in the same cluster, with identical launch arguments and environment, so this is a controlled comparison rather than a strict single-variable experiment. One run per configuration, so the figures below carry no uncertainty estimate; ranges are across all ranks.

Arm	Build	Nodes	#QPs	Exit	Measurement (all ranks)
control	base	1	0/0	1, Diff: nan	combine 1.148 us, reported 542,580 GB/s
fix	patched	1	11/11	0	combine 711-746 GB/s (SU), 837.5-873.5 us
control	base	2	11/11	0 / 0	dispatch 81-82 (SO) / 265-270 (SU) GB/s, 1486-1500 us
fix	patched	2	11/11	0 / 0	dispatch 81-82 (SO) / 264-270 (SU) GB/s, 1485-1504 us

The two-node arms are the regression check, since the unconditional write-back is on a line the scale-out path also executes. Result: no correctness regression, the QP count is unchanged at 11, and the dispatch ranges overlap. With one run per build I would put that as "no observed regression", not as proof of identical performance.

A third single-node arm confirms the two problems stay separate: with this patch applied but a launcher environment that sets neither NCCL_GIN_TYPE nor NCCL_IB_DISABLE, the run still fails earlier, at nccl.cu:189, with Requested properties for GIN GDAKI NIC 6, only 2 GIN GDAKI NICs have been created. The passing arms above set both.

Compile-side checks on the guard:

kNumQPs	0	1	2	8	11	17	64	1025
qp_mapping.cuh, host g++	static_assert	ok	ok	ok	ok	ok	ok	ok
comm.cuh, nvcc -arch=sm_100 with the JIT flags	static_assert	ok	ok	not run	ok	ok	not run	ok

1 is the existing fast path, 2 and 17 are kMinGinContextCnt and kMaxGinContextCnt, 11 is kDefaultGinContextCnt, and 1025 is the hybrid num_sms * 16 + 1 at 64 SMs. The negative sentinel kFlushAllAllocatedQPs = -1 is used only for the barrier instantiations in impls/barrier.cuh and is converted to a runtime value in gin_barrier_wo_local_sync, so it never reaches these templates. The full tree builds and installs with the change (TORCH_CUDA_ARCH_LIST="9.0;10.0"); note that the AOT build does not instantiate the guarded templates, which live in the JIT translation unit, so the GPU runs above are what exercises them.

Two things I did not do

No repository test is added. The natural regression path, tests/elastic/test_ep.py, requires GPUs and distributed initialization, and it now passes on a single node with this change. I did not want to invent a new test shape without knowing your CI. If you would like one, tell me where it should live and I will add it.

No allocation-cost measurement for the default. The single-node branch now goes through resolve_gin_context_cnt(), so an unspecified count becomes kDefaultGinContextCnt = 11. A single node uses no RDMA, so those contexts are allocated and unused. I kept the default because it follows the existing resolver policy rather than introducing a second constant, but kMinGinContextCnt = 2 is a reasonable alternative if the footprint matters. On the unmodified library, 2/2 and 8/17 performed the same in a single run each, so I have no evidence either way. Happy to change it.

[Xuan-1998]

Think we can get rid of the signal here since the code is self explainable

[Xuan-1998]

ditto

[Xuan-1998]

Thanks @KeitaW, this is great. We will also add this check to our nightly to exercise the intra node path as well