enable hip graph on multi-node (#884)

* initial checkin

* enable msccl when hip graph is on

* remove the commented out code of msccl enable check

* clean up the code

* remove the msccl HighestTransportType check logic

src/collectives/device/msccl_kernel_impl.h

@@ -23,6 +23,9 @@ extern __shared__ struct mscclShmemData mscclShmem;
 #define COMPUTE_FLAG(__WORKINDEX__,__GRIDOFFSET_ITER__,__STEP__) \
   MSCCL_MAX_ITER*MSCCL_MAX_NUM_STEPS*(uint64_t)__WORKINDEX__ + ((uint64_t)__GRIDOFFSET_ITER__ * MSCCL_MAX_NUM_STEPS + (uint64_t)__STEP__)
 
+#define GET_WORKINDEX_FROM_FLAG(__FLAG__) \
+  (__FLAG__) / (MSCCL_MAX_ITER*MSCCL_MAX_NUM_STEPS)
+
 // a copy of the volatile load/store from prims_ll
 template<typename U>
 __device__ static U load(U *src) {
@@ -293,7 +296,10 @@ __device__ __forceinline__ void mscclRunInterpreter(
           int8_t dependentBid = mscclShmem.mscclTB.dependentBid[dependentPointer+tid];
           int16_t dependentStep = mscclShmem.mscclTB.dependentStep[dependentPointer+tid];
           uint64_t goalFlag = COMPUTE_FLAG(workIndex, iter, dependentStep);
-          while ((mscclFlags + dependentBid)->flag < goalFlag);
+          while (true){
+            uint64_t curFlag = (mscclFlags + dependentBid)->flag;
+            if (curFlag >= goalFlag && GET_WORKINDEX_FROM_FLAG(curFlag) == workIndex) break;
+          }
         }
         step += numDependencies-1;
         barrier(nthreads, mscclBarrierNext, mscclBarriers);

src/collectives/msccl.cc

@@ -48,6 +48,8 @@ ncclResult_t mscclRunAlgo(
   struct mscclAlgo* hostAlgo = status.hostAlgos[mscclAlgoHandle];
   struct mscclAlgo* devAlgo = status.devAlgos[mscclAlgoHandle];
 
+  NCCLCHECK(mscclGetCaptureStatus(stream));
+
   NCCLCHECK(mscclSetupCount(hostAlgo, comm, count, dataType));
 
   NCCLCHECK(mscclSetupScratch(hostAlgo, stream));
@@ -55,11 +57,14 @@ ncclResult_t mscclRunAlgo(
   NCCLCHECK(mscclSetupSyncFlags(stream));
 
   if (status.connectedAlgos[comm].find(mscclAlgoHandle) == status.connectedAlgos[comm].end()) {
+    hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+    CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
     NCCLCHECK(mscclSetupConnections(hostAlgo, comm));
+    CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
     status.connectedAlgos[comm].insert(mscclAlgoHandle);
   }
 
-  NCCLCHECK(mscclSetupProxy(hostAlgo, comm));
+  NCCLCHECK(mscclSetupProxy(hostAlgo, comm, stream));
 
   NCCLCHECK(mscclSetupKernel(sendBuff, recvBuff, count, dataType, op, hostAlgo, devAlgo, comm, stream));
 

src/include/msccl/msccl_setup.h

@@ -11,6 +11,8 @@
 #include "comm.h"
 #include "msccl/msccl_struct.h"
 
+ncclResult_t mscclGetCaptureStatus(hipStream_t stream);
+
 ncclResult_t mscclSetupScratch(struct mscclAlgo* hostAlgo, hipStream_t stream);
 
 ncclResult_t mscclSetupSyncFlags(hipStream_t stream);
@@ -19,7 +21,7 @@ ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo, ncclComm_t comm);
 
 ncclResult_t mscclSetupCount(struct mscclAlgo* hostAlgo, ncclComm_t comm, size_t count, ncclDataType_t dataType);
 
-ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm);
+ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm, hipStream_t stream);
 
 ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count,
     ncclDataType_t dataType, ncclRedOp_t op, struct mscclAlgo* hostAlgo, struct mscclAlgo* devAlgo,

src/include/msccl/msccl_status.h

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

src/include/msccl/msccl_struct.h

@@ -160,11 +160,29 @@ struct mscclSavedSchedulerParam {
   hipStream_t stream;
 };
 
+enum mscclCaptureStatus {
+  mscclNoCapture,
+  mscclNewCapture,
+  mscclExistingCapture
+};
+
+struct mscclProxyArg {
+  struct mscclAlgo* hostAlgo;
+  ncclComm_t comm;
+  mscclProxyArg(struct mscclAlgo* hostAlgo, ncclComm_t comm) 
+    : hostAlgo(hostAlgo), comm(comm) {}
+};
+
+typedef std::map<unsigned long long, std::vector<struct mscclProxyArg>> mscclSavedProxyArgs;
+
 struct mscclThreadLocalStatus {
   bool mscclIsCallerFlag;
   mscclGroupStatus groupStatus;
   int groupDepth;
   std::vector<struct mscclSavedSchedulerParam> savedSchedulerParams;
+  unsigned long long captureId;
+  mscclCaptureStatus captureStatus;
+  hipGraph_t graph;
 };
 
 struct mscclStatus {
@@ -189,6 +207,9 @@ struct mscclStatus {
   mscclSchedulerInterface* mscclSchedulerPtr;
   std::vector<mscclAlgoMeta> algoMetas;
   std::vector<std::map<int, mscclAlgoHandle_t>> rankToAlgoHandles;
+  bool graphEnabled;
+  bool graphFirstKernel;
+  bool needsProxy;
 };
 
 struct alignas(16) mscclWork {

src/include/transport.h

@@ -101,7 +101,7 @@ struct ncclTransport {
 };
 
 ncclResult_t ncclTransportP2pConnect(struct ncclComm* comm, int channelId, int nrecv, int* peerRecv, int nsend, int* peerSend, int connIndex);
-ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType=NULL);
+ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType=NULL, bool* needsProxy=NULL);
 
 ncclResult_t ncclNvlsInit(struct ncclComm* comm);
 ncclResult_t ncclNvlsSetup(struct ncclComm* comm, struct ncclComm* parent);

src/init.cc

@@ -42,6 +42,7 @@
 // [/RCCL]
 
 #include "msccl/msccl_lifecycle.h"
+#include "msccl/msccl_status.h"
 
 #define STR2(v) #v
 #define STR(v) STR2(v)
@@ -669,10 +670,6 @@ static ncclResult_t devCommSetup(ncclComm_t comm) {
   NCCLCHECK(ncclCudaCalloc(&tmpCommAndChans.comm.devProf, MAXCHANNELS*PROFILE_NUM_LAUNCHES), comm->sideStream);
 #endif
 
-  if (mscclEnabled()) {
-    NCCLCHECK(mscclInit(comm));
-  }
-
   NCCLCHECKGOTO(ncclCudaMemcpyAsync(devCommAndChans, &tmpCommAndChans, 1, comm->sharedRes->deviceStream.cudaStream), ret, fail);
 exit:
   CUDACHECK(cudaStreamSynchronize(comm->sharedRes->deviceStream.cudaStream));
@@ -1046,6 +1043,10 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p
   int *topParentLocalRanks = NULL;
   int tpProxyRank;
 
+  int highestTransportType = TRANSPORT_P2P;
+  bool needsProxy = false;
+  bool mscclNeedsProxy = needsProxy;
+
   // AllGather1 - begin
   NCCLCHECKGOTO(ncclCalloc(&comm->peerInfo, nranks+1), ret, fail); // Extra rank to represent CollNet root
   NCCLCHECKGOTO(fillInfo(comm, comm->peerInfo+rank, comm->commHash), ret, fail);
@@ -1414,7 +1415,8 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p
     if (comm->nRanks == 1) continue;
     NCCLCHECKGOTO(ncclTransportP2pConnect(comm, c, 1, &channel->ring.prev, 1, &channel->ring.next, 0), ret, fail);
   }
-  NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &ringGraph, 0), ret, fail);
+  NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &ringGraph, 0, &highestTransportType, &needsProxy), ret, fail);
+  mscclNeedsProxy |= needsProxy;
   if (ringGraph.nIntraChannels && rcclParamP2pNetDisable() == 0) {
     comm->useIntraNet = 1;
     // Connect NET for intranode use
@@ -1434,7 +1436,8 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p
     NCCLCHECKGOTO(ncclTransportP2pConnect(comm, c, NCCL_MAX_TREE_ARITY, channel->tree.down, 1, &channel->tree.up, 0), ret, fail);
     NCCLCHECKGOTO(ncclTransportP2pConnect(comm, c, 1, &channel->tree.up, NCCL_MAX_TREE_ARITY, channel->tree.down, 0), ret, fail);
   }
