[RCCL] Pick the gfx1250 tuning table by compute-partition mode (RCCL_TUNING_PARTITION_MODE)

projects/rccl/CHANGELOG.md

@@ -24,6 +24,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 * Hierarchical symmetric `ncclAllGather` for multi-node using GIN and LSA (`AllGather_RailRing_LsaST`): a GIN ring over the rails team moves data between nodes and LSA stores broadcast it to the other GPUs in each node.
 * nccl4py: per-call collective configuration (`NCCLCollConfig`, `VendorOption`) including the launch completion event, communicator properties (`NCCLCommProperties`), and the `GIN_ONLY` window flag. On ROCm, the HIP `Event` shim provides only an event handle; `record()`, `sync()` and `query` are not implemented.
 * `RCCL_TUNER_EMBEDDED_CONFIG` (default `1`): the built-in CSV tuner configurations in `tuner/*.csv` are now compiled into `librccl.so` and used as the stock tuning defaults, so a cluster no longer has to stage the CSV on every node or on NFS. Set `0` to ignore the embedded configurations. `NCCL_TUNER_CONFIG_FILE` and the on-disk tuner directories still take precedence, and `NCCL_TUNER_PLUGIN=none` still disables the tuner entirely.
+* `RCCL_TUNING_PARTITION_MODE` (`auto`, `SPX` or `DPX`; default `auto`): selects the gfx1250 tuning table for the GPU's compute-partition mode. `auto` uses the DPX table when the GPU reports DPX, and the SPX table otherwise. A communicator uses the DPX table only when every rank selects it. The DPX table is currently a copy of the SPX table, so tuning is unchanged.
 
 ### Changed
 * `install.sh --all_unrolls` (`-DBUILD_ALL_UNROLLS=ON`) now adds the unroll factors the multi-arch default skips (8 and 16) on top of each architecture's existing set. It no longer replaces that set or clears the gfx1250 pin on unroll 32. The flag stays off by default, so a normal build does not compile unroll 8 or 16.

projects/rccl/docs/api-reference/env-variables.rst

@@ -840,6 +840,19 @@ threshold table introduced for gfx1250.
         | ``1``: Ignore the table; fall back to pre-table compile-time constants.
           This matches the behavior before the arch table was introduced.
 
+    * - | ``RCCL_TUNING_PARTITION_MODE``
+        | Selects which gfx1250 per-architecture table (SPX or DPX) is used for
+          tuning thresholds. It only changes tuning, not results. A communicator
+          uses the DPX table only when every rank selects DPX; otherwise every rank
+          uses the SPX table. Other architectures have one table for every mode.
+          Has no effect when ``RCCL_IGNORE_ARCH_TABLE=1``.
+      - | ``auto`` (default): Follow the compute-partition mode the GPU reports.
+          ``DPX`` selects the DPX table; ``SPX``, ``CPX``, or no reported mode
+          selects the SPX table. Any other value is reported once and treated
+          as ``auto``.
+        | ``SPX``: Use the SPX table.
+        | ``DPX``: Use the DPX table.
+
     * - | ``RCCL_DDA_THRESHOLD``
         | Upper bound in bytes for the DDA VMM/Simple tier per collective.
           Messages above this size exit DDA and fall through to Ring, CE, or the

projects/rccl/src/include/comm.h

