fix(rccl): fit AINIC scatter parts to the P2P channel pool

projects/rccl/CHANGELOG.md

@@ -56,7 +56,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 * Fixed virtual address space exhaustion when symmetric windows backed by the same physical allocation are registered repeatedly.
 * Fixed Copy Engine `ncclAlltoAllv` hanging on intra-node communicators of more than 8 ranks with uneven send sizes. Each rank decided from its own send sizes whether to add intra-batch barriers, so ranks could issue different barrier counts. The decision now uses the full send-size matrix, which every rank shares.
 * Fixed `NCCL_P2P_DISABLE` and `NCCL_P2P_LEVEL` being ignored on Intel, ARM and Zhaoxin hosts, where the CPU-based default P2P level overwrote the user setting. On such hosts `NCCL_P2P_DISABLE=1` now falls back to SHM as documented, and `NCCL_P2P_LEVEL` is honored.
-* Fixed a 15 to 30% `ncclScatter` bandwidth loss from 8 MB upward on multi-node AINIC jobs of up to 8 nodes, caused by the move to two network channels per peer. Scatter now uses and connects one channel per peer when the peers cannot each get their own channels. Point-to-point operations and the other collectives keep two channels.
+* Fixed a 15 to 30% `ncclScatter` bandwidth loss from 16 MB upward on multi-node AINIC jobs, caused by the move to two network channels per peer (for example on 4 nodes of 8 GPUs). Scatter now uses and connects the most channels per peer whose flows still fit the P2P channel pool. Point-to-point operations and the other collectives keep their channel count.
 
 ### Known issues
 * The FP8 ReduceSum and ReduceCopy device APIs are not exported in the LLVM bitcode library.

projects/rccl/src/enqueue/enqueue.cc

@@ -1327,18 +1327,17 @@ static bool rcclP2pBatchEligible(struct ncclComm* comm, ssize_t sendBytes, ssize
 }
 
 // Number of the p2pnChannelsPerPeer parts that a task of collAPI uses and connects.
-// On AINIC ncclTopoComputeP2pChannels raises p2pnChannelsPerPeer to 2 so one peer can saturate a NIC.
-// Scatter already drives all nRanks-1 peers at once from one root and loses 15-30% from 8 MB up with
-// two parts per peer, so it uses only part 0 once the peers cannot each own p2pnChannelsPerPeer
-// channels. It must not connect part 1 either: when scatter connects both parts itself, it stays at
+// Scatter drives all nRanks-1 peers at once from one root, so with p2pnChannelsPerPeer parts its flows
+// can exceed p2pnChannels (on 4 nodes of 8 GPUs with AINIC: 62 flows on 32 channels, 15-30% slower
+// from 16 MB up). ncclTopoComputeP2pChannels gives it p2pnChannelsPerPeerScatter instead: the most
+// parts whose flows fit the pool.
+// It must not connect the extra parts either: when scatter connects both parts itself, it stays at
 // or below the two-part speed even though it sends on part 0 only. Part 1 connected later by another
 // collective does not slow it down. The answer may depend only on values both ends of a pair share
 // (collAPI, comm-wide sizes), never on per-rank task counts, or sender and receiver would disagree on
 // the parts.
 int rcclP2pTaskParts(struct ncclComm* comm, ncclFunc_t collAPI) {
-  if (collAPI == ncclFuncScatter && comm->nNodes > 1 && comm->p2pnChannelsPerPeer > 1 && rcclUseAinic() &&
-      (comm->nRanks - 1) * comm->p2pnChannelsPerPeer > comm->p2pnChannels)
-    return 1;
+  if (collAPI == ncclFuncScatter) return comm->p2pnChannelsPerPeerScatter;
   return comm->p2pnChannelsPerPeer;
 }
 
@@ -1348,11 +1347,12 @@ static int rcclP2pPolicyChannels(struct ncclComm* comm, struct ncclTaskP2p* task
     ncclFunc_t collAPI;
     int nChannels;
   };
-  // This workaround targets direct P2P/IPC traffic only. When P2P is disabled,
-  // preserve the normal channel selection for the fallback transport.
   if (task == nullptr) return -1;
+  // The Scatter part limit is about network flows, so it applies with P2P disabled too.
   int taskParts = rcclP2pTaskParts(comm, task->collAPI);
   if (taskParts < comm->p2pnChannelsPerPeer) return taskParts;