-  NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &treeGraph, 0), ret, fail);
+  NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &treeGraph, 0, &highestTransportType, &needsProxy), ret, fail);
+  mscclNeedsProxy |= needsProxy;
   INFO(NCCL_INIT, "Connected all trees");
 
   // Setup NVLS
@@ -1565,6 +1568,11 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p
   // Call devCommSetup before the last barrier, making sure we don't have a thread running in front and starting to
   // launch NCCL kernels before all cuda mem allocation is complete. That could cause a deadlock.
   NCCLCHECKGOTO(devCommSetup(comm), ret, fail);
+  if (mscclEnabled()) {
+    NCCLCHECK(mscclInit(comm));
+    mscclStatus& status = mscclGetStatus();
+    status.needsProxy |= mscclNeedsProxy;
+  }
 
   /* Local intra-node barrier */
   NCCLCHECKGOTO(bootstrapBarrier(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks, comm->localRankToRank[0]), ret, fail);

src/misc/msccl/msccl_lifecycle.cc

@@ -168,6 +168,8 @@ ncclResult_t mscclInit(ncclComm_t comm) {
   mscclThreadLocalStatus threadLocalStatus = mscclGetThreadLocalStatus();
   threadLocalStatus.groupStatus = mscclNoGroup;
   threadLocalStatus.groupDepth = 0;
+  threadLocalStatus.captureId = ULLONG_MAX;
+  threadLocalStatus.captureStatus = mscclNoCapture;
   comm->mscclCompatible = mscclCommCompatible(comm);
 
   {
@@ -187,6 +189,7 @@ ncclResult_t mscclInit(ncclComm_t comm) {
     }
     NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
     status.lastStream = nullptr;
+    status.needsProxy = false;
     mscclSchedulerTriedLoadAlgo = false;
 
     NCCLCHECK(mscclSchedulerInit());
@@ -390,8 +393,6 @@ ncclResult_t mscclEnqueueCheck(
     size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op,
     mscclFunc_t func, ncclComm_t comm, hipStream_t stream) {
   mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
-  hipStreamCaptureStatus captureStatus;
-  unsigned long long pid;
 
   threadLocalStatus.savedSchedulerParams.push_back({});
   NCCLCHECK(mscclSetSavedSchedulerParam(
@@ -402,29 +403,16 @@ ncclResult_t mscclEnqueueCheck(
   switch (threadLocalStatus.groupStatus) {
     case mscclNoGroup:
       if (comm->mscclCompatible) {
-        if (stream == (hipStream_t)0) {
-          captureStatus = hipStreamCaptureStatusNone;
-        } else {
-          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:
       if (comm->mscclCompatible) {
-        if (stream == (hipStream_t)0) {
-          captureStatus = hipStreamCaptureStatusNone;
-        } else {
-          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
@@ -432,8 +420,8 @@ ncclResult_t mscclEnqueueCheck(
             break;
           }
         }
-      }
       threadLocalStatus.groupStatus = mscclGroupUnsupportedOp;
+      NCCLCHECK(mscclFallBackSavedParams());
     case mscclGroupUnsupportedOp:
       NCCLCHECK(mscclFallBackSavedParams());
       break;

src/misc/msccl/msccl_setup.cc

@@ -19,6 +19,28 @@
 RCCL_PARAM(MscclEnableDoneEvent, "MSCCL_ENABLE_DONE_EVENT", 1);
 #endif
 
+ncclResult_t mscclGetCaptureStatus(hipStream_t stream) {
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  mscclSavedProxyArgs& savedProxyArgs = mscclGetSavedProxyArgs();
+  cudaStreamCaptureStatus captureStatus;
+  unsigned long long captureId;
+  CUDACHECK(hipStreamGetCaptureInfo_v2(stream, &captureStatus, &captureId, &threadLocalStatus.graph, nullptr, nullptr));
+  if (captureStatus == cudaStreamCaptureStatusActive) {
+    if (savedProxyArgs.count(captureId) == 0) {
+      threadLocalStatus.captureStatus = mscclNewCapture;
+      savedProxyArgs[captureId] = std::vector<struct mscclProxyArg>();
+    } else {
+      INFO(NCCL_INIT|NCCL_NET,"mscclGetCaptureStatus: captureId %llu is same with the previous one\n", captureId);
+      threadLocalStatus.captureStatus = mscclExistingCapture;
+    }
+    threadLocalStatus.captureId = captureId;
+  } else {
+    threadLocalStatus.captureStatus = mscclNoCapture;
+  }
+  INFO(NCCL_INIT|NCCL_NET,"mscclGetCaptureStatus: %d, captureId: %llu, size: %lu\n", threadLocalStatus.captureStatus, threadLocalStatus.captureId, mscclGetSavedProxyArgs()[captureId].size());
+  return ncclSuccess;
+}
+
 ncclResult_t mscclSetupCount(struct mscclAlgo* hostAlgo, ncclComm_t comm, size_t count, ncclDataType_t dataType) {
   mscclStatus& status = mscclGetStatus();
   status.stepSize = comm->buffSizes[hostAlgo->protocol] / NCCL_STEPS;
@@ -55,14 +77,19 @@ ncclResult_t mscclSetupScratch(struct mscclAlgo* hostAlgo, hipStream_t stream) {
 
 ncclResult_t mscclSetupSyncFlags(hipStream_t stream) {
   mscclStatus& status = mscclGetStatus();
-  if (status.workIndex > (1ULL << (8*sizeof(status.workIndex))) - 2 * NCCL_MAX_OPS - 1) {
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  if (threadLocalStatus.captureStatus == mscclNewCapture ||
+      status.workIndex > (1ULL << (8*sizeof(status.workIndex))) - 2 * NCCL_MAX_OPS - 1) {
     CUDACHECK(hipMemsetAsync(status.syncFlags, 0, sizeof(struct mscclFlag) * MSCCL_MAX_NUM_THREAD_BLOCKS, stream));
     status.workIndex = 1; // setting the workIndex back to 1 for next iterations
+    status.graphFirstKernel = false;
   }
   return ncclSuccess;
 }
 
 ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo, ncclComm_t comm) {
+  mscclStatus& status = mscclGetStatus();
+
   // Check whether there is enough channels
   if (hostAlgo->nChannels > comm->nChannels) {
     WARN("MSCCL: number of channels available (%d) less than required (%d)", comm->nChannels, hostAlgo->nChannels);
@@ -88,15 +115,19 @@ ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo, ncclComm_t comm)
 
   // Connect MSCCL connections
   mscclSetIsCallerFlag();
-  NCCLCHECK(ncclTransportP2pSetup(comm, NULL, 0));
+  int highestTransportType = TRANSPORT_P2P;
+  bool needsProxy = false;
+  NCCLCHECK(ncclTransportP2pSetup(comm, NULL, 0, &highestTransportType, &needsProxy));
+  status.needsProxy |= needsProxy;
   mscclClearIsCallerFlag();
 
   INFO(NCCL_INIT, "MSCCL: Setup connections finished, used %ld", allocTracker[comm->cudaDev].totalAllocSize);
   return ncclSuccess;
 }
 
-ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm) {
+static ncclResult_t mscclSetupProxyImpl(struct mscclAlgo* hostAlgo, ncclComm_t comm) {
   mscclStatus& status = mscclGetStatus();
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
   struct ncclProxyOp proxyOp = {};
   proxyOp.connIndex = 0;
   proxyOp.sliceSteps = status.sliceSteps;
@@ -147,6 +178,38 @@ ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm) {
   return ncclSuccess;
 }
 
+static void HIPRT_CB mscclSetupProxyCallback(void *args) {
+  std::vector<struct mscclProxyArg>* params = (std::vector<struct mscclProxyArg>*)args;
+  INFO(NCCL_INIT|NCCL_NET,"mscclSetupProxyCallback: proxy args size: %ld\n", params->size());
+  for (auto &p : *params) {
+    mscclSetupProxyImpl(p.hostAlgo, p.comm);
+  }    
+}
+
+ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm, hipStream_t stream) {
+  mscclStatus& status = mscclGetStatus();
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  mscclSavedProxyArgs& savedProxyArgs = mscclGetSavedProxyArgs();
+  if (threadLocalStatus.captureStatus == mscclNoCapture) {
+    INFO(NCCL_INIT|NCCL_NET,"mscclSetupProxy: no capture\n");
+    NCCLCHECK(mscclSetupProxyImpl(hostAlgo, comm));
+  } else if (status.needsProxy) {
+    INFO(NCCL_INIT|NCCL_NET,"mscclSetupProxy: capture\n");
+    if (savedProxyArgs[threadLocalStatus.captureId].size() == 0) {
+      INFO(NCCL_INIT|NCCL_NET,"mscclSetupProxy: adding callback\n");
+
+      hipGraphNode_t callbackNode;
+      hipHostNodeParams p;
+      p.fn = mscclSetupProxyCallback;
+      auto params = &savedProxyArgs[threadLocalStatus.captureId];
+      p.userData = params;
+      CUDACHECK(hipGraphAddHostNode(&callbackNode, threadLocalStatus.graph, nullptr, 0, &p));
+    }
+    mscclGetSavedProxyArgs()[threadLocalStatus.captureId].emplace_back(hostAlgo, comm);
+  }
+  return ncclSuccess;
+}
+
 static ncclResult_t hostToDevRedOp(
     ncclDevRedOpFull *opFull, ncclRedOp_t op, ncclDataType_t datatype, ncclComm *comm
   ) {
@@ -268,7 +331,7 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
 
   dim3 grid = {(uint32_t)hostAlgo->nBlocks, 1, 1};
   dim3 block = {NCCL_MAX_NTHREADS, 1, 1};
-  ncclDevRedOpFull opFull;
+  ncclDevRedOpFull opFull = {};
   NCCLCHECK(hostToDevRedOp(&opFull, op, dataType, comm));
 
   mscclWork work;
@@ -283,7 +346,8 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
   work.maxAllowedCount = status.maxAllowedCount;
   work.hasReduce = hostAlgo->hasReduce;
   work.redOpArgIsPtr = opFull.scalarArgIsPtr;
-
+  INFO(NCCL_INIT, "MSCCL: Setup Kernel finished");
+  
   void *args[3] = {&comm->devComm, &devAlgo, &work};
   void *func = mscclKernelEntries[(opFull.op * ncclNumTypes + dataType) * NCCL_NUM_PROTOCOLS + hostAlgo->protocol];
   if (enableDoneEvent) {

src/misc/msccl/msccl_status.cc

@@ -4,6 +4,7 @@
  ************************************************************************/
 
 #include "msccl/msccl_status.h"
+#include "msccl/msccl_struct.h"
 
 mscclStatus& mscclGetStatus() {
   static mscclStatus status;
@@ -14,3 +15,8 @@ mscclThreadLocalStatus& mscclGetThreadLocalStatus() {
   static thread_local mscclThreadLocalStatus threadLocalStatus;
   return threadLocalStatus;
 }
+
+mscclSavedProxyArgs& mscclGetSavedProxyArgs() {
+  static mscclSavedProxyArgs savedProxyArgs;
+  return savedProxyArgs;
+}

src/transport.cc

@@ -19,7 +19,7 @@ struct ncclTransport* ncclTransports[NTRANSPORTS] = {
 };
 
 template <int type>
-static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclConnect* connect, int channelId, int peer, int connIndex, int* transportType) {
+static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclConnect* connect, int channelId, int peer, int connIndex, int* transportType, bool* needsProxy) {
   struct ncclPeerInfo* myInfo = comm->peerInfo+comm->rank;
   struct ncclPeerInfo* peerInfo = comm->peerInfo+peer;
   struct ncclConnector* connector = (type == 1) ? comm->channels[channelId].peers[peer]->send + connIndex :
@@ -44,6 +44,7 @@ static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph*
       connector->transportComm = transportComm;
       NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, connect, connector, channelId, connIndex));
       if (transportType) *transportType = t;
+      if (needsProxy) *needsProxy = (transportComm->proxyProgress != NULL);
       return ncclSuccess;
     }
   }
@@ -77,10 +78,11 @@ void dumpData(struct ncclConnect* data, int ndata) {
   }
 }
 
-ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType/*=NULL*/) {
+ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType/*=NULL*/, bool* needsProxy/*=NULL*/) {
   // Stream used during transport setup; need for P2P pre-connect + CUDA Graph
   ncclResult_t ret = ncclSuccess;
   int highestType = TRANSPORT_P2P;  // track highest transport type
+  bool needsProxyResult = false;
   struct ncclConnect** data = (ncclConnect**) malloc(sizeof(ncclConnect*) * comm->nRanks); // Store intermediate send/recvData structs for connect
   struct ncclConnect** recvData = (ncclConnect**) malloc(sizeof(ncclConnect*) * comm->nRanks); // Points to entries inside data for given recv connection within a channel
   struct ncclConnect** sendData = (ncclConnect**) malloc(sizeof(ncclConnect*) * comm->nRanks); // Points to entries inside data for given send connection within a channel
@@ -103,10 +105,11 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
     recvData[i] = data[i];
     int sendChannels = 0, recvChannels = 0;
     int type;
+    bool proxy;
     TIME_START(0);
     for (int c=0; c<MAXCHANNELS; c++) {
       if (recvMask & (1UL<<c)) {
-        NCCLCHECKGOTO(selectTransport<0>(comm, graph, recvData[i]+recvChannels++, c, recvPeer, connIndex, &type), ret, fail);
+        NCCLCHECKGOTO(selectTransport<0>(comm, graph, recvData[i]+recvChannels++, c, recvPeer, connIndex, &type, &proxy), ret, fail);
         if (type > highestType) highestType = type;
       }
     }
@@ -115,8 +118,9 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
     sendData[i] = recvData[i]+recvChannels;
     for (int c=0; c<MAXCHANNELS; c++) {
       if (sendMask & (1UL<<c)) {
-        NCCLCHECKGOTO(selectTransport<1>(comm, graph, sendData[i]+sendChannels++, c, sendPeer, connIndex, &type), ret, fail);
+        NCCLCHECKGOTO(selectTransport<1>(comm, graph, sendData[i]+sendChannels++, c, sendPeer, connIndex, &type, &proxy), ret, fail);
         if (type > highestType) highestType = type;
+        needsProxyResult |= proxy;
       }
     }
     TIME_STOP(1);
@@ -207,6 +211,7 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
   free(recvData);
 
   if (highestTransportType != NULL) *highestTransportType = highestType;
+  if (needsProxy != NULL) *needsProxy = needsProxyResult;
   TIME_PRINT("P2P Setup/Connect");
 exit:
   NCCLCHECK(ncclStrongStreamWaitStream(ncclCudaGraphNone(), &comm->sharedRes->deviceStream, &comm->sharedRes->hostStream));