cleaner order of iteration

Author: brb-nv

cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu

@@ -205,15 +205,14 @@ TargetRanksInfo TargetRanksInfoForDP(
     }
 
     std::vector<int> retRanks;
-    for (int i = peerTPRankStart; i < peerTPRankEnd; i++)
+    for (int i = peerCPRankStart; i < peerCPRankEnd; i++)
     {
-        for (int j = peerCPRankStart; j < peerCPRankEnd; j++)
+        for (int j = peerTPRankStart; j < peerTPRankEnd; j++)
         {
             for (int k = peerPPRankStart; k < peerPPRankEnd; k++)
             {
-                // Rank formula: ppRank * (tpNum * cpNum) + tpRank * cpNum + cpRank
-                // where i=tpRank, j=cpRank, k=ppRank
-                int irank = (k * peerTPNum * peerCPNum) + (i * peerCPNum) + j;
+                // Rank formula: ppRank * (tpNum * cpNum) + tpRank * cpNum + cpRank.
+                int irank = (k * peerTPNum * peerCPNum) + (j * peerCPNum) + i;
                 retRanks.push_back(irank);
             }
         }

tensorrt_llm/_torch/device_mesh.py

@@ -118,8 +118,8 @@ def build_mesh(self):
                 "DeviceMesh creation requested but torch.distributed process group "
                 "has not been initialised.")
 
-        # Dimensions go from slowest-varying (outermost) to fastest-varying (innermost)
-        # Layout: pp is outermost, then tp, then cp is innermost (consecutive)
+        # Dimensions go from slowest-varying (outermost) to fastest-varying (innermost).
+        # Layout: pp is outermost, then tp, then cp is innermost (consecutive).
         dims = ["pp", "tp"]
         shape = [self.pp_size, self.tp_size]
 

tensorrt_llm/mapping.py

@@ -292,13 +292,13 @@ def has_cp(self):
         return self.cp_size > 1
 
     def prev_cp_rank(self):
-        # cp ranks are consecutive, so prev is rank - 1 with wraparound within cp group
+        # cp ranks are consecutive, so prev is rank - 1 with wraparound within cp group.
         if self.cp_rank == 0:
             return self.rank + self.cp_size - 1
         return self.rank - 1
 
     def next_cp_rank(self):
-        # cp ranks are consecutive, so next is rank + 1 with wraparound within cp group
+        # cp ranks are consecutive, so next is rank + 1 with wraparound within cp group.
         if self.cp_rank == self.cp_size - 1:
             return self.rank - self.cp_size + 1
         return self.rank + 1
@@ -596,15 +596,15 @@ def _init_parallel_groups(self):
             ranks = range(i, self.world_size, self.tp_size * self.cp_size)
             self.pp_groups.append(list(ranks))
 
-        # init cp group (consecutive ranks within each tp slice)
+        # init cp group (consecutive ranks within each tp slice).
         for i in range(self.pp_size):
             for j in range(self.tp_size):
                 ranks = range(
                     i * self.tp_size * self.cp_size + j * self.cp_size,
                     i * self.tp_size * self.cp_size + (j + 1) * self.cp_size)
                 self.cp_groups.append(list(ranks))
 
-        # init tp group (interleaved ranks with stride of cp_size)
+        # init tp group (interleaved ranks with stride of cp_size).
         for i in range(self.pp_size):
             for j in range(self.cp_size):
                 ranks = range(i * self.tp_size * self.cp_size + j,

tensorrt_llm/models/modeling_utils.py

@@ -740,7 +740,7 @@ def from_checkpoint(
 
         rank = config.mapping.rank
         if config.mapping.cp_size > 1:
-            # cp_tp_pp rank -> tp_pp rank: because different cp ranks share the same ckpt
+            # cp_tp_pp rank -> tp_pp rank: because different cp ranks share the same ckpt.
             tp_size = config.mapping.tp_size
             cp_size = config.mapping.cp_size
             rank = (rank % (tp_size * cp_size)) // cp_size + rank // (