MSCCL: Improve executor and integrate scheduler (#694)

* MSCCL: improve executor and add scheduler for testing

* Use external scheduler

* Fix cmake error

* Address comments

* Fix thread safe issue

* Make MSCCL lifecycle APIs thread safe

* Make MSCCL internal scheduler aware of topology hint

* Revise error message

src/collectives/device/msccl_kernel.cu

@@ -110,6 +110,19 @@ __device__ __forceinline__ static void threadBlockCopy(
   }
 }
 
+#define MSCCL_REDUCE_UNROLL_LOOP_A(numloops) \
+for (int r = 0; r < numloops; r++) { \
+  srcOffset = srcBaseOffset + (ssize_t)mscclShmem.mscclTB.reductionSrcOffsets[t->reductionPointer+r] * sizePerMscclChunk; \
+  reduceInput = load(srcPointer + srcOffset); \
+  o = redFn(reduceInput, o); \
+}
+
+#define MSCCL_REDUCE_UNROLL_LOOP_B(numloops) \
+for (int r = 0; r < numloops; r++) { \
+  srcOffset = srcBaseOffset + (ssize_t)mscclShmem.mscclTB.reductionSrcOffsets[t->reductionPointer+r] * sizePerMscclChunk; \
+  srcs[r] = srcPointer + srcOffset; \
+}
+
 template<typename T, typename RedOp, typename Proto>
 __device__ __forceinline__ void mscclRunInterpreter(
   struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork work) {
@@ -137,7 +150,7 @@ __device__ __forceinline__ void mscclRunInterpreter(
   int channelId = mscclShmem.mscclTB.channelId;
   {
     void *dst, *src;
-    int bytes;
+    int bytes = 0;
     // Use first 3 warps to load comm, channel, and work into shmem
     switch (tid/WARP_SIZE) {
     case 0:
@@ -159,8 +172,11 @@ __device__ __forceinline__ void mscclRunInterpreter(
       bytes = sizeof(mscclWork);
       static_assert(sizeof(mscclWork) <= sizeof(uint64_t) * WARP_SIZE, "mscclWork cannot be loaded by a single warp in one insn.");
       break;
+    case 3:
+      /* set abort flag to 0 */
+      if (tid == 3 * WARP_SIZE) ncclShmem.aborted = 0;
+      break;
     default:
-      bytes = 0;
       break;
     }
     copyToShmem8(tid%WARP_SIZE, dst, src, bytes);
@@ -264,11 +280,76 @@ __device__ __forceinline__ void mscclRunInterpreter(
             if (tid < thisNelem){
               dstOffset = gridOffset + (ssize_t) (t->dstOffset+c) * sizePerMscclChunk;
               T* dstIndex = dstPointer + dstOffset + tid;
+              T reduceInput;
               T o = load(dstIndex);
-              for (int r = 0; r < numReductions; r++){
-                srcOffset = gridOffset + (ssize_t) (mscclShmem.mscclTB.reductionSrcOffsets[t->reductionPointer+r]+c) * sizePerMscclChunk;
-                T t = load(srcPointer + srcOffset + tid);
-                o = redFn(t,o);
+              ssize_t srcBaseOffset = gridOffset + (ssize_t)c * sizePerMscclChunk + tid;
+              switch (numReductions) {
+                case 1:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(1);
+                  break;
+                case 2:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(2);
+                  break;
+                case 3:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(3);
+                  break;
+                case 4:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(4);
+                  break;
+                case 5:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(5);
+                  break;
+                case 6:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(6);
+                  break;
+                case 7:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(7);
+                  break;
+                case 8:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(8);
+                  break;
+                case 9:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(9);
+                  break;
+                case 10:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(10);
+                  break;
+                case 11:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(11);
+                  break;
+                case 12:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(12);
+                  break;
+                case 13:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(13);
+                  break;
+                case 14:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(14);
+                  break;
+                case 15:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(15);
+                  break;
+                case 16:
+                  #pragma unroll
+                  MSCCL_REDUCE_UNROLL_LOOP_A(16);
+                  break;
+                default:
+                  break;
               }
               store(dstIndex, o);
             }
@@ -277,9 +358,74 @@ __device__ __forceinline__ void mscclRunInterpreter(
             T* srcs[MSCCL_MAX_REDUCE_FUSION+1]; // +1 is for SIMPLE protocol as dst is added in the list of srcs
             dstOffset = gridOffset + (ssize_t) (t->dstOffset+c) * sizePerMscclChunk;
             T* dst = dstPointer + dstOffset;
-            for (int r = 0; r < numReductions; r++) {
-              srcOffset = gridOffset + (ssize_t) (mscclShmem.mscclTB.reductionSrcOffsets[t->reductionPointer+r]+c) * sizePerMscclChunk;
-              srcs[r] = srcPointer + srcOffset;
+            ssize_t srcBaseOffset = gridOffset + (ssize_t)c * sizePerMscclChunk;
+            switch (numReductions) {
+              case 1:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(1);
+                break;
+              case 2:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(2);
+                break;
+              case 3:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(3);
+                break;
+              case 4:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(4);
+                break;
+              case 5:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(5);
+                break;
+              case 6:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(6);
+                break;
+              case 7:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(7);
+                break;
+              case 8:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(8);
+                break;
+              case 9:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(9);
+                break;
+              case 10:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(10);
+                break;
+              case 11:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(11);
+                break;
+              case 12:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(12);
+                break;
+              case 13:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(13);
+                break;
+              case 14:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(14);
+                break;
+              case 15:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(15);
+                break;
+              case 16:
+                #pragma unroll
+                MSCCL_REDUCE_UNROLL_LOOP_B(16);
+                break;
+              default:
+                break;
             }
             prims.reduce(srcs, numReductions, &dst, 1, thisNelem);
           }

src/collectives/msccl.cc

@@ -10,26 +10,26 @@
 #include <cstdio>
 #include <cstdlib>
 
-NCCL_API(ncclResult_t, mscclLoadAlgo, const char *mscclAlgoFilePath, mscclAlgoHandle_t *mscclAlgoHandle);
-ncclResult_t mscclLoadAlgo(const char *mscclAlgoFilePath, mscclAlgoHandle_t *mscclAlgoHandle) {
+NCCL_API(ncclResult_t, mscclLoadAlgo, const char *mscclAlgoFilePath, mscclAlgoHandle_t *mscclAlgoHandle, int rank);
+ncclResult_t mscclLoadAlgo(const char *mscclAlgoFilePath, mscclAlgoHandle_t *mscclAlgoHandle, int rank) {
   mscclStatus& status = mscclGetStatus();
 
   if (status.freeAlgoHandles.size() == 0) {
     WARN("MSCCL: MSCCL_MAX_NUM_ALGOS (%d) limit reached", MSCCL_MAX_NUM_ALGOS);
     return ncclInvalidUsage;
   }
-  mscclAlgoHandle_t handle = *status.freeAlgoHandles.rbegin();
+  *mscclAlgoHandle = *status.freeAlgoHandles.rbegin();
   status.freeAlgoHandles.pop_back();
 
   struct mscclAlgo* hostAlgo;
   NCCLCHECK(ncclCalloc(&hostAlgo, 1));
-  NCCLCHECK(mscclGetAlgoFromXmlFile(mscclAlgoFilePath, hostAlgo, status.rank));
-  status.hostAlgos[handle] = hostAlgo;
+  NCCLCHECK(mscclGetAlgoFromXmlFile(mscclAlgoFilePath, hostAlgo, rank));
+  status.hostAlgos[*mscclAlgoHandle] = hostAlgo;
 
   struct mscclAlgo* devAlgo;
   NCCLCHECK(ncclCudaCalloc(&devAlgo, 1));
   CUDACHECK(hipMemcpy(devAlgo, hostAlgo, sizeof(struct mscclAlgo), hipMemcpyHostToDevice));
-  status.devAlgos[handle] = devAlgo;
+  status.devAlgos[*mscclAlgoHandle] = devAlgo;
 
   return ncclSuccess;
 }
@@ -54,7 +54,10 @@ ncclResult_t mscclRunAlgo(
 
   NCCLCHECK(mscclSetupSyncFlags(stream));
 
-  NCCLCHECK(mscclSetupConnections(hostAlgo, comm));
+  if (status.connectedAlgos[comm].find(mscclAlgoHandle) == status.connectedAlgos[comm].end()) {
+    NCCLCHECK(mscclSetupConnections(hostAlgo, comm));
+    status.connectedAlgos[comm].insert(mscclAlgoHandle);
+  }
 
   NCCLCHECK(mscclSetupProxy(hostAlgo, comm));
 
@@ -75,5 +78,9 @@ ncclResult_t mscclUnloadAlgo(mscclAlgoHandle_t mscclAlgoHandle) {
 
   status.freeAlgoHandles.push_back(mscclAlgoHandle);
 
+  for (auto &s : status.connectedAlgos) {
+    s.second.erase(mscclAlgoHandle);
+  }
+
   return ncclSuccess;
 }

src/graph/rome_models.cc

@@ -392,7 +392,7 @@ static struct rcclRomeModel rome_model_56 = {
   .gdrLevel = { },
   .pattern = "40404040",
   .ringBase = "0 1 3 2 6 7 15 14 10 11 9 8 12 13 5 4|0 1 2 3 7 6 13 12 8 9 10 11 15 14 5 4|0 2 3 7 6 14 15 11 10 8 9 13 12 4 5 1|4 5 13 12 8 9 11 10 14 15 7 6 2 3 1 0|4 5 14 15 11 10 9 8 12 13 6 7 3 2 1 0|1 5 4 12 13 9 8 10 11 15 14 6 7 3 2 0",
-  .options = "pivotA2AEnabled=1,pivotA2ANumBiRings=3,tuning=1",
+  .options = "pivotA2AEnabled=1,pivotA2ANumBiRings=3,tuning=1,mscclEnabled=1",
   .treeBase = "10 11|14 15|6 7|2 3|0 1|4 5|12 13|8 9",
 };
 
@@ -842,6 +842,8 @@ static void parseOptions(struct ncclTopoSystem* system, const char *options) {
         system->ll128Enabled = (bool)atol(tokens[i*2+1]);
       } else if (strcmp(tokens[i*2], "baseBw") == 0) {
         system->baseBw = std::stof(tokens[i*2+1]);
+      } else if (strcmp(tokens[i*2], "mscclEnabled") == 0) {
+        system->mscclEnabled = (bool)atol(tokens[i*2+1]);
       }
     }
     free(str_temp);

src/graph/topo.h

@@ -167,6 +167,7 @@ struct ncclTopoSystem {
   int pivotA2ANumBiRings;
   bool ll128Enabled;
   float baseBw;
+  bool mscclEnabled;
 };
 
 ncclResult_t ncclTopoGetNode(struct ncclTopoSystem* system, struct ncclTopoNode** node, int type, uint64_t id);

src/include/comm.h

@@ -319,6 +319,9 @@ struct ncclComm {
   bool finalizeCalled;
   // shared structures for finalization
   int finalizeRankCnt;
+
+  // Whether this comm is compatible with MSCCL
+  bool mscclCompatible;
 };
 
 enum ncclLaunchMode {

src/include/msccl/msccl_parser.h

@@ -100,4 +100,6 @@ static ncclResult_t mscclXmlFindTag(struct mscclXml* xml, const char* tagName, s
 
 ncclResult_t mscclGetAlgoFromXmlFile(const char* xmlGraphFile, struct mscclAlgo* algo, int rank);
 
+ncclResult_t mscclGetAlgoMetaFromXmlFile(const char* xmlGraphFile, struct mscclAlgoMeta* algoMeta);
+
 #endif

src/include/msccl/msccl_scheduler.h

@@ -0,0 +1,52 @@
+/*************************************************************************
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT License.
+ ************************************************************************/
+
+#ifndef MSCCL_SCHEDULER_H_
+#define MSCCL_SCHEDULER_H_
+
+typedef enum { mscclFuncReduce             =  0,
+               mscclFuncBroadcast          =  1,
+               mscclFuncAllReduce          =  2,
+               mscclFuncReduceScatter      =  3,
+               mscclFuncAllGather          =  4,
+               mscclFuncSend               =  5,
+               mscclFuncRecv               =  6,
+               mscclFuncGather             =  7,
+               mscclFuncScatter            =  8,
+               mscclFuncAllToAll           =  9,
+               mscclFuncAllToAllv          =  10,
+               mscclNumFuncs               =  11 } mscclFunc_t;
+
+struct mscclSchedulerParam {
+  const void* sendBuff;
+  const size_t* sendCounts;
+  const size_t* sDisPls;
+  void* recvBuff;
+  const size_t* recvCounts;
+  const size_t* rDisPls;
+  size_t count;
+  ncclDataType_t dataType;
+  int root;
+  int peer;
+  ncclRedOp_t op;
+  mscclFunc_t func;
+  int rank;
+  int nRanks;
+  bool scheduled;
+  mscclAlgoHandle_t handle;
+};
+
+typedef struct {
+  // Name of the scheduler (mainly for logs)
+  const char* name;
+  // Load all algorithms
+  ncclResult_t (*init)();
+  // Select an algorithm
+  ncclResult_t (*selectAlgo)(struct mscclSchedulerParam* param);
+  // Unload all algorithms
+  ncclResult_t (*teardown)();
+} mscclSchedulerInterface;
+
+#endif

src/include/msccl/msccl_status.h

@@ -10,4 +10,6 @@
 
 mscclStatus& mscclGetStatus();
 
+mscclThreadLocalStatus& mscclGetThreadLocalStatus();
+
 #endif

src/include/msccl/msccl_struct.h

@@ -11,8 +11,9 @@
 #include <set>
 #include <vector>
 #include "devcomm.h"
+#include "msccl/msccl_scheduler.h"
 
-#define MSCCL_MAX_NUM_STEPS 256
+#define MSCCL_MAX_NUM_STEPS 64
 #define MSCCL_MAX_NUM_THREAD_BLOCKS_PER_CHANNEL 32
 #define MSCCL_MAX_NUM_THREAD_BLOCKS (MSCCL_MAX_NUM_THREAD_BLOCKS_PER_CHANNEL * MAXCHANNELS)
 #define MSCCL_MAX_COUNT 72 // max concurrent number of msccl chunk transmission
@@ -35,19 +36,6 @@
 #define MSCCL_LOCAL_COPY 6
 #define MSCCL_REDUCE 7
 
-typedef enum { mscclFuncReduce             =  0,
-               mscclFuncBroadcast          =  1,
-               mscclFuncAllReduce          =  2,
-               mscclFuncReduceScatter      =  3,
-               mscclFuncAllGather          =  4,
-               mscclFuncSend               =  5,
-               mscclFuncRecv               =  6,
-               mscclFuncGather             =  7,
-               mscclFuncScatter            =  8,
-               mscclFuncAllToAll           =  9,
-               mscclFuncAllToAllv          =  10,
-               mscclNumFuncs               =  11 } mscclFunc_t;
-
 struct mscclTransmission {
   int16_t dependencePointer; // index to the first dependence
   int16_t numDependencies; // dependencePointer+numDependencies indicate the last dependence
@@ -98,6 +86,27 @@ struct mscclChannelInfo {
   int nRecvPeers;
 };
 
+struct mscclAlgoMeta {
+  // Path to algorithm file
+  std::string filePath;
+  // number of chunks of input/output in each MSCCL algorithm loop
+  int nChunksPerLoop;
+  // number of ranks required by this algorithm
+  int nRanks;
+  // need to times nRanks for all-gather, reduce-scatter and all-to-all
+  int sizeMultiplier;
+  // MSCCL function type
+  mscclFunc_t func;
+  // Min message size allowed for this algorithm.
+  int64_t minBytes;
+  // Max message size allowed for this algorithm, 0 for no limit.
+  int64_t maxBytes;
+  // Whether this algorithm is suitable for in-place.
+  bool inPlace;
+  // Whether this algorithm is suitable for out-of-place.
+  bool outOfPlace;
+};
+
 struct mscclAlgo {
   // number of chunks of input/output in each MSCCL algorithm loop
   int nChunksPerLoop;
@@ -141,29 +150,23 @@ enum mscclGroupStatus {
   mscclGroupUnsupportedOp
 };
 
-struct mscclSchedulerParam {
-  const void* sendBuff;
-  const size_t* sendCounts;
+struct mscclSavedSchedulerParam {
+  struct mscclSchedulerParam p;
   std::vector<size_t> savedSendCounts;
-  const size_t* sDisPls;
   std::vector<size_t> savedSDisPls;
-  void* recvBuff;
-  const size_t* recvCounts;
   std::vector<size_t> savedRecvCounts;
-  const size_t* rDisPls;
   std::vector<size_t> savedRDisPls;
-  size_t count;
-  ncclDataType_t dataType;
-  int root;
-  int peer;
-  ncclRedOp_t op;
-  mscclFunc_t func;
-  bool scheduled;
-  mscclAlgoHandle_t handle;
   ncclComm_t comm;
   hipStream_t stream;
 };
 
+struct mscclThreadLocalStatus {
+  bool mscclIsCallerFlag;
+  mscclGroupStatus groupStatus;
+  int groupDepth;
+  std::vector<struct mscclSavedSchedulerParam> savedSchedulerParams;
+};
+
 struct mscclStatus {
   std::vector<mscclAlgoHandle_t> freeAlgoHandles;
   std::map<mscclAlgoHandle_t, mscclAlgo *> hostAlgos;
@@ -177,13 +180,15 @@ struct mscclStatus {
   int sliceSteps;
   int chunkSize;
   int chunkEffectiveSize;
-  int rank;
   uint32_t workIndex;
   uint32_t maxAllowedCount;
   ncclDataType_t dataType;
-  mscclGroupStatus groupStatus;
-  int groupDepth;
-  std::vector<struct mscclSchedulerParam> savedSchedulerParams;
+  std::map<ncclComm_t, std::set<mscclAlgoHandle_t>> connectedAlgos;
+  hipStream_t lastStream;
+  void* mscclSchedulerLib;
+  mscclSchedulerInterface* mscclSchedulerPtr;
+  std::vector<mscclAlgoMeta> algoMetas;
+  std::vector<std::map<int, mscclAlgoHandle_t>> rankToAlgoHandles;
 };
 
 struct alignas(16) mscclWork {

src/init.cc

@@ -859,6 +859,7 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
     int nc;
     bool pivotA2AEnabled;
     bool ll128Enabled;
+    bool mscclEnabled;
   };
 
   int nChannelsOrig;
@@ -951,6 +952,8 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
   comm->topo->pivotA2ANumBiRings = 0;
   // LL128
   comm->topo->ll128Enabled = false;
+  // Topology hint for MSCCL internal scheduler about whether to enable MSCCL
+  comm->topo->mscclEnabled = false;
   // Compute paths between GPUs and NICs
   NCCLCHECKGOTO(ncclTopoComputePaths(comm->topo, comm), ret, fail);
   // Remove inaccessible GPUs and unused NICs
@@ -1096,6 +1099,7 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
   allGather3Data[rank].pivotA2AEnabled = comm->topo->pivotA2AEnabled && rcclParamPivotAlltoallEnable();
   comm->topo->ll128Enabled =  comm->topo->ll128Enabled || rcclParamLL128ForceEnable();
   allGather3Data[rank].ll128Enabled = comm->topo->ll128Enabled;
+  allGather3Data[rank].mscclEnabled = comm->topo->mscclEnabled;
 
   comm->nChannels = (comm->topo->nodes[GPU].count != comm->topo->nRanks && comm->topo->nodes[NET].count)
     ? std::min(treeGraph.nChannels, ringGraph.nChannels) : ringGraph.nChannels;
@@ -1183,6 +1187,7 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
     comm->collNetSupport = std::min(allGather3Data[i].collNetSupport, comm->collNetSupport);
     comm->topo->pivotA2AEnabled = comm->topo->pivotA2AEnabled && allGather3Data[i].pivotA2AEnabled;
     comm->topo->ll128Enabled = comm->topo->ll128Enabled && allGather3Data[i].ll128Enabled;
+    comm->topo->mscclEnabled = comm->topo->mscclEnabled && allGather3Data[i].mscclEnabled;
   }
 
   comm->nChannels = treeGraph.nChannels = ringGraph.nChannels =

src/misc/msccl/msccl_lifecycle.cc

@@ -4,116 +4,192 @@
  ************************************************************************/
 
 #include <atomic>
+#include <map>
+#include <mutex>
+#include <set>
+
+#include <dirent.h>
+#include <dlfcn.h>
+#include <error.h>
+#include <link.h>
 
 #include "alloc.h"
 #include "checks.h"
+#include "graph/topo.h"
 
 #include "msccl/msccl_lifecycle.h"
 #include "msccl/msccl_parser.h"
 #include "msccl/msccl_setup.h"
 #include "msccl/msccl_status.h"
 
-RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 0);
+RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 1);
 static const char* mscclAlgoFilePathEnv = "MSCCL_ALGO_FILE_PATH";
 static std::atomic<bool> mscclInitialized;
 static bool mscclSchedulerTriedLoadAlgo = false;
+static std::mutex mscclLifecycleMutex;
 
 bool mscclEnabled() {
   return rcclParamMscclEnabled();
 }
 
-static bool mscclIsCallerFlag = false;
-
 void mscclSetIsCallerFlag() {
-  mscclIsCallerFlag = true;
+  mscclGetThreadLocalStatus().mscclIsCallerFlag = true;
 }
 
 void mscclClearIsCallerFlag() {
-  mscclIsCallerFlag = false;
+  mscclGetThreadLocalStatus().mscclIsCallerFlag = false;
 }
 
 bool mscclIsCaller() {
-  return mscclIsCallerFlag;
+  return mscclGetThreadLocalStatus().mscclIsCallerFlag;
 }
 
 bool mscclAvailable() {
   return mscclEnabled() && mscclInitialized.load(std::memory_order_acquire);
 }
 
-ncclResult_t mscclInit(ncclComm_t comm) {
-  if (comm->intraRanks > 1) {
-    mscclInitialized.store(false, std::memory_order_release);
-    INFO(NCCL_INIT, "MSCCL doesn't support multiple GPUs in one process and is not available");
-    return ncclSuccess;
+static bool mscclCommCompatible(ncclComm_t comm) {
+  std::map<uint64_t, std::set<uint64_t>> hostHashToPidHashes;
+  for (int i = 0; i < comm->nRanks; i++) {
+    uint64_t hostHash = comm->peerInfo[i].hostHash;
+    uint64_t pidHash = comm->peerInfo[i].pidHash;
+    if (hostHashToPidHashes.find(hostHash) != hostHashToPidHashes.end()) {
+      auto& pidHashSet = hostHashToPidHashes[hostHash];
+      if (pidHashSet.find(pidHash) != pidHashSet.end()) {
+        return false;
+      }
+    }
+    hostHashToPidHashes[hostHash].insert(pidHash);
+  }
+  return true;
+}
+
+static const char* mscclSchedulerPathEnv = "MSCCL_SCHEDULER";
+static const char* mscclSchedulerDefaultPath = "libmsccl-scheduler.so";
+static const char* mscclAlgoDirEnv = "MSCCL_ALGO_DIR";
+static const char* mscclAlgoDefaultDir = "msccl-algorithms";
+extern "C" bool mscclUnitTestMode() __attribute__((__weak__));
+static const char* mscclUnitTestAlgoDefaultDir = "msccl-unit-test-algorithms";
+
+static ncclResult_t mscclInternalSchedulerInit() {
+  mscclStatus& status = mscclGetStatus();
+  const char* mscclAlgoDir = getenv(mscclAlgoDirEnv);
+  std::string mscclAlgoDirStr;
+  if (mscclAlgoDir == nullptr) {
+    // Try to find default algorithm directory based on librccl.so path
+    Dl_info dl_info;
+    struct link_map *link_map_ptr = nullptr;
+    if (!dladdr1((void *)mscclInternalSchedulerInit, &dl_info, (void **)&link_map_ptr, RTLD_DL_LINKMAP)) {
+      WARN("MSCCL Internal Scheduler: dladdr1 failed");
+      return ncclInvalidUsage;
+    }
+    std::string selfLibPath = link_map_ptr->l_name;
+    mscclAlgoDirStr = selfLibPath.substr(0, selfLibPath.find_last_of("/\\") + 1);
+    mscclAlgoDirStr += (mscclUnitTestMode && mscclUnitTestMode()) ? mscclUnitTestAlgoDefaultDir : mscclAlgoDefaultDir;
+    mscclAlgoDir = mscclAlgoDirStr.c_str();
+  }
+  struct dirent *entry = nullptr;
+  DIR *dp = nullptr;
+  dp = opendir(mscclAlgoDir);
+  if (dp == nullptr) {
+    WARN("MSCCL Internal Scheduler: open algorithm directory %s failed", mscclAlgoDir);
+    return ncclInvalidUsage;
+  }
+  while ((entry = readdir(dp))) {
+    if (entry->d_type != DT_LNK && entry->d_type != DT_REG) {
+      continue;
+    }
+    status.algoMetas.emplace_back();
+    std::string fullPath = mscclAlgoDir;
+    fullPath += "/";
+    fullPath += entry->d_name;
+    NCCLCHECK(mscclGetAlgoMetaFromXmlFile(fullPath.c_str(), &(status.algoMetas.back())));
+  }
+  if (closedir(dp)) {
+    WARN("MSCCL Internal Scheduler: closedir failed, error %d", errno);
+    return ncclInvalidUsage;
+  }
+  status.rankToAlgoHandles.resize(status.algoMetas.size());
+  return ncclSuccess;
+}
+
+static ncclResult_t mscclSchedulerInit() {
+  mscclStatus& status = mscclGetStatus();
+  bool useInternalScheduler = false;
+
+  const char* mscclSchedulerPath = getenv(mscclSchedulerPathEnv);
+  if (mscclSchedulerPath) {
+    status.mscclSchedulerLib = dlopen(mscclSchedulerPath, RTLD_NOW | RTLD_LOCAL);
   } else {
+    status.mscclSchedulerLib = dlopen(mscclSchedulerDefaultPath, RTLD_NOW | RTLD_LOCAL);
+  }
+  if (status.mscclSchedulerLib == nullptr) {
+    INFO(NCCL_INIT, "MSCCL: No external scheduler found, using internal implementation");
+    useInternalScheduler = true;
+  } else {
+    status.mscclSchedulerPtr = (mscclSchedulerInterface *)dlsym(status.mscclSchedulerLib, "mscclScheduler");
+    if (status.mscclSchedulerPtr == nullptr) {
+      INFO(NCCL_INIT, "MSCCL: Failed to find mscclScheduler symbol, using internal implementation");
+      useInternalScheduler = true;
+    }
+  }
+  if (useInternalScheduler) {
+    NCCLCHECK(mscclInternalSchedulerInit());
+  } else {
+    NCCLCHECK(status.mscclSchedulerPtr->init());
+  }
+  return ncclSuccess;
+}
+
+ncclResult_t mscclInit(ncclComm_t comm) {
+  // Always initialize thread local status
+  mscclThreadLocalStatus threadLocalStatus = mscclGetThreadLocalStatus();
+  threadLocalStatus.groupStatus = mscclNoGroup;
+  threadLocalStatus.groupDepth = 0;
+  comm->mscclCompatible = mscclCommCompatible(comm);
+
+  {
+    std::lock_guard<std::mutex> lock(mscclLifecycleMutex);
+
+    if (mscclInitialized.load(std::memory_order_acquire)) {
+      return ncclSuccess;
+    }
+
+    mscclStatus& status = mscclGetStatus();
+    status.scratchBuffer = nullptr;
+    status.scratchBufferSize = 0;
+    status.workIndex = 1;
+    status.freeAlgoHandles.resize(MSCCL_MAX_NUM_ALGOS);
+    for (int i = 0; i < MSCCL_MAX_NUM_ALGOS; i++) {
+      status.freeAlgoHandles[i] = MSCCL_MAX_NUM_ALGOS - i - 1;
+    }
+    NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
+    status.lastStream = nullptr;
+    mscclSchedulerTriedLoadAlgo = false;
+
+    NCCLCHECK(mscclSchedulerInit());
+
     mscclInitialized.store(true, std::memory_order_release);
   }
 
-  mscclStatus& status = mscclGetStatus();
-  status.scratchBuffer = nullptr;
-  status.scratchBufferSize = 0;
-  status.rank = comm->rank;
-  status.workIndex = 1;
-  status.freeAlgoHandles.resize(MSCCL_MAX_NUM_ALGOS);
-  for (int i = 0; i < MSCCL_MAX_NUM_ALGOS; i++) {
-    status.freeAlgoHandles[i] = MSCCL_MAX_NUM_ALGOS - i - 1;
-  }
-  NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
-  status.groupStatus = mscclNoGroup;
-  status.groupDepth = 0;
-  mscclSchedulerTriedLoadAlgo = false;
+  INFO(NCCL_INIT, "MSCCL: Initialization finished");
   return ncclSuccess;
 }
 
 ncclResult_t mscclGroupStart() {
-  mscclStatus& status = mscclGetStatus();
-  status.groupDepth++;
-  if (status.groupStatus == mscclNoGroup) {
-    status.groupStatus = mscclGroupSupportedOp;
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  threadLocalStatus.groupDepth++;
+  if (threadLocalStatus.groupStatus == mscclNoGroup) {
+    threadLocalStatus.groupStatus = mscclGroupSupportedOp;
   }
   return ncclSuccess;
 }
 
-static ncclResult_t mscclScheduler(struct mscclSchedulerParam* param) {
-  static bool algoAvailable = false;
-  static mscclAlgoHandle_t loadedAlgoHandle;
-  static mscclAlgo* loadedHostAlgo = nullptr;
-
+static ncclResult_t mscclInternalSchedulerSelectAlgo(struct mscclSchedulerParam* param) {
+  mscclStatus& status = mscclGetStatus();
   param->scheduled = false;
 
-  if (!mscclSchedulerTriedLoadAlgo) {
-    mscclSchedulerTriedLoadAlgo = true;
-    const char* mscclAlgoFilePath = getenv(mscclAlgoFilePathEnv);
-    if (mscclAlgoFilePath != nullptr) {
-      NCCLCHECK(mscclLoadAlgo(mscclAlgoFilePath, &loadedAlgoHandle));
-      mscclStatus& status = mscclGetStatus();
-      loadedHostAlgo = status.hostAlgos[loadedAlgoHandle];
-      algoAvailable = true;
-    }
-  }
-  if (!algoAvailable) {
-    return ncclSuccess;
-  }
-
-  bool mscclAlgoFuncIsValid = loadedHostAlgo->func == param->func;
-  if (!mscclAlgoFuncIsValid) {
-    return ncclSuccess;
-  }
-
-  bool numGpusIsValid = loadedHostAlgo->nRanks == param->comm->nRanks;
-  if (!numGpusIsValid) {
-    return ncclSuccess;
-  }
-
-  size_t nBytes = param->count * ncclTypeSize(param->dataType) * loadedHostAlgo->sizeMultiplier;
-  bool msgSizeIsValid =
-    param->count > 0 && (param->count % loadedHostAlgo->nChunksPerLoop) == 0 &&
-    nBytes >= loadedHostAlgo->minBytes &&
-    (loadedHostAlgo->maxBytes == 0 || nBytes <= loadedHostAlgo->maxBytes);
-  if (!msgSizeIsValid) {
-    return ncclSuccess;
-  }
-
+  // Whether the algorithm is in-place
   bool isInPlace = false;
   if (param->func == mscclFuncReduce ||
       param->func == mscclFuncBroadcast ||
@@ -123,120 +199,153 @@ static ncclResult_t mscclScheduler(struct mscclSchedulerParam* param) {
     isInPlace = param->sendBuff == param->recvBuff;
   } else if (param->func == mscclFuncAllGather ||
              param->func == mscclFuncGather) {
-    isInPlace = (char*)param->sendBuff == (char*)param->recvBuff + param->comm->rank * param->count * ncclTypeSize(param->dataType);
+    isInPlace = (char*)param->sendBuff == (char*)param->recvBuff + param->rank * param->count * ncclTypeSize(param->dataType);
   } else if (param->func == mscclFuncReduceScatter ||
              param->func == mscclFuncScatter) {
-    isInPlace = (char*)param->recvBuff == (char*)param->sendBuff + param->comm->rank * param->count * ncclTypeSize(param->dataType);
-  }
-  bool inPlaceOutOfPlaceIsValid = isInPlace ? loadedHostAlgo->inPlace : loadedHostAlgo->outOfPlace;
-  if (!inPlaceOutOfPlaceIsValid) {
-    return ncclSuccess;
+    isInPlace = (char*)param->recvBuff == (char*)param->sendBuff + param->rank * param->count * ncclTypeSize(param->dataType);
+  }
+
+  // Search suitable algorithms
+  for (size_t i = 0; i < status.algoMetas.size(); i++) {
+    auto &m = status.algoMetas[i];
+    size_t nBytes = param->count * ncclTypeSize(param->dataType) * m.sizeMultiplier;
+    bool msgSizeIsValid =
+      param->count > 0 && (param->count % m.nChunksPerLoop) == 0 &&
+      nBytes >= m.minBytes && (m.maxBytes == 0 || nBytes <= m.maxBytes);
+    if (msgSizeIsValid &&
+        m.nRanks == param->nRanks &&
+        m.func == param->func &&
+        (isInPlace ? m.inPlace : m.outOfPlace)) {
+      // If not loaded for current rank, load it
+      if (status.rankToAlgoHandles[i].find(param->rank) == status.rankToAlgoHandles[i].end()) {
+        mscclAlgoHandle_t algoHandle;
+        NCCLCHECK(mscclLoadAlgo(m.filePath.c_str(), &algoHandle, param->rank));
+        status.rankToAlgoHandles[i][param->rank] = algoHandle;
+      }
+      param->handle = status.rankToAlgoHandles[i][param->rank];
+      param->scheduled = true;
+      return ncclSuccess;
+    }
   }
 
-  param->handle = loadedAlgoHandle;
-  param->scheduled = true;
   return ncclSuccess;
 }
 
-static ncclResult_t mscclSetSchedulerParam(
+static ncclResult_t mscclSchedulerSelectAlgo(struct mscclSavedSchedulerParam* param) {
+  mscclStatus& status = mscclGetStatus();
+  if (status.mscclSchedulerPtr) {
+    NCCLCHECK(status.mscclSchedulerPtr->selectAlgo(&(param->p)));
+  } else {
+    if (param->comm->topo->mscclEnabled) {
+      NCCLCHECK(mscclInternalSchedulerSelectAlgo(&(param->p)));
+    } else {
+      param->p.scheduled = false;
+    }
+  }
+  return ncclSuccess;
+}
+
+static ncclResult_t mscclSetSavedSchedulerParam(
   const void* sendBuff, const size_t sendCounts[], const size_t sDisPls[],
   void* recvBuff, const size_t recvCounts[], const size_t rDisPls[],
   size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op,
   mscclFunc_t func, ncclComm_t comm, hipStream_t stream,
-  struct mscclSchedulerParam* param) {
-  param->sendBuff = sendBuff;
-  param->sendCounts = sendCounts;
-  param->sDisPls = sDisPls;
-  param->recvBuff = recvBuff;
-  param->recvCounts = recvCounts;
-  param->rDisPls = rDisPls;
-  param->count = count;
-  param->dataType = dataType;
-  param->root = root;
-  param->peer = peer;
-  param->op = op;
-  param->func = func;
+  struct mscclSavedSchedulerParam* param) {
+  param->p.sendBuff = sendBuff;
+  param->p.sendCounts = sendCounts;
+  param->p.sDisPls = sDisPls;
+  param->p.recvBuff = recvBuff;
+  param->p.recvCounts = recvCounts;
+  param->p.rDisPls = rDisPls;
+  param->p.count = count;
+  param->p.dataType = dataType;
+  param->p.root = root;
+  param->p.peer = peer;
+  param->p.op = op;
+  param->p.func = func;
+  param->p.rank = comm->rank;
+  param->p.nRanks = comm->nRanks;
   param->comm = comm;
   param->stream = stream;
   return ncclSuccess;
 }
 
-static ncclResult_t mscclSaveCountsAndDispls(struct mscclSchedulerParam* param) {
-  if (param->sendCounts) {
-    param->savedSendCounts.assign(param->sendCounts, param->sendCounts + param->comm->nRanks);
-    param->sendCounts = param->savedSendCounts.data();
-    param->savedSDisPls.assign(param->sDisPls, param->sDisPls + param->comm->nRanks);
-    param->sDisPls = param->savedSDisPls.data();
-    param->savedRecvCounts.assign(param->recvCounts, param->recvCounts + param->comm->nRanks);
-    param->recvCounts = param->savedRecvCounts.data();
-    param->savedRDisPls.assign(param->rDisPls, param->rDisPls + param->comm->nRanks);
-    param->rDisPls = param->savedRDisPls.data();
+static ncclResult_t mscclSaveCountsAndDispls(struct mscclSavedSchedulerParam* param) {
+  if (param->p.sendCounts) {
+    param->savedSendCounts.assign(param->p.sendCounts, param->p.sendCounts + param->p.nRanks);
+    param->p.sendCounts = param->savedSendCounts.data();
+    param->savedSDisPls.assign(param->p.sDisPls, param->p.sDisPls + param->p.nRanks);
+    param->p.sDisPls = param->savedSDisPls.data();
+    param->savedRecvCounts.assign(param->p.recvCounts, param->p.recvCounts + param->p.nRanks);
+    param->p.recvCounts = param->savedRecvCounts.data();
+    param->savedRDisPls.assign(param->p.rDisPls, param->p.rDisPls + param->p.nRanks);
+    param->p.rDisPls = param->savedRDisPls.data();
   }
   return ncclSuccess;
 }
 
 static ncclResult_t mscclRunSavedParams() {
-  mscclStatus& status = mscclGetStatus();
-  for (auto& param : status.savedSchedulerParams) {
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  for (auto& param : threadLocalStatus.savedSchedulerParams) {
     NCCLCHECK(mscclRunAlgo(
-      param.sendBuff, param.sendCounts, param.sDisPls,
-      param.recvBuff, param.recvCounts, param.rDisPls,
-      param.count, param.dataType, param.root, param.peer, param.op, param.handle, param.comm, param.stream));
+      param.p.sendBuff, param.p.sendCounts, param.p.sDisPls,
+      param.p.recvBuff, param.p.recvCounts, param.p.rDisPls,
+      param.p.count, param.p.dataType, param.p.root, param.p.peer, param.p.op, param.p.handle, param.comm, param.stream));
   }
-  status.savedSchedulerParams.clear();
+  threadLocalStatus.savedSchedulerParams.clear();
   return ncclSuccess;
 }
 
 static ncclResult_t mscclFallBackSavedParams() {
-  mscclStatus& status = mscclGetStatus();
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
   mscclSetIsCallerFlag();
-  for (auto& param : status.savedSchedulerParams) {
-    switch (param.func) {
+  for (auto& param : threadLocalStatus.savedSchedulerParams) {
+    switch (param.p.func) {
       case mscclFuncReduce:
-        NCCLCHECK(ncclReduce(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.op, param.root, param.comm, param.stream));
+        NCCLCHECK(ncclReduce(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.op, param.p.root, param.comm, param.stream));
         break;
       case mscclFuncBroadcast:
-        NCCLCHECK(ncclBroadcast(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.root, param.comm, param.stream));
+        NCCLCHECK(ncclBroadcast(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.root, param.comm, param.stream));
         break;
       case mscclFuncAllReduce:
-        NCCLCHECK(ncclAllReduce(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.op, param.comm, param.stream));
+        NCCLCHECK(ncclAllReduce(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.op, param.comm, param.stream));
         break;
       case mscclFuncReduceScatter:
-        NCCLCHECK(ncclReduceScatter(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.op, param.comm, param.stream));
+        NCCLCHECK(ncclReduceScatter(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.op, param.comm, param.stream));
         break;
       case mscclFuncAllGather:
-        NCCLCHECK(ncclAllGather(param.sendBuff, param.recvBuff, param.count, param.dataType,
+        NCCLCHECK(ncclAllGather(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
           param.comm, param.stream));
         break;
       case mscclFuncSend:
-        NCCLCHECK(ncclSend(param.sendBuff, param.count, param.dataType,
-          param.peer, param.comm, param.stream));
+        NCCLCHECK(ncclSend(param.p.sendBuff, param.p.count, param.p.dataType,
+          param.p.peer, param.comm, param.stream));
         break;
       case mscclFuncRecv:
-        NCCLCHECK(ncclRecv(param.recvBuff, param.count, param.dataType,
-          param.peer, param.comm, param.stream));
+        NCCLCHECK(ncclRecv(param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.peer, param.comm, param.stream));
         break;
       case mscclFuncGather:
-        NCCLCHECK(ncclGather(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.root, param.comm, param.stream));
+        NCCLCHECK(ncclGather(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.root, param.comm, param.stream));
         break;
       case mscclFuncScatter:
-        NCCLCHECK(ncclScatter(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.root, param.comm, param.stream));
+        NCCLCHECK(ncclScatter(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.root, param.comm, param.stream));
         break;
       case mscclFuncAllToAll:
-        NCCLCHECK(ncclAllToAll(param.sendBuff, param.recvBuff, param.count, param.dataType,
+        NCCLCHECK(ncclAllToAll(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
           param.comm, param.stream));
         break;
       case mscclFuncAllToAllv:
         NCCLCHECK(ncclAllToAllv(
-          param.sendBuff, param.sendCounts, param.sDisPls,
-          param.recvBuff, param.recvCounts, param.rDisPls,
-          param.dataType, param.comm, param.stream));
+          param.p.sendBuff, param.p.sendCounts, param.p.sDisPls,
+          param.p.recvBuff, param.p.recvCounts, param.p.rDisPls,
+          param.p.dataType, param.comm, param.stream));
         break;
       default:
         WARN("Invalid MSCCL function type in saved parameter");
@@ -244,7 +353,7 @@ static ncclResult_t mscclFallBackSavedParams() {
     }
   }
   mscclClearIsCallerFlag();
-  status.savedSchedulerParams.clear();
+  threadLocalStatus.savedSchedulerParams.clear();
   return ncclSuccess;
 }
 
@@ -253,40 +362,43 @@ ncclResult_t mscclEnqueueCheck(
     void* recvBuff, const size_t recvCounts[], const size_t rDisPls[],
     size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op,
     mscclFunc_t func, ncclComm_t comm, hipStream_t stream) {
-  mscclStatus& status = mscclGetStatus();
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
   hipStreamCaptureStatus captureStatus;
   unsigned long long pid;
 
-  status.savedSchedulerParams.push_back({});
-  NCCLCHECK(mscclSetSchedulerParam(
+  threadLocalStatus.savedSchedulerParams.push_back({});
+  NCCLCHECK(mscclSetSavedSchedulerParam(
     sendBuff, sendCounts, sDisPls, recvBuff, recvCounts, rDisPls,
     count, dataType, root, peer, op, func, comm, stream,
-    &status.savedSchedulerParams.back()));
+    &threadLocalStatus.savedSchedulerParams.back()));
 
-  switch (status.groupStatus) {
+  switch (threadLocalStatus.groupStatus) {
     case mscclNoGroup:
-      CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
-      if (captureStatus == hipStreamCaptureStatusNone) {
-        NCCLCHECK(mscclScheduler(&status.savedSchedulerParams.back()));
-        if (status.savedSchedulerParams.back().scheduled) {
-          NCCLCHECK(mscclRunSavedParams());
-          break;
+      if (comm->mscclCompatible) {
+        CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
+        if (captureStatus == hipStreamCaptureStatusNone) {
+          NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
+          if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
+            NCCLCHECK(mscclRunSavedParams());
+            break;
+          }
         }
       }
       NCCLCHECK(mscclFallBackSavedParams());
       break;
     case mscclGroupSupportedOp:
-      CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
-      if (captureStatus == hipStreamCaptureStatusNone) {
-        NCCLCHECK(mscclScheduler(&status.savedSchedulerParams.back()));
-        if (status.savedSchedulerParams.back().scheduled) {
-          // Only save counts and displs when there is suitable MSCCL algorithm for this
-          NCCLCHECK(mscclSaveCountsAndDispls(&status.savedSchedulerParams.back()));
-          break;
+      if (comm->mscclCompatible) {
+        CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
+        if (captureStatus == hipStreamCaptureStatusNone) {
+          NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
+          if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
+            // Only save counts and displs when there is suitable MSCCL algorithm for this
+            NCCLCHECK(mscclSaveCountsAndDispls(&threadLocalStatus.savedSchedulerParams.back()));
+            break;
+          }
         }
       }
-      NCCLCHECK(mscclFallBackSavedParams());
-      break;
+      threadLocalStatus.groupStatus = mscclGroupUnsupportedOp;
     case mscclGroupUnsupportedOp:
       NCCLCHECK(mscclFallBackSavedParams());
       break;
@@ -297,37 +409,71 @@ ncclResult_t mscclEnqueueCheck(
 }
 
 ncclResult_t mscclGroupEnd() {
-  mscclStatus& status = mscclGetStatus();
-  status.groupDepth--;
-  if (status.groupDepth == 0) {
-    if (status.groupStatus == mscclGroupSupportedOp) {
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  threadLocalStatus.groupDepth--;
+  if (threadLocalStatus.groupDepth == 0) {
+    if (threadLocalStatus.groupStatus == mscclGroupSupportedOp) {
       NCCLCHECK(mscclRunSavedParams());
     }
-    status.groupStatus = mscclNoGroup;
+    threadLocalStatus.groupStatus = mscclNoGroup;
   }
   return ncclSuccess;
 }
 
-ncclResult_t mscclTeardown() {
-  if (!mscclInitialized.load(std::memory_order_acquire)) {
-    return ncclSuccess;
-  }
+static ncclResult_t mscclInternalSchedulerTeardown() {
+  ncclResult_t ret = ncclSuccess, tmpRet = ncclSuccess;
   mscclStatus& status = mscclGetStatus();
-  for (auto &p : status.hostAlgos) {
-    free(p.second);
-    status.freeAlgoHandles.push_back(p.first);
+  for (auto &m : status.rankToAlgoHandles) {
+    for (auto &p : m) {
+      tmpRet = mscclUnloadAlgo(p.second);
+      if (ret == ncclSuccess) {
+        ret = tmpRet;
+      }
+    }
   }
-  for (auto &p : status.devAlgos) {
-    CUDACHECK(hipFree(p.second));
+  status.algoMetas.clear();
+  status.rankToAlgoHandles.clear();
+  return ret;
+}
+
+ncclResult_t mscclTeardown() {
+  // Always teardown thread local status
+  mscclThreadLocalStatus threadLocalStatus = mscclGetThreadLocalStatus();
+  threadLocalStatus.savedSchedulerParams.clear();
+
+  {
+    std::lock_guard<std::mutex> lock(mscclLifecycleMutex);
+
+    if (!mscclInitialized.load(std::memory_order_acquire)) {
+      return ncclSuccess;
+    }
+    mscclStatus& status = mscclGetStatus();
+    for (auto &p : status.hostAlgos) {
+      free(p.second);
+      status.freeAlgoHandles.push_back(p.first);
+    }
+    for (auto &p : status.devAlgos) {
+      CUDACHECK(hipFree(p.second));
+    }
+    CUDACHECK(hipFree(status.scratchBuffer));
+    CUDACHECK(hipFree(status.syncFlags));
+    status.hostAlgos.clear();
+    status.devAlgos.clear();
+    status.freeAlgoHandles.clear();
+    status.scratchBuffer = nullptr;
+    status.scratchBufferSize = 0;
+    status.connectedAlgos.clear();
+    if (status.mscclSchedulerPtr) {
+      NCCLCHECK(status.mscclSchedulerPtr->teardown());
+      status.mscclSchedulerPtr = nullptr;
+      dlclose(status.mscclSchedulerLib);
+      status.mscclSchedulerLib = nullptr;
+    } else {
+      NCCLCHECK(mscclInternalSchedulerTeardown());
+    }
+    mscclInitialized.store(false, std::memory_order_release);
   }
-  CUDACHECK(hipFree(status.scratchBuffer));
-  CUDACHECK(hipFree(status.syncFlags));
-  status.hostAlgos.clear();
-  status.devAlgos.clear();
-  status.freeAlgoHandles.clear();
-  status.scratchBuffer = nullptr;
-  status.scratchBufferSize = 0;
-  status.workIndex = 1;
-  mscclInitialized.store(false, std::memory_order_release);
+
+  INFO(NCCL_INIT, "MSCCL: Teardown finished");
   return ncclSuccess;
 }

src/misc/msccl/msccl_parser.cc

@@ -701,3 +701,88 @@ ncclResult_t mscclGetAlgoFromXmlFile(const char* str, struct mscclAlgo* algo, in
   free(xml);
   return ncclSuccess;
 }
+
+ncclResult_t mscclXmlLoadSingleNode(FILE* file, struct mscclXmlNode* node) {
+  memset(node, 0, sizeof(struct mscclXmlNode));
+  return mscclXmlGetNode(file, node);
+}
+
+ncclResult_t mscclAlgoMetaXmlLoad(const char* xmlFilePath, struct mscclXmlNode* node) {
+  ncclResult_t ret = ncclSuccess;
+  FILE* file = fopen(xmlFilePath, "r");
+  if (file == NULL) {
+    fprintf(stderr, "Could not open MSCCL XML algorithm file %s : %s", xmlFilePath, strerror(errno));
+    return ncclSystemError;
+  }
+  NCCLCHECK(mscclXmlLoadSingleNode(file, node));
+  fclose(file);
+  return ncclSuccess;
+}
+
+ncclResult_t mscclGetAlgoMetaFromXmlFile(const char* str, struct mscclAlgoMeta* algoMeta) {
+  ncclResult_t ret = ncclSuccess;
+  struct mscclXmlNode* node;
+  node = (struct mscclXmlNode *)malloc(sizeof(struct mscclXmlNode));
+  NCCLCHECK(mscclAlgoMetaXmlLoad(str, node));
+
+  algoMeta->filePath = str;
+
+  int nChunksPerLoop;
+  NCCLCHECK(mscclXmlGetAttrInt(node, "nchunksperloop", &nChunksPerLoop));
+  algoMeta->nChunksPerLoop  = nChunksPerLoop;
+
+  int nGpus;
+  NCCLCHECK(mscclXmlGetAttrInt(node, "ngpus", &nGpus));
+  algoMeta->nRanks = nGpus;
+
+  const char* coll;
+  NCCLCHECK(mscclXmlGetAttrStr(node, "coll", &coll));
+  algoMeta->sizeMultiplier = 1;
+  if (strcmp(coll, "reduce") == 0) {
+    algoMeta->func = mscclFuncReduce;
+  } else if (strcmp(coll, "broadcast") == 0) {
+    algoMeta->func = mscclFuncBroadcast;
+  } else if (strcmp(coll, "allreduce") == 0) {
+    algoMeta->func = mscclFuncAllReduce;
+  } else if (strcmp(coll, "reducescatter") == 0) {
+    algoMeta->sizeMultiplier = nGpus;
+    algoMeta->func = mscclFuncReduceScatter;
+  } else if (strcmp(coll, "allgather") == 0) {
+    algoMeta->sizeMultiplier = nGpus;
+    algoMeta->func = mscclFuncAllGather;
+  } else if (strcmp(coll, "send") == 0) {
+    algoMeta->func = mscclFuncSend;
+  } else if (strcmp(coll, "recv") == 0) {
+    algoMeta->func = mscclFuncRecv;
+  } else if (strcmp(coll, "gather") == 0) {
+    algoMeta->func = mscclFuncGather;
+  } else if (strcmp(coll, "scatter") == 0) {
+    algoMeta->func = mscclFuncScatter;
+  } else if (strcmp(coll, "alltoall") == 0) {
+    algoMeta->sizeMultiplier = nGpus;
+    algoMeta->func = mscclFuncAllToAll;
+  } else if (strcmp(coll, "alltoallv") == 0) {
+    algoMeta->func = mscclFuncAllToAllv;
+  } else {
+    return ncclInvalidUsage;
+  }
+
+  int64_t minBytes;
+  NCCLCHECK(mscclXmlGetAttrInt64(node, "minBytes", &minBytes));
+  algoMeta->minBytes = minBytes;
+
+  int64_t maxBytes;
+  NCCLCHECK(mscclXmlGetAttrInt64(node, "maxBytes", &maxBytes));
+  algoMeta->maxBytes = maxBytes;
+
+  int inplace;
+  NCCLCHECK(mscclXmlGetAttrInt(node, "inplace", &inplace));
+  algoMeta->inPlace = (bool)inplace;
+
+  int outofplace;
+  NCCLCHECK(mscclXmlGetAttrInt(node, "outofplace", &outofplace));
+  algoMeta->outOfPlace = (bool)outofplace;
+
+  free(node);
+  return ncclSuccess;
+}

src/misc/msccl/msccl_setup.cc

@@ -258,6 +258,11 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
     ncclDataType_t dataType, ncclRedOp_t op, struct mscclAlgo* hostAlgo, struct mscclAlgo* devAlgo,
     ncclComm_t comm, hipStream_t stream) {
   mscclStatus& status = mscclGetStatus();
+
+  if (status.lastStream != stream && status.lastStream != nullptr) {
+    CUDACHECK(hipStreamWaitEvent(stream, comm->doneEvent, 0));
+  }
+
   dim3 grid = {(uint32_t)hostAlgo->nBlocks, 1, 1};
   dim3 block = {NCCL_MAX_NTHREADS, 1, 1};
   ncclDevRedOpFull opFull;
@@ -278,7 +283,8 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
 
   void *args[3] = {&comm->devComm, &devAlgo, &work};
   void *func = mscclKernelEntries[(opFull.op * ncclNumTypes + dataType) * NCCL_NUM_PROTOCOLS + hostAlgo->protocol];
-  CUDACHECK(hipExtLaunchKernel(func, grid, block, args, 0, stream, NULL, NULL,0));
+  CUDACHECK(hipExtLaunchKernel(func, grid, block, args, 0, stream, NULL, comm->doneEvent, 0));
   status.workIndex++;
+  status.lastStream = stream;
   return ncclSuccess;
 }

src/misc/msccl/msccl_status.cc

@@ -9,3 +9,8 @@ mscclStatus& mscclGetStatus() {
   static mscclStatus status;
   return status;
 }
+
+mscclThreadLocalStatus& mscclGetThreadLocalStatus() {
+  static thread_local mscclThreadLocalStatus threadLocalStatus;
+  return threadLocalStatus;
+}

src/nccl.h.in

@@ -508,8 +508,8 @@ typedef int mscclAlgoHandle_t;
  * its handle via mscclAlgoHandle. This API is expected to be called by MSCCL
  * scheduler instead of end users.
  */
-ncclResult_t  mscclLoadAlgo(const char *mscclAlgoFilePath, mscclAlgoHandle_t *mscclAlgoHandle);
-ncclResult_t pmscclLoadAlgo(const char *mscclAlgoFilePath, mscclAlgoHandle_t *mscclAlgoHandle);
+ncclResult_t  mscclLoadAlgo(const char *mscclAlgoFilePath, mscclAlgoHandle_t *mscclAlgoHandle, int rank);
+ncclResult_t pmscclLoadAlgo(const char *mscclAlgoFilePath, mscclAlgoHandle_t *mscclAlgoHandle, int rank);
 
 /*! @brief MSCCL Run Algorithm
  *