From 4be06f04d847535e27b4d37d2a5b3a46eb2528aa Mon Sep 17 00:00:00 2001
From: Mustafa Abduljabbar <mustafa.abduljabbar@amd.com>
Date: Wed, 2 Apr 2025 16:26:23 -0400
Subject: [PATCH] Add AllGather LL128 multi-node tuning and include LL cutoff
 points in tuning models (#1618)

* Enable LL/LL128 cutoff points in tuning models

* Initializing ll/ll128 model cutoffs for MI300

* Use RCCL_LL_LIMITS_UNDEFINED

---------

Co-authored-by: PedramAlizadeh <pmohamma@amd.com>
---
 src/enqueue.cc      | 53 +++++++++++++++++++++++++++------------------
 src/graph/topo.h    |  7 ++++++
 src/graph/tuning.cc | 28 +++++++++++++++++++++++-
 src/include/comm.h  |  3 ++-
 4 files changed, 68 insertions(+), 23 deletions(-)

diff --git a/src/enqueue.cc b/src/enqueue.cc
index cc9d4e1665..d2d80fc60c 100644
--- a/src/enqueue.cc
+++ b/src/enqueue.cc
@@ -1903,14 +1903,6 @@ static ncclResult_t updateCollCostTable(
   return ncclSuccess;
 }
 
-// The following parameters are based on the observation that LL and LL128
-// performance is determined by how much data goes out of a GPU
-// TODO: these numbers can be part of topo detection.
-RCCL_PARAM(RsLLMinSizePerRank,    "REDUCE_SCATTER_LL_MIN_SIZE_PER_RANK", 0);
-RCCL_PARAM(RsLLMaxSizePerRank,    "REDUCE_SCATTER_LL_MAX_SIZE_PER_RANK", 655360);
-RCCL_PARAM(RsLL128MinSizePerRank, "REDUCE_SCATTER_LL128_MIN_SIZE_PER_RANK", 131072);
-RCCL_PARAM(RsLL128MaxSizePerRank, "REDUCE_SCATTER_LL128_MAX_SIZE_PER_RANK", 3211264);
-
 static ncclResult_t topoGetAlgoInfo(
     struct ncclComm* comm, struct ncclTaskColl* info, size_t nBytes,
     float** collCostTable, int backupAlgo, int backupProto, float backupTime, ncclSimInfo_t* simInfo
@@ -1953,23 +1945,42 @@ static ncclResult_t topoGetAlgoInfo(
     const char *protoStr = getenv("NCCL_PROTO");
     userProtocolInput = !protoStr ? 0 : 1;
   }
-  if(!userProtocolInput && comm->nNodes >= 2 && info->func == ncclFuncReduceScatter && IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx942")) {
-    // Keep it simple unless otherwise required
-    info->protocol = NCCL_PROTO_SIMPLE;
-    // Normalize the comparison to sizePerRank as this is essentially what matters in determining protocol choice
-    size_t sizePerRank = nBytes / comm->nRanks;
 
-    if(sizePerRank <= rcclParamRsLLMaxSizePerRank() && sizePerRank >= rcclParamRsLLMinSizePerRank()) {
-      info->protocol = NCCL_PROTO_LL;
-    }
+  if(!userProtocolInput && comm->nNodes >= 2 && (info->func == ncclFuncReduceScatter || info->func == ncclFuncAllGather)) {
+    auto llMin = comm->minMaxLLRange[info->func][NCCL_PROTO_LL][0];
+    auto llMax = comm->minMaxLLRange[info->func][NCCL_PROTO_LL][1];
+
+    auto ll128Min = comm->minMaxLLRange[info->func][NCCL_PROTO_LL128][0];
+    auto ll128Max = comm->minMaxLLRange[info->func][NCCL_PROTO_LL128][1];
+
+    // Only override model choices if min/max cutoff points are set in the tuning models
+    if((ll128Max != RCCL_LL_LIMITS_UNDEFINED) || (llMax != RCCL_LL_LIMITS_UNDEFINED)) {
+      // Keep it simple unless otherwise required
+      info->protocol = NCCL_PROTO_SIMPLE;
+      // Normalize the comparison to sizePerRank as this is essentially what matters in determining protocol choice
+      size_t sizePerRank = nBytes / comm->nRanks;
+
+      if(sizePerRank <= llMax && sizePerRank > llMin) {
+        info->protocol = NCCL_PROTO_LL;
+      }
 #if defined(ENABLE_LL128)
-    // LL128 RS performance is better than LL when enabled, so the next condition overrides the previous LL choice
-    if(comm->topo->ll128Enabled) {
-      if(sizePerRank <= rcclParamRsLL128MaxSizePerRank() && sizePerRank >= rcclParamRsLL128MinSizePerRank()) {
-        info->protocol = NCCL_PROTO_LL128;
+      // When applicable, LL128 RS performance is better than LL, so the next condition overrides the previous LL choice
+      if(comm->topo->ll128Enabled) {
+        if(sizePerRank <= ll128Max && sizePerRank > ll128Min) {
+          info->protocol = NCCL_PROTO_LL128;
+        }
+      }
+#endif
+    } else if (IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx942")) {
+      // Warn that model detection for MI300 (or future others) did not work as expected
+      // Add supported archs to this condition as they come (e.g. gfx950)
+      // Also make sure the tuning_model and model detection are updated for new archs
+      static bool failedWarn = false;
+      if (!failedWarn) {
+        WARN("LL cutoff points not detected for a supported arch %s", comm->topo->nodes[GPU].nodes[0].gpu.gcn);
+        failedWarn = true;
       }
     }
-#endif
   }
 #endif
   if (comm->rank == 0) INFO(NCCL_TUNING, "%s: %ld Bytes -> Algo %d proto %d time %f", ncclFuncToString(info->func), nBytes, info->algorithm, info->protocol, time);
diff --git a/src/graph/topo.h b/src/graph/topo.h
index 4e65f2b472..e488bd1e0a 100644
--- a/src/graph/topo.h
+++ b/src/graph/topo.h
@@ -123,6 +123,13 @@ struct ncclTopoLinkList {
 #define RCCL_TOPO_FORCE_INTRA 16
 #define RCCL_TOPO_XGMI_ALL  32
 
+#define RCCL_LL_TUNABLE_COLLS 4 // LL/LL64/LL128 tunable Collectives
+#define RCCL_RS_TUNABLE 0       // reduce_scatter index
+#define RCCL_AG_TUNABLE 1       // all_gather index
+#define RCCL_AR_TUNABLE 2       // all_reduce index
+#define RCCL_RE_TUNABLE 3       // reduce index
+#define RCCL_LL_LIMITS_UNDEFINED 0
+
 #define GCN_ARCH_NAME_LEN 16
 
 struct ncclTopoNode {
diff --git a/src/graph/tuning.cc b/src/graph/tuning.cc
index e89592db88..c550e7ae05 100644
--- a/src/graph/tuning.cc
+++ b/src/graph/tuning.cc
@@ -64,11 +64,14 @@ static const float baseLat  [NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS] = {
 #define NCCL_HW_PCI 1
 #define NCCL_HW_NET 2
 
+
+
 struct tuningModel {
   float hwLat [3][NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
   float bwRatio [2][NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
   float treeCorrectionFactor[NCCL_NUM_PROTOCOLS][27];
   float ringCorrectionFactor[NCCL_NUM_PROTOCOLS][27];
+  uint64_t llProtoRanges[RCCL_LL_TUNABLE_COLLS][NCCL_NUM_PROTOCOLS - 1][2];
 };
 
 static struct tuningModel tuning_model_0 {
@@ -99,6 +102,8 @@ static struct tuningModel tuning_model_0 {
     { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.3, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.8, 0.7, 0.5, 0.4, 0.4, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, },
     { 1.0, 0.8, 0.2, 1.0, 1.0, 0.3, 1.0, 0.1, 0.1, 0.2, 0.2, 0.1, 0.5, 1.0, 0.8, 0.8, 1.0, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, },
   },
+
+  .llProtoRanges = {RCCL_LL_LIMITS_UNDEFINED},
 };
 
 static struct tuningModel tuning_model_1 {
@@ -129,6 +134,8 @@ static struct tuningModel tuning_model_1 {
     { 1.0, 0.5, 1.0, 1.0, 0.6, 0.7, 1.0, 1.0, 0.2, 1.0, 0.9, 0.7, 1.0, 1.0, 1.0, 0.9, 0.9, 0.8, 0.8, 0.7, 0.6, 0.5, 0.5, 0.3, 0.2, 0.1, 0.1, },
     { 0.3, 1.0, 0.3, 0.1, 0.1, 0.1, 0.3, 0.7, 1.0, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.2, 0.3, 0.5, 0.9, 1.0, 1.0, 1.0, 1.0, },
   },
+
+  .llProtoRanges = {RCCL_LL_LIMITS_UNDEFINED},
 };
 
 static struct tuningModel tuning_model_2 {
@@ -159,6 +166,8 @@ static struct tuningModel tuning_model_2 {
     { 0.1, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.4, 1.0, 1.0, 1.0, 1.0, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, },
     { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.2, 0.4, 0.5, 0.6, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, },
   },
+
+  .llProtoRanges = {RCCL_LL_LIMITS_UNDEFINED},
 };
 
 static struct tuningModel tuning_model_3 {
@@ -189,6 +198,8 @@ static struct tuningModel tuning_model_3 {
     { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, },
     { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.5, 1.0, 0.1, 0.3, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.3, 0.4, 0.7, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, },
   },
+
+  .llProtoRanges = {RCCL_LL_LIMITS_UNDEFINED},
 };
 
 static struct tuningModel tuning_model_4 {
@@ -219,6 +230,8 @@ static struct tuningModel tuning_model_4 {
     { 0.4, 0.5, 0.5, 0.4, 0.4, 0.4, 0.4, 0.2, 0.2, 0.1, 0.3, 1.0, 1.0, 0.7, 0.8, 0.5, 1.0, 1.0, 1.0, 1.0, 1.0, 0.9, 0.8, 0.5, 0.4, 0.3, 0.3, },
     { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0, 1.0, 0.8, 0.5, 0.1, 0.7, 0.2, 0.4, 0.4, 0.6, 0.7, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, },
   },
+
+  .llProtoRanges = {RCCL_LL_LIMITS_UNDEFINED},
 };
 
 static struct tuningModel tuning_model_5 {
@@ -249,6 +262,11 @@ static struct tuningModel tuning_model_5 {
     { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0, 0.2, 1.0, 0.4, 0.4, 0.1, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, },
     { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.8, 1.0, 1.0, 1.0, },
   },
+
+  .llProtoRanges = {
+    /*ReduceScatter*/ {/* LL (Min/Max) */ {0, 655360} , /* LL128 (Min/Max) */ {131072, 3211264}},
+    /*AllGather*/     {/* LL (Min/Max) */ {0, 98304} , /* LL128 (Min/Max) */ {98304, 5046272}},
+  },
 };
 
 static struct tuningModel rcclTuningModel[] = {
@@ -354,6 +372,14 @@ ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCom
   for (int a=0; a<NCCL_NUM_ALGORITHMS; a++) intraHw[a] = graphs[a]->typeIntra == LINK_NVL ? NCCL_HW_NVLINK : NCCL_HW_PCI;
   for (int a=0; a<NCCL_NUM_ALGORITHMS; a++) hw[a] = nNodes == 1 ? intraHw[a] : NCCL_HW_NET;
 
+  memcpy(comm->minMaxLLRange[ncclFuncReduceScatter],
+        rcclTuningModel[comm->topo->tuning].llProtoRanges[RCCL_RS_TUNABLE],
+        sizeof(rcclTuningModel[comm->topo->tuning].llProtoRanges[RCCL_RS_TUNABLE]));
+
+  memcpy(comm->minMaxLLRange[ncclFuncAllGather],
+        rcclTuningModel[comm->topo->tuning].llProtoRanges[RCCL_AG_TUNABLE],
+        sizeof(rcclTuningModel[comm->topo->tuning].llProtoRanges[RCCL_AG_TUNABLE]));
+
   for (int coll=0; coll<NCCL_NUM_FUNCTIONS; coll++) {
     int nsteps = coll == ncclFuncAllReduce ? 2*(nRanks-1) :
       coll == ncclFuncReduceScatter || coll == ncclFuncAllGather ? nRanks-1 :
@@ -690,4 +716,4 @@ ncclResult_t ncclTopoGetAlgoTime(struct ncclComm* comm, int coll, int algorithm,
   int latCount = algorithm == NCCL_ALGO_RING ? numPipeOps : DIVUP(numPipeOps, NCCL_MAX_DEV_WORK_BATCH_COLLS);
   *time = lat * latCount + nBytes / (1000 * bw);
   return ncclSuccess;
-}
\ No newline at end of file
+}
diff --git a/src/include/comm.h b/src/include/comm.h
index f5241cf3ec..52c98d4bc0 100644
--- a/src/include/comm.h
+++ b/src/include/comm.h
@@ -517,6 +517,7 @@ struct ncclComm {
   float bandwidths[NCCL_NUM_FUNCTIONS][NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
   float ringbdw[NCCL_NUM_FUNCTIONS][NCCL_NUM_PROTOCOLS];
   int maxThreads[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
+  uint64_t minMaxLLRange[NCCL_NUM_FUNCTIONS][NCCL_NUM_PROTOCOLS - 1][2];
 
   /* This attribute can indicate the states of communicators and return code of
   * asynchronous NCCL operations. */
@@ -599,7 +600,7 @@ struct ncclComm {
   struct ncclKernelPlanner planner;
 
   hipStream_t sideStream; // [RCCL] Cached non-captured stream
-  
+
   cudaMemPool_t memPool;
   // Queue of events and associated callbacks for cleaning up asynchronous work.
   // Using this is preferable to using CUDA host callbacks because host callbacks