Use user passed parameter to determine if it is in combine

bingzheliu · meta-codesync[bot] · commit 683a4418dd3e · 2025-12-09T17:30:35.000-08:00
Summary: Instead of calculate if it is 2nd a2a inside a2avD, we use user passing parameter combine to determine if it is 2nd a2a.

Reviewed By: cenzhaometa

Differential Revision: D86346838

fbshipit-source-id: 8567fd95877b36c179ece264b4e19db8a958ba1c
diff --git a/comms/ctran/algos/AllToAll/AllToAllvDynamic.cuh b/comms/ctran/algos/AllToAll/AllToAllvDynamic.cuh
@@ -164,7 +164,7 @@ __device__ __forceinline__ void recvImplNonContig(
     int groupIdx,
     int ngroups,
     size_t maxRecvcount,
-    bool nonContigIndices) {
+    bool combine) {
   const auto localRank = statex->localRank();
   const auto nLocalRanks = statex->nLocalRanks();
 
@@ -189,7 +189,7 @@ __device__ __forceinline__ void recvImplNonContig(
     // writes it to the recvCountsTmpbufGPU buffer.
     devSyncWaitStep(sync, groupIdx, 0);
     mySendIndicesBlockLength = recvIndicesPeerAllToAllvDynamicBufsMap[0];
-    if (threadIdx.x == 0 && groupIdx == 0 && !nonContigIndices) {
+    if (threadIdx.x == 0 && groupIdx == 0 && !combine) {
       for (int i = 0; i < sendcountsLength; i++) {
         recvCountsTmpbufGPU[recvPeerGlobal * sendcountsLength + i] =
             recvcountsPeerAllToAllvDynamicBufsMap[i];
@@ -198,7 +198,7 @@ __device__ __forceinline__ void recvImplNonContig(
     devSyncSetStep(sync, groupIdx, CTRAN_ALGO_STEP_RESET);
 
     size_t recvOffsets = 0, lastRecvIndex = 0;
-    if (nonContigIndices) {
+    if (combine) {
       lastRecvIndex = sendcountsLength * statex->rank() / statex->nRanks();
     }
     for (int i = 0; i < mySendIndicesBlockLength; i++) {
@@ -304,7 +304,7 @@ __device__ __forceinline__ void selfCopyNonContig(
     int groupIdx,
     bool groupType,
     size_t maxRecvcount,
-    bool nonContigIndices) {
+    bool combine) {
   // Now we calculate the startSendIndex on-the-fly,
   // which may not be efficient. If the inputChunkCountPerRank can be
   // on CPU, we can calculate it on CPU and pass it to GPU.
@@ -317,7 +317,7 @@ __device__ __forceinline__ void selfCopyNonContig(
     startSendIndex += inputChunkCountPerRank[i];
   }
 
-  if (!nonContigIndices && groupIdx == 0 && groupType == GROUP_RECV) {
+  if (!combine && groupIdx == 0 && groupType == GROUP_RECV) {
     ctranKernCopy<size_t>(
         sendcounts,
         recvCountsTmpbufGPU + rank * sendcountsLength,
@@ -326,7 +326,7 @@ __device__ __forceinline__ void selfCopyNonContig(
         1);
   }
 
-  if (nonContigIndices) {
+  if (combine) {
     curOffsetIndex = sendcountsLength * rank / nRanks;
   }
 
@@ -383,7 +383,7 @@ __device__ __forceinline__ void ncclKernelAllToAllvDynamicCommon(
     int* flag,
     CtranKernelAllToAllvDynamicArgs args,
     ALGOTYPE algoType,
-    bool nonContigIndices = false) {
+    bool combine = false) {
   const auto gtIdx = blockDim.x * blockIdx.x + threadIdx.x;
 
   const auto rank = statex->rank();
@@ -461,7 +461,7 @@ __device__ __forceinline__ void ncclKernelAllToAllvDynamicCommon(
         groupIdx,
         groupType,
         args.nonContig.maxRecvcount,
-        nonContigIndices);
+        combine);
     if (groupType == GROUP_RECV) {
       recvImplNonContig(
           recvbuffs,
@@ -471,7 +471,7 @@ __device__ __forceinline__ void ncclKernelAllToAllvDynamicCommon(
           groupIdx,
           ngroups,
           args.nonContig.maxRecvcount,
-          nonContigIndices);
+          combine);
     } else {
       sendImplNonContig(
           sendbuffs,
@@ -510,7 +510,7 @@ __device__ __forceinline__ void ncclKernelAllToAllvDynamicCommon(
   // Copy back to recvcounts for DYNAMIC and DYNAMIC_SPLIT
   // or if it is first a2a for DYNAMIC_SPLIT_NON_CONTIG
   if (groupIdx == 0 && groupType == GROUP_RECV &&
-      (algoType != DYNAMIC_SPLIT_NON_CONTIG || !nonContigIndices)) {
+      (algoType != DYNAMIC_SPLIT_NON_CONTIG || !combine)) {
     ctranKernCopy<size_t>(
         recvCountsTmpbufGPU,
         reinterpret_cast<size_t*>(args.actualRecvcounts),
@@ -528,7 +528,7 @@ __device__ __forceinline__ void ncclKernelAllToAllvDynamicCommon(
 template <typename T>
 __device__ __forceinline__ void generateSendbuffs(
     CtranKernelAllToAllvDynamicArgs& args,
-    bool nonContigIndices = false) {
+    bool combine = false) {
   const auto gtIdx = blockDim.x * blockIdx.x + threadIdx.x;
   const size_t* sendSplitLengths = (size_t*)args.sendcounts;
   args.split.sendbuffsPtrShmDev =
@@ -548,7 +548,7 @@ __device__ __forceinline__ void generateSendbuffs(
       // and hence need to reset the sendbuff offset.
       // The length of each rank is equal to maxsendcounts/ranks.
       // i / numCountsPerRank is the rank number.
-      if (nonContigIndices && (i % numCountsPerRank == 0)) {
+      if (combine && (i % numCountsPerRank == 0)) {
         sendbuffsGPU[i] = sendbuffsGPU[0] +
             (args.nonContig.maxSendcount / statex->nRanks()) *
                 (i / numCountsPerRank);
@@ -592,14 +592,14 @@ __global__ void ncclKernelAllToAllvDynamicSplitNonContig(
     CtranKernelAllToAllvDynamicArgs args) {
   devStateLoadToShm(devState);
 
-  bool nonContigIndices = false;
   int totalSendIndicesLength = 0;
   for (int i = 0; i < statex->nRanks(); i++) {
     totalSendIndicesLength += args.nonContig.inputChunkCountPerRank[i];
   }
-  nonContigIndices = (totalSendIndicesLength < args.sendcountsLength);
 
-  generateSendbuffs<T>(args, nonContigIndices);
+  bool combine = args.nonContig.combine;
+
+  generateSendbuffs<T>(args, combine);
 
   ctranKernCopy<size_t>(
       args.nonContig.inputChunkIndices,
@@ -630,7 +630,7 @@ __global__ void ncclKernelAllToAllvDynamicSplitNonContig(
   }
 
   ncclKernelAllToAllvDynamicCommon<T>(
-      flag, args, DYNAMIC_SPLIT_NON_CONTIG, nonContigIndices);
+      flag, args, DYNAMIC_SPLIT_NON_CONTIG, combine);
 }
 
 #define DECL_CTRAN_ALLTOALLVDYNAMIC_KERN(T)               \
diff --git a/comms/ctran/algos/AllToAll/AllToAllvDynamicCommon.cc b/comms/ctran/algos/AllToAll/AllToAllvDynamicCommon.cc
@@ -31,7 +31,8 @@
               ibPutReqs,                                                      \
               ibRecvCtrlReqs,                                                 \
               maxRecvcount,                                                   \
-              maxSendcount));                                                 \
+              maxSendcount,                                                   \
+              combine));                                                      \
     } else {                                                                  \
       FB_COMMCHECK(peerPutContig(                                             \
           comm,                                                               \
@@ -218,7 +219,8 @@ commResult_t ctranAllToAllvDynamicIbImpl(
     CtranComm* comm,
     std::unique_ptr<CtranMapperTimestamp> timestamp,
     KernelElem* elem,
-    void* recvbuff) {
+    void* recvbuff,
+    bool combine) {
   const auto& statex = comm->statex_;
   const int myRank = statex->rank();
   const int nRanks = statex->nRanks();
@@ -391,7 +393,8 @@ commResult_t opIbImpl(
       comm,
       std::move(timestamp),
       op->alltoallv_dynamic.kElem,
-      op->alltoallv_dynamic.recvbuff);
+      op->alltoallv_dynamic.recvbuff,
+      op->alltoallv_dynamic.combine);
 }
 
 commResult_t setupGpeOp(
@@ -406,7 +409,8 @@ commResult_t setupGpeOp(
     uint64_t opCount,
     std::vector<std::unique_ptr<struct OpElem>>& opGroup,
     KernelElem* elem,
-    void* recvbuff) {
+    void* recvbuff,
+    bool combine) {
   std::unique_ptr<struct OpElem> op =
       std::unique_ptr<struct OpElem>(new OpElem(opType, comm, opCount));
   op->alltoallv_dynamic.sendbuffs = sendbuffs;
@@ -417,6 +421,7 @@ commResult_t setupGpeOp(
   op->alltoallv_dynamic.maxRecvcount = maxRecvcount;
   op->alltoallv_dynamic.kElem = elem;
   op->alltoallv_dynamic.recvbuff = recvbuff;
+  op->alltoallv_dynamic.combine = combine;
 
   opGroup.push_back(std::move(op));
 
diff --git a/comms/ctran/algos/AllToAll/AllToAllvDynamicCommon.h b/comms/ctran/algos/AllToAll/AllToAllvDynamicCommon.h
@@ -45,7 +45,8 @@ commResult_t setupGpeOp(
     uint64_t opCount,
     std::vector<std::unique_ptr<struct OpElem>>& opGroup,
     KernelElem* elem,
-    void* recvbuff = nullptr);
+    void* recvbuff = nullptr,
+    bool combine = false);
 
 template <typename PerfConfig = DefaultPerfCollConfig>
 commResult_t peerPutNonContig(
@@ -64,6 +65,7 @@ commResult_t peerPutNonContig(
     std::vector<std::unique_ptr<CtranMapperRequest>>& ibRecvCtrlReqs,
     size_t maxRecvcount,
     size_t maxSendcount,
+    bool combine,
     bool skipWaitRecvCtrl = false) {
   // Prepare basic info for nonContig send
   size_t* sendIndices = reinterpret_cast<size_t*>(comm->ctran_->algo->getTmpBuf(
@@ -97,15 +99,14 @@ commResult_t peerPutNonContig(
   for (int r = 0; r < comm->statex_->nRanks(); r++) {
     totalBlock += sendIndicesBlockLengthsTmpbufCPU[r];
   }
-  bool nonContigIndices = (totalBlock < sendcountsLength);
 
   // Calculate the offset of each recvbuff, considering if it is 1st or 2nd
   // all2allv.
   std::vector<size_t> remoteRecvBuffsBytesOffset(sendcountsLength);
   remoteRecvBuffsBytesOffset[0] = 0;
   int numCountsPerRank = sendcountsLength / nRanks;
   for (int i = 1; i < sendcountsLength; i++) {
-    if (nonContigIndices && (i % numCountsPerRank == 0)) {
+    if (combine && (i % numCountsPerRank == 0)) {
       remoteRecvBuffsBytesOffset[i] = 0;
     } else {
       remoteRecvBuffsBytesOffset[i] += remoteRecvBuffsBytesOffset[i - 1] +
@@ -153,7 +154,7 @@ commResult_t peerPutNonContig(
       // Allgather sendcounts
       // Skip sending sendcounts if it is second all2allv.
       // TODO: using hints instead of nonContigIndices to determine this.
-      if (!nonContigIndices) {
+      if (!combine) {
         puts.emplace_back(
             CtranMapperPutMsg{
                 .sbuf = reinterpret_cast<size_t*>(sendCountsTmpbufGPU),
diff --git a/comms/ctran/algos/AllToAll/AllToAllvDynamicPImpl.cc b/comms/ctran/algos/AllToAll/AllToAllvDynamicPImpl.cc
@@ -31,6 +31,7 @@ using ctran::alltoallvdynamicp::PersistArgs;
             completedIbRecvCtrlReqs,                                          \
             pArgs->maxRecvCount,                                              \
             pArgs->maxSendCount,                                              \
+            op->alltoallv_dynamic.combine,                                    \
             /* skipWaitRecvCtrl */ true));                                    \
     /* Wait for all puts to complete */                                       \
     for (auto& req : ibPutReqs) {                                             \
diff --git a/comms/ctran/algos/AllToAll/AllToallvDynamicSplitNonContig.cc b/comms/ctran/algos/AllToAll/AllToallvDynamicSplitNonContig.cc
@@ -84,6 +84,7 @@ commResult_t ctranAlltoallvDynamicSplitNonContig(
       maxRecvcount;
   config.args.collective.alltoallv_dynamic.nonContig.maxSendcount =
       maxSendcount;
+  config.args.collective.alltoallv_dynamic.nonContig.combine = combine;
 
   if (recvbuff != nullptr) {
     for (int i = 0; i < comm->statex_->nRanks(); i++) {
@@ -106,7 +107,8 @@ commResult_t ctranAlltoallvDynamicSplitNonContig(
       opCount,
       opGroup,
       elem,
-      recvbuff));
+      recvbuff,
+      combine));
 
   XCHECK(alltoallvDynamicSplitNonContigKerns.contains(datatype))
       << "alltoallvDynamicSplitNonContigKerns does not contain datatype "
diff --git a/comms/ctran/gpe/CtranGpe.h b/comms/ctran/gpe/CtranGpe.h
@@ -163,6 +163,7 @@ struct OpElem {
       KernelElem* kElem;
       // Persistent args for persistent alltoallv_dynamic.
       void* pArgs;
+      bool combine;
     } alltoallv_dynamic;
     struct {
       const void* sendbuff;
diff --git a/comms/ctran/gpe/CtranGpeDev.h b/comms/ctran/gpe/CtranGpeDev.h
@@ -260,6 +260,7 @@ struct CtranKernelAllToAllvDynamicArgs {
       size_t maxInputChunkCountPerRank{0};
       size_t maxRecvcount{0};
       size_t maxSendcount{0};
+      bool combine;
     } nonContig;
     struct {
     } contig;
