From aebafe7f55f306c963278e55a7278ddef6d446af Mon Sep 17 00:00:00 2001
From: Tianxing Wu <chi0chu319@gmail.com>
Date: Thu, 11 Sep 2025 10:32:12 +0000
Subject: [PATCH] xcd remap

---
 .../ops/gemm/kernel/gemm_tile_partitioner.hpp | 27 +++++++++++++++++++
 .../moe_flatmm/kernel/moe_flatmm_kernel.hpp   | 15 +++++++----
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp b/include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp
index 28e8bee908..c606160bc4 100644
--- a/include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp
+++ b/include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp
@@ -265,6 +265,32 @@ struct GemmSpatiallyLocalTilePartitioner
         return integer_divide_ceil(K, KPerBlock);
     }
 
+    CK_TILE_HOST_DEVICE static auto RemapXCD(int pid, int GRID_MN, int NUM_XCDS = 8) {
+        // Number of pids per XCD in the new arrangement
+        int pids_per_xcd = (GRID_MN + NUM_XCDS - 1) / NUM_XCDS;
+
+        // When GRID_MN cannot divide NUM_XCDS, some xcds will have
+        // pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
+        // We calculate the number of xcds that have pids_per_xcd pids as tall_xcds
+        int tall_xcds = GRID_MN % NUM_XCDS;
+        tall_xcds = (tall_xcds == 0) ? NUM_XCDS : tall_xcds;
+
+        // Compute current XCD and local pid within the XCD
+        int xcd = pid % NUM_XCDS;
+        int local_pid = pid / NUM_XCDS;
+
+        // Calculate new pid based on the new grouping
+        if (xcd < tall_xcds) {
+            pid = xcd * pids_per_xcd + local_pid;
+        } else {
+            pid = tall_xcds * pids_per_xcd
+                + (xcd - tall_xcds) * (pids_per_xcd - 1)
+                + local_pid;
+        }
+
+        return pid;
+    }
+
     /**
      * @brief Calculate workgroup 1D index mapping into 2D output C-tile space.
      *
@@ -276,6 +302,7 @@ struct GemmSpatiallyLocalTilePartitioner
     {
         const auto M0 = integer_divide_ceil(M, MPerBlock);
         const auto N0 = integer_divide_ceil(N, NPerBlock);
+        // index_t block_1d_id = RemapXCD(_block_1d_id, M0 * N0)
 
         if(M0 == 1)
         {
diff --git a/include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp b/include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp
index d0886b6a23..c158be1ea1 100644
--- a/include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp
+++ b/include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp
@@ -753,12 +753,20 @@ struct MoeFlatmmKernel
     template <class MoeFlatmmKernelArgs>
     CK_TILE_DEVICE void operator()(MoeFlatmmKernelArgs kargs) const
     {
+        // total number of tokens: sorted tokens + delimiter tokens + trailing padding tokens
+        // we launch the grid based on the total number of tokens which needs to be static
         int partition_idx       = blockIdx.x;
-        int total_work_tile_cnt = TilePartitioner::GridSize(kargs.M, kargs.N);
+        auto valid_token_cnt = kargs.p_max_token_id[0]; // sorted tokens + delimiter tokens
+        int total_valid_tile_cnt = TilePartitioner::GridSize(valid_token_cnt, kargs.N);
+        auto tilePartitioner = TilePartitioner{valid_token_cnt, kargs.N};
         do
         {
+            if (partition_idx >= total_valid_tile_cnt) {
+                return; // early exit for trailing padding tokens
+            }
+            partition_idx = tilePartitioner.RemapXCD(partition_idx, total_valid_tile_cnt);
             const auto [block_offset_m, block_offset_n] =
-                TilePartitioner{kargs.M, kargs.N}.GetOutputTileIndex(partition_idx);
+                tilePartitioner.GetOutputTileIndex(partition_idx);
 
             this->operator()(kargs, block_offset_m, block_offset_n);
             partition_idx += gridDim.x;
@@ -772,7 +780,6 @@ struct MoeFlatmmKernel
         // const auto [iM, iN]   = TilePartitioner{kargs.M, kargs.N}.GetOutputTileIndex(blockIdx.x);
         const index_t coord_m = __builtin_amdgcn_readfirstlane(iM * TilePartitioner::MPerBlock);
         const index_t coord_n = __builtin_amdgcn_readfirstlane(iN * TilePartitioner::NPerBlock);
-        const index_t max_token_id = kargs.p_max_token_id[0];
         // allocate LDS
         __shared__ char smem_ptr_ping[GetSmemPingSize()];
         __shared__ char smem_ptr_pong[GetSmemPongSize()];
@@ -800,8 +807,6 @@ struct MoeFlatmmKernel
             return gather_token_id;
         };
         
-        if(coord_m >= max_token_id)
-            return;
         static_for<0, DramMRepeat, 1>{}([&](auto m0) {
             const auto row_idx =
                 coord_m + m0 * (TilePartitioner::MPerBlock / DramMRepeat) + a_coord[I0];