device: fix single-node NVLS AllGather/ReduceScatter rank ordering (#1906)

[EylonKrause]

Fixes #1906.

The bug

Single-node NVLS AllGather returns each rank's chunk in the wrong output position when CUDA_VISIBLE_DEVICES permutes GPUs out of dense-head order — a silent data corruption (the output is a permutation of the correct result). @sjeaugey confirmed the root cause on the issue: the single-node NVLS AllGather is "missing the user rank table which ensures proper data ordering."

In the non-registered oneNode path:

src/device/all_gather.h:     prims.gather(offset, nvls->nHeads * count, nelem, count, -1, 0);
  ⟶ src/device/prims_simple.h (ScatterGatherOp, Recv branch):  pOffset = i * peerOffset;

the peer loop index i is the dense head index, so head i's gathered slice is written to output slot i. But head i physically holds the data of user rank denseToUserRank[i], so slot i is wrong whenever denseToUserRank[i] != i.

Every correct sibling path already remaps through denseToUserRank:

• multi-node NVLS — all_gather.h:317: int rank = denseToUserRank[node*nRails+rail]; userOneBeg = rank*countPerRank;
• CollNetDirect — all_gather.h:546 (same form)
• the PAT helper nvlsDenseToUserOffset in prims_simple.h

ReduceScatter has the exact symmetric mirror (reduce_scatter.h, oneNode/!regUsed scatter → pOffset = i*peerOffset on the source side). The oneNode regUsed paths use rank*count directly and are already correct; multi-node paths are already correct; and AllReduce's NVLS scatter+gather round-trip is order-invariant (its output is the fully-reduced array), so it must not change.

Why this is device-only

The denseToUserRank table is already built and copied to the device for pure single-node NVLS on current master — ncclTransportInitRankMap (init.cc:1573) is gated only on comm->nvlsSupport (no nNodes>1 guard), and devCommSetup copies it (init.cc:698, null-guarded). So no host / comm.h / init.cc / connect.cc changes are needed — only the device-side consumption was missing.

The fix

• ScatterGatherOp gains an opt-in template <..., bool Remap = false> plus a const int* remap = nullptr; the per-peer offset becomes (Remap && remap != nullptr) ? remap[i] : i. Because Remap is a template constant, every existing caller compiles to byte-identical code (AllReduce, multi-node, CollNet, the two reg-sync gather/scatter(0,0,0,0,-1,0), directScatter/directGather) — they don't pass the new arg.
• New scatterRemap() / gatherRemap() wrappers set Remap=true.
• The two oneNode NVLS sites pass ncclShmem.comm.denseToUserRank, so peer i's slice is placed at slot denseToUserRank[i] — matching the multi-node/CollNet/PAT direction exactly.

Direction (the load-bearing detail)

Specialized to a single node (node=0, nRails=nHeads), the correct multi-node formula rank = denseToUserRank[node*nRails+rail]; slot = rank*count reduces to slot(i) = denseToUserRank[i]. denseToUserRank[i] is the user rank that owns head i (generic.cc builds userToDenseRank then inverts it; for a head rank r, userToDenseRank[r]=h where nvlsHeads[h]==r). So peer i (head i) → output slot denseToUserRank[i]. Using the inverse table would double-permute; this uses the same table and direction as every correct path.

Testing

I cannot validate this on hardware — I don't have an NVSwitch box and can't do a full NCCL build. I verified: (1) the fix matches the multi-node / CollNet / PAT remap direction line-for-line; (2) the C++ template/overload/wrapper mechanics compile clean and a standalone harness confirms scatterRemap/gatherRemap permute the per-peer offset by remap[] and that an identity remap is byte-identical to plain scatter/gather.

Please validate on NVSwitch before merging — I'd suggest all_gather_perf / reduce_scatter_perf forcing NCCL_ALGO=NVLS, with #wrong == 0 across the identity and several non-identity CUDA_VISIBLE_DEVICES permutations (e.g. 1,2,3,4,5,6,7,0 and 3,4,5,6,7,0,1,2), plus a non-NVLS regression pass.

Two notes for reviewers: the ReduceScatter change is the symmetric mirror (the issue reproduces AllGather only — happy to drop it if you'd rather land AllGather first), and both sites keep the existing nvls->nHeads*count span, which assumes nHeads == nRanks on this single-node path (as the current code already does).

Disclosure: this contribution was authored with an AI coding assistant (Claude) and reviewed before submission.