+  // The policies below target direct P2P/IPC traffic only. When P2P is disabled,
+  // preserve the normal channel selection for the fallback transport.
   if (ncclParamP2pDisable()) return -1;
   constexpr Policy policies[] = {
     // Two-channel full-mesh traffic has high variance on gfx110x. Keep the

projects/rccl/src/graph/paths.cc

@@ -1504,9 +1504,19 @@ ncclResult_t ncclTopoComputeP2pChannels(struct ncclComm* comm) {
     NCCLCHECK(ncclTopoReconcileGrowChannels(comm, &comm->p2pnChannels));
     comm->p2pnChannelsPerPeer = std::min(comm->p2pnChannelsPerPeer, comm->p2pnChannels);
   }
-
-  INFO(NCCL_INIT | NCCL_TUNING, "P2P channels: p2pnChannels=%d p2pnChannelsPerPeer=%d maxP2pPeers=%d",
-       comm->p2pnChannels, comm->p2pnChannelsPerPeer, maxP2pPeers);
+  // Scatter drives all nRanks-1 peers from one root at once. On multi-node AINIC, give it the most parts
+  // per peer whose flows still fit the pool; more parts oversubscribe it. A user-set
+  // nChannelsPerNetPeer is kept as is.
+  comm->p2pnChannelsPerPeerScatter = comm->p2pnChannelsPerPeer;
+  if (comm->nNodes > 1 && comm->config.nChannelsPerNetPeer == NCCL_CONFIG_UNDEF_INT && rcclUseAinic()) {
+    int partsThatFit = comm->p2pnChannels / (comm->nRanks - 1);
+    partsThatFit = partsThatFit >= 1 ? pow2Down(partsThatFit) : 1;
+    comm->p2pnChannelsPerPeerScatter = std::min(comm->p2pnChannelsPerPeer, partsThatFit);
+  }
+
+  INFO(NCCL_INIT | NCCL_TUNING,
+       "P2P channels: p2pnChannels=%d p2pnChannelsPerPeer=%d p2pnChannelsPerPeerScatter=%d maxP2pPeers=%d",
+       comm->p2pnChannels, comm->p2pnChannelsPerPeer, comm->p2pnChannelsPerPeerScatter, maxP2pPeers);
 
   // Init channels that weren't used so far
   for (int c = comm->nChannels; c < std::max(comm->nChannels, comm->p2pnChannels); c++) NCCLCHECK(initChannel(comm, c));

projects/rccl/src/include/comm.h

@@ -821,6 +821,8 @@ struct ncclComm {
   // Channels (per peer) for p2p
   int p2pnChannels;
   int p2pnChannelsPerPeer;
+  // Parts per peer for Scatter, at most p2pnChannelsPerPeer (see ncclTopoComputeP2pChannels).
+  int p2pnChannelsPerPeerScatter;
   int p2pChannelShiftSize;
   int p2pSchedGroupSize;
   int p2pMaxPeers;

projects/rccl/test/host/enqueue-test.cc

@@ -339,6 +339,48 @@ TEST_F(EnqueueMicrotest, ShmemDynamicSize_DeviceLinker_IsAlwaysZero) {
 }
 #endif
 
+// ===========================================================================
+// rcclP2pTaskParts (enqueue.cc:1340) and rcclP2pPolicyChannels (enqueue.cc:1345)
+//
+// Scatter takes p2pnChannelsPerPeerScatter, every other collAPI takes
+// p2pnChannelsPerPeer. Both ends of a pair must get the same answer, so it may
+// depend only on comm-wide fields.
+// ===========================================================================
+
+TEST_F(EnqueueMicrotest, P2pTaskParts_ScatterUsesPreRaiseCount) {
+  std::unique_ptr<ncclComm> storage{new ncclComm{}};
+  ncclComm& comm = *storage;
+  comm.p2pnChannelsPerPeer = 2;
+  comm.p2pnChannelsPerPeerScatter = 1;
+  EXPECT_EQ(1, rcclP2pTaskParts(&comm, ncclFuncScatter));
+  for (ncclFunc_t f : {ncclFuncSendRecv, ncclFuncSend, ncclFuncRecv, ncclFuncAlltoAll, ncclFuncGather,
+                       ncclFuncAlltoAllv}) {
+    EXPECT_EQ(2, rcclP2pTaskParts(&comm, f)) << "collAPI=" << int(f);
+  }
+}
+
+TEST_F(EnqueueMicrotest, P2pTaskParts_ScatterUnchangedWhenNotRaised) {
+  std::unique_ptr<ncclComm> storage{new ncclComm{}};
+  ncclComm& comm = *storage;
+  comm.p2pnChannelsPerPeer = 4;
+  comm.p2pnChannelsPerPeerScatter = 4;
+  for (ncclFunc_t f : {ncclFuncScatter, ncclFuncSend, ncclFuncAlltoAll, ncclFuncGather}) {
+    EXPECT_EQ(4, rcclP2pTaskParts(&comm, f)) << "collAPI=" << int(f);
+  }
+}
+
+TEST_F(EnqueueMicrotest, P2pPolicyChannels_ScatterBelowPerPeerReturnsItsParts) {
+  // Returns before the arch policy table, so comm.topo is never read.
+  std::unique_ptr<ncclComm> storage{new ncclComm{}};
+  ncclComm& comm = *storage;
+  comm.p2pnChannelsPerPeer = 2;
+  comm.p2pnChannelsPerPeerScatter = 1;
+  ncclTaskP2p task{};
+  task.collAPI = ncclFuncScatter;
+  EXPECT_EQ(1, rcclP2pPolicyChannels(&comm, &task));
+  EXPECT_EQ(-1, rcclP2pPolicyChannels(&comm, nullptr));
+}
+
 // ===========================================================================
 // calcP2pChannelCount (enqueue.cc:1550)
 //
@@ -3892,6 +3934,7 @@ struct Gfx1250AddP2pScene {
     c->cudaArch = 1250;
     c->p2pnChannels = 4;
     c->p2pnChannelsPerPeer = 1;
+    c->p2pnChannelsPerPeerScatter = 1;
     c->p2pChannelShiftSize = 0;
     c->p2pChunkSize = 128 << 10;
     c->ll128LineElems = 16;

projects/rccl/test/host/fakes/enqueue_fakes.cc

@@ -17,7 +17,7 @@
 
 #include "signature-drift.h"
 
-#include "channel.h"            // rcclEffectiveP2pBatchEnable
+#include "channel.h"            // rcclEffectiveP2pBatchEnable, rcclP2pTaskParts
 #include "enqueue_fakes.h"
 #include "nccl_fakes.h"         // g_loadParam, for the NCCL_PARAM body this file stands in for
 #include "sym_kernels_fakes.h"  // g_symkAvailable and the sym_kernels.cc seams below's canonical home
@@ -31,6 +31,7 @@ ASSERT_HOOK_MATCHES_PROD(g_addProxyOpIfNeeded, ncclAddProxyOpIfNeeded);
 ASSERT_HOOK_MATCHES_PROD(g_getCollNetSupport, ncclGetCollNetSupport);
 ASSERT_HOOK_MATCHES_PROD(g_getRegBuff, ncclGetRegBuff);
 ASSERT_HOOK_MATCHES_PROD(g_rcclEffectiveP2pBatchEnable, rcclEffectiveP2pBatchEnable);
+ASSERT_HOOK_MATCHES_PROD(g_rcclP2pTaskParts, rcclP2pTaskParts);
 #undef ASSERT_HOOK_MATCHES_PROD
 
 // Generous default: a deny-everything default would make even a single small task look over budget.
@@ -75,6 +76,12 @@ std::function<ncclResult_t(struct ncclComm*, struct ncclKernelPlan*, struct nccl
 static int DefaultEffectiveP2pBatchEnable(struct ncclComm*) { return 0; }
 std::function<int(struct ncclComm*)> g_rcclEffectiveP2pBatchEnable = DefaultEffectiveP2pBatchEnable;
 
+// Same field choice as enqueue.cc's rcclP2pTaskParts.
+static int DefaultP2pTaskParts(struct ncclComm* comm, ncclFunc_t collAPI) {
+  return collAPI == ncclFuncScatter ? comm->p2pnChannelsPerPeerScatter : comm->p2pnChannelsPerPeer;
+}
+std::function<int(struct ncclComm*, ncclFunc_t)> g_rcclP2pTaskParts = DefaultP2pTaskParts;
+
 // enqueue.cc:437-441 splits the default on ROCm version; mirrored rather than pinned to one arm.
 #if ROCM_VERSION >= 71200
 int64_t ncclParamGraphRegister() { return g_loadParam("GRAPH_REGISTER", 1); }
@@ -105,11 +112,15 @@ void ResetEnqueueFakes() {
   g_getCollNetSupport = DefaultGetCollNetSupport;
   g_getRegBuff = DefaultGetRegBuff;
   g_rcclEffectiveP2pBatchEnable = DefaultEffectiveP2pBatchEnable;
+  g_rcclP2pTaskParts = DefaultP2pTaskParts;
 }
 
 int rcclEffectiveP2pBatchEnable(struct ncclComm* comm) {
   return g_rcclEffectiveP2pBatchEnable(comm);
 }
+int rcclP2pTaskParts(struct ncclComm* comm, ncclFunc_t collAPI) {
+  return g_rcclP2pTaskParts(comm, collAPI);
+}
 
 // src/enqueue/enqueue.cc
 ncclResult_t ncclGetAlgoInfo(struct ncclComm* comm, struct ncclTaskColl* task, int collNetSupport, int nvlsSupport,

projects/rccl/test/host/fakes/enqueue_fakes.h

@@ -52,6 +52,9 @@ extern std::function<ncclResult_t(struct ncclComm*, struct ncclTaskColl*, int*)>
 // enqueue.cc's rcclEffectiveP2pBatchEnable: default 0 is batching off, the single-node answer.
 extern std::function<int(struct ncclComm*)> g_rcclEffectiveP2pBatchEnable;
 
+// enqueue.cc's rcclP2pTaskParts: default reads the same comm fields as the real one, so scenes set parts there.
+extern std::function<int(struct ncclComm*, ncclFunc_t)> g_rcclP2pTaskParts;
+
 void ResetEnqueueFakes();
 
 #endif  // RCCL_TEST_HOST_FAKES_ENQUEUE_FAKES_H_