WarpSpeed Refactor and Gating to SPX

projects/rccl/src/enqueue.cc

@@ -2474,16 +2474,7 @@ static ncclResult_t topoGetAlgoInfo(struct ncclComm* comm, struct ncclTaskColl*
   TRACE(NCCL_COLL, "%ld Bytes -> Algo %d proto %d time %f", nBytes, info->algorithm, info->protocol, time);
   int nc = comm->nChannels;
 #ifdef ENABLE_WARP_SPEED
-  if (comm->topo->warpSpeedEnabled) {
-    nc /= comm->warpSpeedChannelMultiplier;
-    // Temporary check as we reduce CU usage for all collectives
-    // TODO: Remove this condition after optimizing all collectives
-    if (IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx950") && comm->nNodes == 1 && comm->nRanks == 8 &&
-        info->func != ncclFuncAllReduce && info->func != ncclFuncAllGather && info->func != ncclFuncReduceScatter &&
-        ncclParamMaxNchannels() < 0) {
-      nc *= 2;
-    }
-  }
+  nc = rcclWarpSpeedAdjustChannels(comm, info, nc);
 #endif
   int nt = comm->maxThreads[info->algorithm][info->protocol];
   int threadThreshold = comm->threadThresholds[info->algorithm][info->protocol];

projects/rccl/src/graph/connect.cc

@@ -1198,23 +1198,8 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
 #ifdef ENABLE_WARP_SPEED
   channelMultiplier = comm->warpSpeedChannelMultiplier = wsEnabled ? rcclGetMaxWarpsPerBlock(comm) : 1;
   if (wsEnabled) {
-    // If user didn't override, use requested channels; otherwise keep capped max.
-    if (!userUpdatedMaxChannels) {
-      maxNchannels = nc * comm->nChannels * channelMultiplier;
-      // Cap the single-node count at MAXCHANNELS/2, the maximum WarpSpeed supports. Left unbounded it
-      // overruns the fixed comm->channels[MAXCHANNELS] array (for example nc reaches 1760 in CPX mode at 64 ranks).
-      nc = singleNode ? std::min(maxNchannels, MAXCHANNELS / 2) : std::min(maxNchannels, maxChannels);
-    } else {
-      nc = maxNchannels = std::min(adjustedMaxNchannels * channelMultiplier, MAXCHANNELS);
-    }
-
-    if (!userUpdatedMaxChannels && isGfx950 && singleNode && comm->nRanks == 8) {
-      // For gfx950 single-node, use half the channels since they are doubled on a single node
-      // Remove when all collectives have been optimized
-      nc /= 2;
-    }
-    INFO(NCCL_TUNING, "WarpSpeed enabled: warpSpeedChannelMultiplier %d, maxNchannels %d, nc %d", channelMultiplier,
-         maxNchannels, nc);
+    nc = rcclWarpSpeedComputeNChannels(comm, nc, channelMultiplier, maxChannels, adjustedMaxNchannels,
+                                       userUpdatedMaxChannels);
   } else {
     maxChannels = std::min(
       (singleNode ? (isGfx950 ? RCCL_MI3XX_MAX_SINGLE_NODE_CHANNELS * 2 : RCCL_MI3XX_MAX_SINGLE_NODE_CHANNELS) :

projects/rccl/src/include/rccl_common.h

@@ -158,5 +158,8 @@ bool rcclWarpSpeedSupported(struct ncclComm* comm, struct ncclKernelPlan* plan);
 ncclResult_t rcclSetWarpSpeedAuto(struct ncclComm* comm, struct ncclTaskColl* info, size_t nBytes);
 int rcclGetMaxWarpsPerBlock(struct ncclComm* comm);
 bool rcclCanUseWarpSpeedAuto(struct ncclComm* comm, int nNodes);
+int rcclWarpSpeedComputeNChannels(struct ncclComm* comm, int nc, int channelMultiplier, int maxChannels,
+                                  int adjustedMaxNchannels, bool userUpdatedMaxChannels);
+int rcclWarpSpeedAdjustChannels(struct ncclComm* comm, struct ncclTaskColl* info, int nc);
 #endif
 #endif

projects/rccl/src/init.cc

@@ -1339,6 +1339,13 @@ static ncclResult_t ncclP2pSchedule(struct ncclComm* comm) {
   return ncclSuccess;
 }
 
+#ifdef ENABLE_WARP_SPEED
+static bool willEnableWarpSpeed(struct ncclComm* parent, struct ncclComm* comm, int nNodes) {
+  return rcclParamWarpSpeedForceEnable() > 0 ||
+         ((!parent || comm->isGrow) && rcclCanUseWarpSpeedAuto(comm, nNodes));
+}
+#endif
+
 static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* parent,
                                        uint64_t timers[TIMERS_INIT_COUNT]) {
   // We use 2 AllGathers
@@ -1787,8 +1794,7 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p
     }
   }
 #ifdef ENABLE_WARP_SPEED
-  comm->topo->warpSpeedEnabled =
-    (rcclParamWarpSpeedForceEnable() > 0 || ((!parent || comm->isGrow) && rcclCanUseWarpSpeedAuto(comm, nNodes)));
+  comm->topo->warpSpeedEnabled = willEnableWarpSpeed(parent, comm, nNodes);
 #endif
 
   // For single node communicators that do not uses the full xgmi links per gpu, i.e., nranks < 8

projects/rccl/src/rccl_wrap.cc

@@ -743,7 +743,7 @@ bool rcclIsAboveWarpSpeedThreshold(struct ncclComm* comm, struct ncclTaskColl* i
 
 bool rcclCanUseWarpSpeedAuto(struct ncclComm* comm, int nNodes) {
   return IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx950") && (nNodes == 1) &&
-         (rcclParamWarpSpeedAutoMode() != 0);
+         (rcclParamWarpSpeedAutoMode() != 0) && comm-> cuCount > 128; // Only use in SPX mode, 256 CU on gfx950
 }
 
 ncclResult_t validChannelsForWarpSpeed(struct ncclComm* comm, struct ncclTaskColl* info) {
@@ -805,6 +805,47 @@ int rcclGetMaxWarpsPerBlock(struct ncclComm* comm) {
   }
   return warpsPerBlock;
 }
+
+// Compute the bandwidth channel count (nc) when WarpSpeed is enabled, scaling the
+// base channel count by the per-block warp multiplier. 
+int rcclWarpSpeedComputeNChannels(struct ncclComm* comm, int nc, int channelMultiplier, int maxChannels,
+                                  int adjustedMaxNchannels, bool userUpdatedMaxChannels) {
+  const bool singleNode = comm->nNodes == 1;
+  const bool isGfx950 = IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx950");
+  int maxNchannels;
+  // If user didn't override, use requested channels; otherwise keep capped max.
+  if (!userUpdatedMaxChannels) {
+    maxNchannels = nc * comm->nChannels * channelMultiplier;
+    nc = singleNode ? maxNchannels : std::min(maxNchannels, maxChannels);
+  } else {
+    nc = maxNchannels = std::min(adjustedMaxNchannels * channelMultiplier, MAXCHANNELS);
+  }
+
+  if (!userUpdatedMaxChannels && isGfx950 && singleNode && comm->nRanks == 8) {
+    // For gfx950 single-node, use half the channels since they are doubled on a single node
+    // Remove when all collectives have been optimized
+    nc /= 2;
+  }
+  INFO(NCCL_TUNING, "WarpSpeed enabled: warpSpeedChannelMultiplier %d, maxNchannels %d, nc %d", channelMultiplier,
+       maxNchannels, nc);
+  return nc;
+}
+
+// Adjust the per-collective channel count (nc) for WarpSpeed during algo/channel
+// tuning. No-op when WarpSpeed is disabled. 
+int rcclWarpSpeedAdjustChannels(struct ncclComm* comm, struct ncclTaskColl* info, int nc) {
+  if (comm->topo->warpSpeedEnabled) {
+    nc /= comm->warpSpeedChannelMultiplier;
+    // Temporary check as we reduce CU usage for all collectives
+    // TODO: Remove this condition after optimizing all collectives
+    if (IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx950") && comm->nNodes == 1 && comm->nRanks == 8 &&
+        info->func != ncclFuncAllReduce && info->func != ncclFuncAllGather && info->func != ncclFuncReduceScatter &&
+        ncclParamMaxNchannels() < 0) {
+      nc *= 2;
+    }
+  }
+  return nc;
+}
 #endif
 
 void rcclGetMaxNthreads(struct ncclComm* comm, int maxNthreads[]) {