@@ -645,6 +645,9 @@ struct ncclPeerInfo {
   bool gpuCftMulticastSupport;
   bool gpuCftCountedSupport;
   uint32_t gitVersionHash;
+  // [RCCL] The rcclTuningPartition this rank asks for (RCCL_TUNING_PARTITION_MODE). The comm
+  // uses DPX only when every rank asks for it. Append later fields after this one.
+  int tuningPartition;
 };
 
 typedef enum ncclGroupTaskType {

projects/rccl/src/include/rccl_arch_thresholds.h

@@ -129,6 +129,14 @@ struct rcclArchThresholds {
   bool forceCeScratch;
 };
 
-const rcclArchThresholds* rcclGetArchThresholds(const char* gcn);
+// Compute-partition mode a gfx1250 comm is tuned for (RCCL_TUNING_PARTITION_MODE). It picks
+// the table only; other archs have one table for every mode.
+enum rcclTuningPartition {
+  RCCL_TUNING_PARTITION_SPX = 0,
+  RCCL_TUNING_PARTITION_DPX = 1,
+};
+
+const rcclArchThresholds* rcclGetArchThresholds(const char* gcn,
+                                                rcclTuningPartition partition = RCCL_TUNING_PARTITION_SPX);
 
 #endif  // RCCL_ARCH_THRESHOLDS_H_

projects/rccl/src/init.cc

@@ -1264,6 +1264,46 @@ static const char* mnnvlDegradedBwStr(unsigned int healthMask) {
 }
 #endif
 
+// RCCL_TUNING_PARTITION_MODE: the gfx1250 tuning table this rank asks for. "auto" (or unset)
+// follows the compute-partition mode the GPU reports; "SPX" / "DPX" force one. It only picks
+// tuning thresholds. Any mode other than DPX (CPX, or none reported) asks for SPX. An invalid
+// value is reported once and treated as "auto".
+static rcclTuningPartition rcclTuningPartitionRequest(const char* computePartition) {
+  const char* mode = ncclGetEnv("RCCL_TUNING_PARTITION_MODE");
+  if (mode != nullptr && mode[0] != '\0' && strcasecmp(mode, "auto") != 0) {
+    if (strcasecmp(mode, "SPX") == 0) return RCCL_TUNING_PARTITION_SPX;
+    if (strcasecmp(mode, "DPX") == 0) return RCCL_TUNING_PARTITION_DPX;
+    static std::once_flag invalidModeOnce;
+    std::call_once(invalidModeOnce, [mode] {
+      ATTN("Invalid value %s for RCCL_TUNING_PARTITION_MODE (expected auto, SPX or DPX), using auto", mode);
+    });
+  }
+  return computePartition != nullptr && strcasecmp(computePartition, "DPX") == 0 ? RCCL_TUNING_PARTITION_DPX
+                                                                                  : RCCL_TUNING_PARTITION_SPX;
+}
+
+// Ranks pick algorithms on their own, so every rank of a comm must read the same table: DPX only
+// when all of them ask for it, SPX otherwise. *mixed reports ranks that disagree.
+static rcclTuningPartition rcclTuningPartitionForComm(const struct ncclPeerInfo* peers, int nRanks, bool* mixed) {
+  int nDpx = 0;
+  for (int i = 0; i < nRanks; i++) nDpx += peers[i].tuningPartition == RCCL_TUNING_PARTITION_DPX;
+  *mixed = nDpx > 0 && nDpx < nRanks;
+  return nDpx == nRanks ? RCCL_TUNING_PARTITION_DPX : RCCL_TUNING_PARTITION_SPX;
+}
+
+// Re-picks comm->archThresholds once every rank's request is known (after AllGather1).
+static void rcclSetTuningPartition(struct ncclComm* comm) {
+  bool mixed = false;
+  const rcclTuningPartition partition = rcclTuningPartitionForComm(comm->peerInfo, comm->nRanks, &mixed);
+  comm->archThresholds = rcclGetArchThresholds(comm->archName, partition);
+  if (comm->rank != 0 || !IsArchMatch(comm->archName, "gfx1250")) return;
+  if (mixed) {
+    ATTN("gfx1250 tuning table: ranks ask for different tables (RCCL_TUNING_PARTITION_MODE or GPU partition mode), "
+         "using SPX on every rank");
+  }
+  INFO(NCCL_INIT | NCCL_TUNING, "gfx1250 tuning table: %s", partition == RCCL_TUNING_PARTITION_DPX ? "DPX" : "SPX");
+}
+
 static ncclResult_t fillInfo(struct ncclComm* comm, struct ncclPeerInfo* info, uint64_t commHash) {
   cudaDeviceProp prop;
   info->rank = comm->rank;
@@ -1311,18 +1351,19 @@ static ncclResult_t fillInfo(struct ncclComm* comm, struct ncclPeerInfo* info, u
   // partition 0 are both function 0 with an accelerator class, so read the mode from sysfs and
   // leave mloPart undefined on an unpartitioned GPU, which must not get the DEV overlay (it breaks
   // Rome gpuId matching and disables GIN/GDR).
+  // The partition mode lives on the physical device. A CPX alias at .1-.7 is usually absent
+  // from sysfs entirely, so ask function 0 rather than our own BDF.
+  char physBusIdStr[NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE];
+  char partition[MAX_STR_LEN];
+  partition[0] = '\0';
+  if (int64ToBusId(info->busId & ~0xfLL, physBusIdStr) == ncclSuccess) {
+    (void)ncclOsGetPciDeviceComputePartitionByBusId(physBusIdStr, partition, sizeof(partition));
+  }
+  info->tuningPartition = rcclTuningPartitionRequest(partition);
   if (info->mloPart == NCCL_TOPO_UNDEF) {
     int fn = (int)(info->busId & 0xf);
     if (fn < NCCL_TOPO_MLOPART_DEV_MAX) {
       char busIdStr[NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE];
-      // The partition mode lives on the physical device. A CPX alias at .1-.7 is usually absent
-      // from sysfs entirely, so ask function 0 rather than our own BDF.
-      char physBusIdStr[NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE];
-      char partition[MAX_STR_LEN];
-      partition[0] = '\0';
-      if (int64ToBusId(info->busId & ~0xfLL, physBusIdStr) == ncclSuccess) {
-        (void)ncclOsGetPciDeviceComputePartitionByBusId(physBusIdStr, partition, sizeof(partition));
-      }
       // A partitioned device makes every function a partition, function 0 included: partition 0 is
       // a partition, not an unpartitioned GPU, and must carry index 0 so all of a device's
       // partitions share one DEV overlay group. An empty string means the platform does not report
@@ -1854,6 +1895,7 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p
       }
     }
   }
+  rcclSetTuningPartition(comm);
 
   // AllGather1 - end
   timers[TIMER_INIT_ALLGATHER] = clockNano() - timers[TIMER_INIT_ALLGATHER];
@@ -3109,6 +3151,7 @@ static ncclResult_t ncclCommInitRankFunc(struct ncclAsyncJob* job_) {
   comm->ll128LineElems = rcclLL128LineElemsFromArch(comm->archName);
   comm->ll128DataElems = rcclLL128DataElemsFromArch(comm->archName);
   comm->ll128ShmemElemsPerThread = rcclLL128ShmemElemsPerThreadFromArch(comm->archName);
+  // SPX until initTransportsRank knows every rank's RCCL_TUNING_PARTITION_MODE request (rcclSetTuningPartition).
   comm->archThresholds = rcclGetArchThresholds(comm->archName);
   comm->maxSharedMemOptin = maxSharedMem;