submission 495528

  #!POPCORN leaderboard nvfp4_group_gemm
  #!POPCORN gpu NVIDIA
   
+ from __future__ import annotations
+  
+ from functools import lru_cache
+ from typing import cast
+  
  import torch
- from task import input_t, output_t
  from torch.utils.cpp_extension import load_inline
   
+ from task import input_t, output_t
+  
  """
  g: 8; k: [7168, 7168, 7168, 7168, 7168, 7168, 7168, 7168]; m: [80, 176, 128, 72, 64, 248, 96, 160]; n: [4096, 4096, 4096, 4096, 4096, 4096, 4096, 4096]; seed: 1111
- ⏱ 89.3 ± 0.09 µs
- ⚡ 88.7 µs 🐌 89.6 µs
+  ⏱ 47.5 ± 0.00 µs
+  ⚡ 47.4 µs 🐌 47.5 µs
   
  g: 8; k: [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048]; m: [40, 76, 168, 72, 164, 148, 196, 160]; n: [7168, 7168, 7168, 7168, 7168, 7168, 7168, 7168]; seed: 1111
- ⏱ 82.7 ± 0.05 µs
- ⚡ 82.7 µs 🐌 82.8 µs
+  ⏱ 45.0 ± 0.04 µs
+  ⚡ 44.6 µs 🐌 45.3 µs
   
  g: 2; k: [4096, 4096]; m: [192, 320]; n: [3072, 3072]; seed: 1111
- ⏱ 29.5 ± 0.03 µs
- ⚡ 29.1 µs 🐌 29.8 µs
+  ⏱ 14.3 ± 0.01 µs
+  ⚡ 13.9 µs 🐌 14.6 µs
   
  g: 2; k: [1536, 1536]; m: [128, 384]; n: [4096, 4096]; seed: 1111
- ⏱ 16.4 ± 0.02 µs
- ⚡ 16.0 µs 🐌 16.6 µs
+  ⏱ 10.5 ± 0.01 µs
+  ⚡ 10.4 µs 🐌 10.7 µs
  """
   
  CUDA_SRC = """
  #include <vector>
  #include <unordered_map>
  #include <cstdint>
- #include <cstdio>
  #include <cuda.h>
  #include <cudaTypedefs.h>
  #include <cuda_runtime.h>
  ⋯ 13 unchanged lines
    } while (0)
   
  static inline uint64_t hash_combine_u64(uint64_t h, uint64_t x) {
-   // 64-bit FNV-1a variant
    h ^= x;
    h *= 1099511628211ULL;
    return h;
  ⋯ 2 unchanged lines
  constexpr int WARP_SIZE = 32;
  constexpr int MMA_K = 64;
   
- // Cache hints
- constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;
  constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;
  constexpr uint64_t EVICT_LAST = 0x14F0000000000000;
   
- // Work item for persistent kernel
  struct WorkItem {
    int problem_idx;
    int tile_m;
    int tile_n;
  };
   
- // Global Problem Info stored in Global Memory
+ // Per-problem metadata.
+ // Align on 128 byte boundary, useful since this is read by many CTAs.
  struct __align__(128) ProblemInfo {
    CUtensorMap A_tmap;
    CUtensorMap B_tmap;
  ⋯ 2 unchanged lines
    const char* SFB_ptr;
    half* C_ptr;
    int M, N, K;
-   int64_t Cs0, Cs1, Cs2;
+   int64_t Cs0, Cs1;
  };
   
- __device__ inline
- constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };
+ // tcgen05 descriptors encode shared-memory addresses in 16-byte units.
+ // Mask to the HW-supported address width and drop the 16B alignment bits.
+ __device__ inline constexpr uint64_t desc_encode(uint64_t x) {
+   return (x & 0x3'FFFFULL) >> 4ULL;
+ }
   
- __device__
- uint32_t elect_sync() {
+ // elect.sync: use this to have a single lane issue TMA/tcgen05 instructions while the
+ // whole warp stays converged.
+ __device__ uint32_t elect_sync() {
    uint32_t pred = 0;
    asm volatile(
      "{\\n\\t"
  ⋯ 7 unchanged lines
    return pred;
  }
   
+ __device__ inline uint32_t get_cluster_ctarank() {
+   uint32_t rank;
+   asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank));
+   return rank;
+ }
+  
+ __device__ inline void cluster_sync() {
+   asm volatile("barrier.cluster.arrive;" ::: "memory");
+   asm volatile("barrier.cluster.wait;" ::: "memory");
+ }
+  
+ // Shared-memory mbarrier helpers.
+ // Used for:
+ // - TMA completion barrier: consumer waits for bytes to arrive in shared memory.
+ // - Stage reuse barrier: producer waits until MMA is done with a stage before
+ //   overwriting it.
  __device__ inline void mbarrier_init(int mbar_addr, int count) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
  }
   
+ // Program the expected byte count for a TMA stage and arrive.
+ // This must happen before issuing any cp.async.bulk.* that completes to the
+ // barrier.
  __device__ inline void mbarrier_arrive_expect_tx(int mbar_addr, int size) {
-   asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" 
+   asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                :: "r"(mbar_addr), "r"(size) : "memory");
  }
   
  ⋯ 10 unchanged lines
    );
  }
   
+ // TMA: 3D tensor-map load from global -> shared memory.
+ // The (x,y,z) coordinates correspond to the CUtensorMap encoding in
+ // init_AB_tmap_u4.
  template <int CTA_GROUP = 1>
  __device__ inline void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z,
                                           int mbar_addr, uint64_t cache_policy) {
  ⋯ 5 unchanged lines
    );
  }
   
+ __device__ inline int mapa_cta_to_cluster(int cta_addr, int dest_cta) {
+   int cluster_addr;
+   asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(cluster_addr) : "r"(cta_addr), "r"(dest_cta));
+   return cluster_addr;
+ }
+  
+ __device__ inline void mbarrier_arrive_cluster(int mbar_cluster_addr) {
+   asm volatile("mbarrier.arrive.shared::cluster.b64 _, [%0];"
+     :: "r"(mbar_cluster_addr) : "memory");
+ }
+  
+ // Cluster multicast variant of 3D TMA.
+ // dst and mbar_addr are in shared::cluster address space.
+ // The multicast mask selects which CTAs in the cluster receive the data.
+ __device__ inline void tma_3d_gmem2smem_multicast(int dst, const void *tmap_ptr, int x, int y, int z,
+                                                    int mbar_addr, uint16_t multicast_mask) {
+   asm volatile(
+     "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::1 "
+     "[%0], [%1, {%2, %3, %4}], [%5], %6;"
+     :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "h"(multicast_mask)
+     : "memory"
+   );
+ }
+  
+ // Linear bulk copy global -> shared.
+ // Used for scale-factor tensors (SFA/SFB).
  __device__ inline void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
    asm volatile(
      "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint "
  ⋯ 30 unchanged lines
    );
  }
   
- // TMEM load helpers
  struct SHAPE {
    static constexpr char _16x256b[] = ".16x256b";
  };
  ⋯ 36 unchanged lines
    );
  }
   
- // ============================================================================
- // KERNEL CONFIGURATION
- // ============================================================================
-  
  constexpr int TMA_BLOCK_N = 128;
  constexpr int TMA_BLOCK_K = 256;
- constexpr int NUM_STAGES = 4;
  constexpr int TMA_NUM_WARPS = 8;
  constexpr int MMA_M = 128;
- constexpr int MBAR_BYTES = ((2 * NUM_STAGES * 8 + 63) & ~63);
   
  constexpr int LOW_M_THRESHOLD = 96;
   
- // Warp assignments (8 warps total):
- // Warp 0-3: epilogue helpers
- // Warp 4: TMA producer
- // Warp 5: MMA consumer (single warp issues tcgen05.mma.cta_group::1)
- // Warp 6-7: additional helpers
  constexpr int TMA_WARP = 4;
+ // Use a second (otherwise idle) warp to issue B/SFB TMA in parallel.
+ // This cuts producer-side latency and reduces MMA-side barrier stalls.
+ constexpr int TMA_WARP_B = 6;
  constexpr int MMA_WARP = 5;
   
- // ============================================================================
- // EPILOGUE
- // ============================================================================
-  
+ // Epilogue: read fp32 accumulators from TMEM (tcgen05.ld) and store fp16 C.
+ // Only threads with tid < BLOCK_M participate; this maps 4 warps (0..3) to the
+ // 128 output rows, with each warp handling a 32-row stripe.
  template <int BLOCK_M, int BLOCK_N, bool LOW_M_EPILOGUE>
  __device__ __forceinline__ void epilogue_store(
      const ProblemInfo& prob,
  ⋯ 25 unchanged lines
      m_iters = (remaining <= 16) ? 1 : 2;
    }
   
+   // Lane mapping: each lane owns two columns (half2) and one of 8 rows.
    const int lane_row = lane_id >> 2;
    const int lane_col = (lane_id & 3) * 2;
+  
+   // We load/store in 128-column halves so tcgen05.ld has a fixed shape.
    constexpr int HALF_N = 128;
    const int halves = BLOCK_N / HALF_N;
   
  ⋯ 1 unchanged lines
      for (int half_idx = 0; half_idx < halves; ++half_idx) {
        float tmp[HALF_N / 2];
        const int col_base = half_idx * HALF_N;
+       // TMEM coordinates are relative to the CTA's output tile.
        tcgen05_ld_16x256b_x16(tmp, warp_id * 32 + m * 16, col_base);
        asm volatile("tcgen05.wait::ld.sync.aligned;\\n");
   
  ⋯ 86 unchanged lines
    }
  }
   
- // ============================================================================
- // MAIN KERNEL
- // ============================================================================
-  
- template <bool PERSISTENT, int BLOCK_M, int BLOCK_N>
+ template <bool PERSISTENT, int BLOCK_M, int BLOCK_N, int NS, int CLUSTER_SIZE = 1>
  __global__ __launch_bounds__(TMA_NUM_WARPS * WARP_SIZE)
  void grouped_gemm_kernel_v4(
    const ProblemInfo* __restrict__ global_probs,
    const WorkItem* __restrict__ work_items,
-   int num_items,
-   int* __restrict__ work_counter
+   int num_items
  ) {
    constexpr int TMA_A_SMEM_BYTES = BLOCK_M * (TMA_BLOCK_K / 2);
    constexpr int TMA_B_SMEM_BYTES = BLOCK_N * (TMA_BLOCK_K / 2);
  ⋯ 4 unchanged lines
    const int lane_id = tid % WARP_SIZE;
    const int warp_id = tid / WARP_SIZE;
   
-   // Shared Memory Setup
+   uint32_t cta_rank = 0;
+   if constexpr (CLUSTER_SIZE > 1) {
+     cta_rank = get_cluster_ctarank();
+   }
+  
+   // Shared memory is used as a multi-stage ring buffer.
+   // Per stage: [A tile][B tile][SFA][SFB]. After all stages we place mbarriers.
    extern __shared__ __align__(1024) char smem_ptr[];
    const int smem_base = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
   
-   // Offsets within each stage
    constexpr int B_off = TMA_A_SMEM_BYTES;
    constexpr int SFA_off = B_off + TMA_B_SMEM_BYTES;
    constexpr int SFB_off = SFA_off + TMA_SFA_SMEM_BYTES;
   
-   // Mbarriers
-   const int mbar_base = smem_base + STAGE_SIZE * NUM_STAGES;
-   
-   // TMEM addresses
+   const int mbar_base = smem_base + STAGE_SIZE * NS;
+  
+   // TMEM allocation is in columns. We need 2 columns per output column because
+   // accumulators are fp32.
    constexpr int TMEM_COLS = BLOCK_N * 2;
    constexpr int SFA_tmem = BLOCK_N;
    constexpr int SFB_tmem = SFA_tmem + 4 * (TMA_BLOCK_K / MMA_K);
  ⋯ 2 unchanged lines
                             | ((uint32_t)BLOCK_N >> 3U << 17U)
                             | ((uint32_t)MMA_M >> 7U << 27U);
   
-   // Allocate TMEM
    if (warp_id == 0) {
      asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem_base), "r"(TMEM_COLS));
-   }
-   else if (warp_id == 1 && elect_sync()) {
+   } else if (warp_id == 1 && elect_sync()) {
+     // Best-effort tensormap prefetch for the first few work items.
+  
      for (int i = 0; i < num_items && i < 8; ++i) {
        const ProblemInfo* prob = &global_probs[work_items[i].problem_idx];
        asm volatile("prefetch.tensormap [%0];" :: "l"(&prob->A_tmap) : "memory");
  ⋯ 6 unchanged lines
    }
    __syncthreads();
   
-   // Persistent work counter
    __shared__ int shared_work_idx;
   
-   // Get first work item
-   int work_idx;
-   if constexpr (PERSISTENT) {
+   int work_idx = blockIdx.x;
+  
+   if (work_idx < num_items) {
      if (tid == 0) {
-       shared_work_idx = atomicAdd(work_counter, 1);
+       for (int i = 0; i < NS; ++i) {
+         // mbarrier[stage]: TMA completion barrier.
+         // Two producer warps arrive (A/SFA and B/SFB).
+         mbarrier_init(mbar_base + i * 8, 2);
+  
+         // mbarrier[NS+stage]: stage reuse barrier.
+         // The MMA warp commits once per stage.
+         mbarrier_init(mbar_base + (NS + i) * 8, 1);
+  
+         if constexpr (CLUSTER_SIZE > 1) {
+           if (cta_rank == 0) {
+             // Only CTA rank 0 initializes the cluster-wide barrier used to
+             // guard multicast stage reuse.
+             mbarrier_init(mbar_base + (2*NS + i) * 8, CLUSTER_SIZE);
+           }
+         }
+       }
+       asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
      }
      __syncthreads();
-     work_idx = shared_work_idx;
-   } else {
-     work_idx = blockIdx.x;
    }
   
-   // Main processing loop
+   // Ensure all CTAs have initialized mbarriers before any multicast TMA.
+   if constexpr (CLUSTER_SIZE > 1) {
+     cluster_sync();
+   }
+  
    while (work_idx < num_items) {
      const WorkItem& work = work_items[work_idx];
      const ProblemInfo& prob = global_probs[work.problem_idx];
  ⋯ 3 unchanged lines
      const int K = prob.K;
      const int num_k_iters = K / TMA_BLOCK_K;
   
-     // Initialize mbarriers per tile
-     if (tid == 0) {
-       for (int i = 0; i < NUM_STAGES; ++i) {
-         mbarrier_init(mbar_base + i * 8, 1);
-         mbarrier_init(mbar_base + (NUM_STAGES + i) * 8, 1);
-       }
-       asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
-     }
-     __syncthreads();
-  
-     // ================================================================
-     // PRODUCER WARP: TMA (warp 4)
-     // ================================================================
-     if (warp_id == TMA_WARP && elect_sync()) {
-       // With M-tile-major work ordering, A stays in L2 across N-tiles
+     if ((warp_id == TMA_WARP || warp_id == TMA_WARP_B) && elect_sync()) {
        constexpr uint64_t cache_A = EVICT_LAST;
        constexpr uint64_t cache_B = EVICT_FIRST;
   
-       auto issue_tma = [&](int k_iter, int stage) {
-         const int mbar_addr = mbar_base + stage * 8;
-         const int stage_base = smem_base + stage * STAGE_SIZE;
-         const int off_k = k_iter * TMA_BLOCK_K;
+         const bool do_A = (warp_id == TMA_WARP);
+         const bool do_B = (warp_id == TMA_WARP_B);
   
-         // TMA loads
-         tma_3d_gmem2smem<1>(stage_base, &prob.A_tmap, 0, m_offset, off_k / 256, mbar_addr, cache_A);
-         if constexpr (BLOCK_N == 256) {
-           tma_3d_gmem2smem<1>(stage_base + B_off, &prob.B_tmap_256, 0, n_offset, off_k / 256, mbar_addr, cache_B);
-         } else {
-           tma_3d_gmem2smem<1>(stage_base + B_off, &prob.B_tmap, 0, n_offset, off_k / 256, mbar_addr, cache_B);
-         }
+         auto issue_tma = [&](int k_iter, int stage) {
+           const int mbar_addr = mbar_base + stage * 8;
+           const int stage_base = smem_base + stage * STAGE_SIZE;
+           const int off_k = k_iter * TMA_BLOCK_K;
   
-         // Scale factor loads
-         const int rest_k = K / 16 / 4;
-         const int k_blk = off_k / (16 * 4);
-         const char* SFA_src = prob.SFA_ptr + ((m_offset / 128) * rest_k + k_blk) * 512;
-         tma_gmem2smem(stage_base + SFA_off, SFA_src, TMA_SFA_SMEM_BYTES, mbar_addr, cache_A);
-         if constexpr (BLOCK_N == 256) {
-           constexpr int SFB_HALF_BYTES = 128 * (TMA_BLOCK_K / 16);
-           const char* SFB_src0 = prob.SFB_ptr + ((n_offset / 128) * rest_k + k_blk) * 512;
-           const char* SFB_src1 = prob.SFB_ptr + (((n_offset / 128) + 1) * rest_k + k_blk) * 512;
-           tma_gmem2smem(stage_base + SFB_off, SFB_src0, SFB_HALF_BYTES, mbar_addr, cache_B);
-           tma_gmem2smem(stage_base + SFB_off + SFB_HALF_BYTES, SFB_src1, SFB_HALF_BYTES, mbar_addr, cache_B);
-         } else {
-           const char* SFB_src = prob.SFB_ptr + ((n_offset / 128) * rest_k + k_blk) * 512;
-           tma_gmem2smem(stage_base + SFB_off, SFB_src, TMA_SFB_SMEM_BYTES, mbar_addr, cache_B);
-         }
+           // Program expect_tx before issuing any TMA that completes to this
+           // barrier. A completion arriving before expect_tx is set can leave the
+           // consumer stuck in mbarrier_wait.
+           const int expect_bytes = do_A
+             ? (TMA_A_SMEM_BYTES + TMA_SFA_SMEM_BYTES)
+             : (TMA_B_SMEM_BYTES + TMA_SFB_SMEM_BYTES);
+           mbarrier_arrive_expect_tx(mbar_addr, expect_bytes);
   
-         mbarrier_arrive_expect_tx(mbar_addr, STAGE_SIZE);
-       };
+           if (do_A) {
+             if constexpr (CLUSTER_SIZE > 1) {
+               // Cluster path: CTA rank 0 multicasts A to the whole cluster.
+               // dst and mbarrier are passed as shared::cluster addresses.
+               if (cta_rank == 0) {
+                 uint16_t mc = (1 << CLUSTER_SIZE) - 1;
+                 int cluster_dst = mapa_cta_to_cluster(stage_base, 0);
+                 int cluster_mbar = mapa_cta_to_cluster(mbar_addr, 0);
+                 tma_3d_gmem2smem_multicast(cluster_dst, &prob.A_tmap, 0, m_offset, off_k / 256, cluster_mbar, mc);
+               }
+             } else {
+               tma_3d_gmem2smem<1>(stage_base, &prob.A_tmap, 0, m_offset, off_k / 256, mbar_addr, cache_A);
+             }
   
-       // Pipeline priming: issue first NUM_STAGES TMAs without waiting
-       for (int k_iter = 0; k_iter < NUM_STAGES && k_iter < num_k_iters; k_iter++) {
+             // SFA scale blocks are indexed by (m_tile, k_blk) and stored as
+             // 512B blocks (matching tcgen05_cp_nvfp4 granularity).
+             const int rest_k = K / 16 / 4;
+             const int k_blk = off_k / (16 * 4);
+             const char* SFA_src = prob.SFA_ptr + ((m_offset / 128) * rest_k + k_blk) * 512;
+             tma_gmem2smem(stage_base + SFA_off, SFA_src, TMA_SFA_SMEM_BYTES, mbar_addr, cache_A);
+           } else if (do_B) {
+             if constexpr (BLOCK_N == 256) {
+               tma_3d_gmem2smem<1>(stage_base + B_off, &prob.B_tmap_256, 0, n_offset, off_k / 256, mbar_addr, cache_B);
+             } else {
+               tma_3d_gmem2smem<1>(stage_base + B_off, &prob.B_tmap, 0, n_offset, off_k / 256, mbar_addr, cache_B);
+             }
+  
+             // SFB scale blocks are indexed by (n_tile, k_blk).
+             const int rest_k = K / 16 / 4;
+             const int k_blk = off_k / (16 * 4);
+             if constexpr (BLOCK_N == 256) {
+               constexpr int SFB_HALF_BYTES = 128 * (TMA_BLOCK_K / 16);
+               const char* SFB_src0 = prob.SFB_ptr + ((n_offset / 128) * rest_k + k_blk) * 512;
+               const char* SFB_src1 = prob.SFB_ptr + (((n_offset / 128) + 1) * rest_k + k_blk) * 512;
+               tma_gmem2smem(stage_base + SFB_off, SFB_src0, SFB_HALF_BYTES, mbar_addr, cache_B);
+               tma_gmem2smem(stage_base + SFB_off + SFB_HALF_BYTES, SFB_src1, SFB_HALF_BYTES, mbar_addr, cache_B);
+             } else {
+               const char* SFB_src = prob.SFB_ptr + ((n_offset / 128) * rest_k + k_blk) * 512;
+               tma_gmem2smem(stage_base + SFB_off, SFB_src, TMA_SFB_SMEM_BYTES, mbar_addr, cache_B);
+             }
+           }
+         };
+  
+       for (int k_iter = 0; k_iter < NS && k_iter < num_k_iters; k_iter++) {
          issue_tma(k_iter, k_iter);
        }
   
-       // Steady state: wait for MMA, then issue TMA
-       for (int k_iter = NUM_STAGES; k_iter < num_k_iters; k_iter++) {
-         const int stage = k_iter % NUM_STAGES;
-         const int mma_phase = (k_iter / NUM_STAGES - 1) % 2;
-         mbarrier_wait(mbar_base + (NUM_STAGES + stage) * 8, mma_phase);
+       for (int k_iter = NS; k_iter < num_k_iters; k_iter++) {
+         const int stage = k_iter % NS;
+         const int mma_phase = (k_iter / NS - 1) % 2;
+         if constexpr (CLUSTER_SIZE > 1) {
+             if (do_A && cta_rank == 0) {
+               // A is shared across the cluster via multicast. Before reusing a
+               // ring-buffer stage for the next multicast, rank 0 must wait for
+               // all CTAs to finish consuming the current stage.
+               mbarrier_wait(mbar_base + (2*NS + stage) * 8, mma_phase);
+             } else {
+               mbarrier_wait(mbar_base + (NS + stage) * 8, mma_phase);
+             }
+         } else {
+           mbarrier_wait(mbar_base + (NS + stage) * 8, mma_phase);
+         }
          issue_tma(k_iter, stage);
        }
      }
   
-     // ================================================================
-     // CONSUMER WARP: MMA (warp 5)
-     // Single elected thread issues tcgen05.mma.cta_group::1
-     // ================================================================
      else if (warp_id == MMA_WARP && elect_sync()) {
        auto make_desc_AB = [](int addr) -> uint64_t {
          const int SBO = 8 * 128;
+         // Descriptor encoding is coupled to the shared-memory swizzle and the
+         // tcgen05 operand layout. SBO matches 128B swizzle.
          return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
        };
        auto make_desc_SF = [](int addr) -> uint64_t {
+         // Scale-factor loads use a different stride (16B) but the same address
+         // encoding (16B units).
          const int SBO = 8 * 16;
          return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
        };
   
        for (int k_iter = 0; k_iter < num_k_iters; k_iter++) {
-         const int stage = k_iter % NUM_STAGES;
-         const int tma_phase = (k_iter / NUM_STAGES) % 2;
+         const int stage = k_iter % NS;
+         const int tma_phase = (k_iter / NS) % 2;
          mbarrier_wait(mbar_base + stage * 8, tma_phase);
   
          const int stage_base = smem_base + stage * STAGE_SIZE;
   
-         // Copy scale factors to TMEM
          const uint64_t SF_desc = make_desc_SF(0);
          const uint64_t SFA_desc = SF_desc + ((uint64_t)(stage_base + SFA_off) >> 4ULL);
          const uint64_t SFB_desc = SF_desc + ((uint64_t)(stage_base + SFB_off) >> 4ULL);
   
+         // Copy scale factors from shared memory into TMEM.
+         #pragma unroll
          for (int k = 0; k < TMA_BLOCK_K / MMA_K; k++) {
            tcgen05_cp_nvfp4(SFA_tmem + k * 4, SFA_desc + (uint64_t)k * (512ULL >> 4ULL));
            if constexpr (BLOCK_N == 256) {
  ⋯ 5 unchanged lines
            }
          }
   
-         // Issue MMA
-         for (int k1 = 0; k1 < TMA_BLOCK_K / 256; k1++) {
-           for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
-             uint64_t a_desc = make_desc_AB(stage_base + k1 * BLOCK_M * 128 + k2 * 32);
-             uint64_t b_desc = make_desc_AB(stage_base + B_off + k1 * BLOCK_N * 128 + k2 * 32);
+         // MMA loop over the 256-wide K tile in 64-wide chunks.
+         #pragma unroll
+         for (int k = 0; k < TMA_BLOCK_K / MMA_K; k++) {
+           uint64_t a_desc = make_desc_AB(stage_base + k * 32);
+           uint64_t b_desc = make_desc_AB(stage_base + B_off + k * 32);
   
-             int k_sf = k1 * 4 + k2;
-             const int scale_A_tmem = SFA_tmem + k_sf * 4 + (work.tile_m % (MMA_M / BLOCK_M)) * (BLOCK_M / 32);
-             int scale_B_tmem;
-             if constexpr (BLOCK_N == 256) {
-               scale_B_tmem = SFB_tmem + k_sf * 8;
-             } else {
-               scale_B_tmem = SFB_tmem + k_sf * 4;
-             }
-  
-             const int enable_input_d = (k_iter == 0 && k1 == 0 && k2 == 0) ? 0 : 1;
-             tcgen05_mma_nvfp4(a_desc, b_desc, idesc, scale_A_tmem, scale_B_tmem, enable_input_d);
+           const int scale_A_tmem = SFA_tmem + k * 4 + (work.tile_m % (MMA_M / BLOCK_M)) * (BLOCK_M / 32);
+           int scale_B_tmem;
+           if constexpr (BLOCK_N == 256) {
+             scale_B_tmem = SFB_tmem + k * 8;
+           } else {
+             scale_B_tmem = SFB_tmem + k * 4;
            }
+  
+           // First MMA uses D=0, subsequent MMAs accumulate.
+           const int enable_input_d = (k_iter == 0 && k == 0) ? 0 : 1;
+           tcgen05_mma_nvfp4(a_desc, b_desc, idesc, scale_A_tmem, scale_B_tmem, enable_input_d);
          }
   
-         tcgen05_commit(mbar_base + (NUM_STAGES + stage) * 8);
+         tcgen05_commit(mbar_base + (NS + stage) * 8);
+         if constexpr (CLUSTER_SIZE > 1) {
+           int cm = mapa_cta_to_cluster(mbar_base + (2*NS + stage) * 8, 0);
+           mbarrier_arrive_cluster(cm);
+         }
        }
   
-       // Wait for final commit
-       if (num_k_iters > 0) {
-         const int last_stage = (num_k_iters - 1) % NUM_STAGES;
-         const int last_phase = ((num_k_iters - 1) / NUM_STAGES) % 2;
-         mbarrier_wait(mbar_base + (NUM_STAGES + last_stage) * 8, last_phase);
-       }
+       const int last_stage = (num_k_iters - 1) % NS;
+       const int last_phase = ((num_k_iters - 1) / NS) % 2;
+       mbarrier_wait(mbar_base + (NS + last_stage) * 8, last_phase);
      }
   
-     // ================================================================
-     // SYNCHRONIZATION & EPILOGUE
-     // ================================================================
      __syncthreads();
      asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
   
  ⋯ 4 unchanged lines
        epilogue_store<BLOCK_M, BLOCK_N, false>(prob, m_offset, n_offset, tid, warp_id, lane_id);
      }
   
-     __syncthreads();
+     // In cluster mode we must synchronize across CTAs before re-initializing
+     // mbarriers; __syncthreads is CTA-local and does not order cluster-wide
+     // mbarrier arrivals.
+     if constexpr (PERSISTENT && CLUSTER_SIZE > 1) {
+       cluster_sync();
+     }
   
-     // Get next work item
      if constexpr (PERSISTENT) {
-       if (tid == 0) {
-         shared_work_idx = atomicAdd(work_counter, 1);
+       if (warp_id == TMA_WARP && elect_sync()) {
+         // Static grid-stride work distribution avoids global atomics and
+         // smooths the tail when num_items slightly exceeds one wave.
+         shared_work_idx = work_idx + gridDim.x;
+  
+         if (shared_work_idx < num_items) {
+           for (int i = 0; i < NS; ++i) {
+             mbarrier_init(mbar_base + i * 8, 2);
+             mbarrier_init(mbar_base + (NS + i) * 8, 1);
+             if constexpr (CLUSTER_SIZE > 1) {
+               if (cta_rank == 0) {
+                 mbarrier_init(mbar_base + (2*NS + i) * 8, CLUSTER_SIZE);
+               }
+             }
+           }
+           asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+         }
        }
-       __syncthreads();
+     }
+  
+     __syncthreads();
+  
+     // Cluster barrier after re-init: ensure all CTAs see re-initialized mbarriers.
+     if constexpr (PERSISTENT && CLUSTER_SIZE > 1) {
+       cluster_sync();
+     }
+  
+     if constexpr (PERSISTENT) {
        work_idx = shared_work_idx;
      } else {
        break;
      }
    }
   
-   // Deallocate TMEM
    if (warp_id == 0) {
      tcgen05_dealloc_cols_cta1(0, TMEM_COLS);
    }
  }
   
- // ============================================================================
- // TENSOR MAP INITIALIZATION
- // ============================================================================
-  
  void init_AB_tmap_u4(
    CUtensorMap *tmap,
    const void *ptr,
  ⋯ 3 unchanged lines
    TORCH_CHECK(ptr != nullptr, "ptr is null");
    TORCH_CHECK(((uintptr_t)ptr % 16) == 0, "ptr must be 16-byte aligned");
    TORCH_CHECK(global_width >= 256 && (global_width % 256) == 0, "K must be multiple of 256");
+   TORCH_CHECK(shared_width == 256, "shared_width must be 256");
   
    constexpr uint32_t rank = 3;
    uint64_t globalDim[rank]       = {256, global_height, global_width / 256};
    uint64_t globalStrides[rank-1] = {global_width / 2, 128};
-   uint32_t boxDim[rank]          = {256, shared_height, shared_width / 256};
+   uint32_t boxDim[rank]          = {256, shared_height, 1};
    uint32_t elementStrides[rank]  = {1, 1, 1};
   
-   // cuTensorMapEncodeTiled is a relatively expensive driver call.
-   // Cache a per-shape template and then patch only the base address.
+   // Swizzle must match the shared-memory layout expected by tcgen05.
+   constexpr CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_128B;
+  
+   // Cache cuTensorMap templates by shape.
    struct ShapeKey { uint64_t gh, gw; uint32_t sh, sw; };
    struct ShapeHash {
      size_t operator()(const ShapeKey& k) const noexcept {
  ⋯ 38 unchanged lines
      auto err = cuTensorMapEncodeTiled(
        &tmp, CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
        rank, (void*)ptr, globalDim, globalStrides, boxDim, elementStrides,
-       CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+       CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle,
        CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
      );
      TORCH_CHECK(err == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed");
  ⋯ 7 unchanged lines
    *tmap = tmp;
  }
   
- struct PadCacheEntry {
-   at::Tensor buf;
-   size_t zeroed_from = 0; // byte offset; bytes in [zeroed_from, end) are guaranteed zero
- };
-  
- static at::Tensor pad_u4_tensor_cached(const at::Tensor& src, int64_t padded_m, PadCacheEntry* entry) {
-   TORCH_CHECK(entry != nullptr, "pad_u4_tensor_cached: entry is null");
-   if (src.size(0) == padded_m && ((uintptr_t)src.data_ptr() & 0xF) == 0) return src;
-  
-   auto new_sizes = src.sizes().vec();
-   new_sizes[0] = padded_m;
-  
-   bool reuse_ok = entry->buf.defined() && entry->buf.dim() == (int)new_sizes.size();
-   if (reuse_ok) {
-     // Allow reusing a larger leading-dimension buffer (avoid reallocs when M/N shrink),
-     // but require trailing dimensions to match exactly.
-     if (entry->buf.size(0) < padded_m) reuse_ok = false;
-     for (int d = 1; d < entry->buf.dim(); d++) {
-       if (entry->buf.size(d) != new_sizes[(size_t)d]) { reuse_ok = false; break; }
-     }
-     if (reuse_ok && (((uintptr_t)entry->buf.data_ptr() & 0xF) != 0)) reuse_ok = false;
-   }
-  
-   const bool need_new =
-     !entry->buf.defined() ||
-     entry->buf.device() != src.device() ||
-     entry->buf.scalar_type() != src.scalar_type() ||
-     !reuse_ok;
-  
-   if (need_new) {
-     entry->buf = at::empty(new_sizes, src.options());
-     entry->zeroed_from = (size_t)entry->buf.nbytes(); // nothing guaranteed yet
-   }
-  
-   const size_t copy_bytes = (size_t)src.nbytes();
-   const size_t total_bytes = (size_t)entry->buf.nbytes();
-   TORCH_CHECK(copy_bytes <= total_bytes, "pad_u4_tensor_cached: size mismatch");
-  
-   CUDA_CHECK(cudaMemcpyAsync(entry->buf.data_ptr(), src.data_ptr(), copy_bytes, cudaMemcpyDeviceToDevice));
-   if (copy_bytes < entry->zeroed_from) {
-     CUDA_CHECK(cudaMemsetAsync((char*)entry->buf.data_ptr() + copy_bytes, 0, entry->zeroed_from - copy_bytes));
-     entry->zeroed_from = copy_bytes;
-   } else {
-     entry->zeroed_from = copy_bytes;
-   }
-   return entry->buf;
- }
-  
- // ============================================================================
- // HOST ENTRY POINT
- // ============================================================================
-  
  std::vector<at::Tensor> group_gemm(
      std::vector<at::Tensor> A_list,
      std::vector<at::Tensor> B_list,
  ⋯ 7 unchanged lines
      c10::cuda::CUDAGuard device_guard(dev);
      auto sizes_accessor = sizes_cpu.accessor<int64_t, 2>();
   
+     // SM count is used for occupancy-based launch shaping.
+     // Hardcoded for the target environment (B200 / SM100) to keep the logic
+     // simple and deterministic.
+     constexpr int sm_count = 148;
+  
      static bool attrs_set = false;
      if (!attrs_set) {
-       constexpr int STAGE_SIZE_128 = 128 * 128 + 128 * 128 + 128 * 16 + 128 * 16;
-       constexpr int SMEM_SIZE_128 = STAGE_SIZE_128 * NUM_STAGES + MBAR_BYTES;
-       constexpr int STAGE_SIZE_64 = 64 * 128 + 128 * 128 + 128 * 16 + 128 * 16;
-       constexpr int SMEM_SIZE_64 = STAGE_SIZE_64 * NUM_STAGES + MBAR_BYTES;
-       constexpr int STAGE_SIZE_256 = 128 * 128 + 256 * 128 + 128 * 16 + 256 * 16;
-       constexpr int SMEM_SIZE_256 = STAGE_SIZE_256 * NUM_STAGES + MBAR_BYTES;
-       CUDA_CHECK(cudaFuncSetAttribute(grouped_gemm_kernel_v4<true, 128, 128>, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_128));
-       CUDA_CHECK(cudaFuncSetAttribute(grouped_gemm_kernel_v4<false, 128, 128>, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_128));
-       CUDA_CHECK(cudaFuncSetAttribute(grouped_gemm_kernel_v4<true, 64, 128>, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_64));
-       CUDA_CHECK(cudaFuncSetAttribute(grouped_gemm_kernel_v4<false, 64, 128>, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_64));
-       CUDA_CHECK(cudaFuncSetAttribute(grouped_gemm_kernel_v4<true, 128, 256>, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_256));
-       CUDA_CHECK(cudaFuncSetAttribute(grouped_gemm_kernel_v4<false, 128, 256>, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_256));
+       // Two pipeline depths:
+       // - HI: deeper pipeline, more overlap, higher shared-memory footprint.
+       // - LO: shallower pipeline, lower shared-memory footprint (can improve
+       //   occupancy when shared memory is the limiter).
+       constexpr int NS_DEEP_HI = 6;
+       constexpr int NS_DEEP_LO = 3;
+       constexpr int NS_WIDE_HI = 4;
+  
+       constexpr int MBAR_BYTES_DEEP_HI = ((2 * NS_DEEP_HI * 8 + 63) & ~63);
+       constexpr int MBAR_BYTES_DEEP_LO = ((2 * NS_DEEP_LO * 8 + 63) & ~63);
+       constexpr int MBAR_BYTES_WIDE_HI = ((3 * NS_WIDE_HI * 8 + 63) & ~63);
+  
+       // BLOCK_N=128
+       constexpr int STAGE_128_K256 = 128 * (256 / 2) + 128 * (256 / 2) + 128 * (256 / 16) + 128 * (256 / 16);
+       constexpr int SMEM_128_HI_K256 = STAGE_128_K256 * NS_DEEP_HI + MBAR_BYTES_DEEP_HI;
+       constexpr int SMEM_128_LO_K256 = STAGE_128_K256 * NS_DEEP_LO + MBAR_BYTES_DEEP_LO;
+  
+       // BLOCK_N=256
+       constexpr int STAGE_256_K256 = 128 * (256 / 2) + 256 * (256 / 2) + 128 * (256 / 16) + 256 * (256 / 16);
+       constexpr int SMEM_256_HI_K256 = STAGE_256_K256 * NS_WIDE_HI + MBAR_BYTES_WIDE_HI;
+  
+       // Variants: {persistent} x {BLOCK_N} x {NS}
+       // BLOCK_N=128
+       CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<true, 128, 128, NS_DEEP_HI>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_128_HI_K256));
+       CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<false, 128, 128, NS_DEEP_HI>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_128_HI_K256));
+       CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<true, 128, 128, NS_DEEP_LO>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_128_LO_K256));
+       CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<false, 128, 128, NS_DEEP_LO>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_128_LO_K256));
+  
+       // BLOCK_N=256 with cluster multicast (CLUSTER_SIZE=4)
+       constexpr int CL = 4;
+       CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<true, 128, 256, NS_WIDE_HI, CL>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_256_HI_K256));
+       CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<false, 128, 256, NS_WIDE_HI, CL>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_256_HI_K256));
+       CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<true, 128, 256, NS_WIDE_HI, CL>), cudaFuncAttributeNonPortableClusterSizeAllowed, 1));
+       CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<false, 128, 256, NS_WIDE_HI, CL>), cudaFuncAttributeNonPortableClusterSizeAllowed, 1));
+  
        attrs_set = true;
      }
   
+     // Cache occupancy for launch shaping (per process).
+     static thread_local bool occ_set = false;
+     static thread_local int occ_128_hi = 0, occ_128_lo = 0;
+     static thread_local int occ_256_hi = 0;
+     if (!occ_set) {
+       constexpr int NS_DEEP_HI = 6;
+       constexpr int NS_DEEP_LO = 3;
+       constexpr int NS_WIDE_HI = 4;
+       constexpr int MBAR_BYTES_DEEP_HI = ((2 * NS_DEEP_HI * 8 + 63) & ~63);
+       constexpr int MBAR_BYTES_DEEP_LO = ((2 * NS_DEEP_LO * 8 + 63) & ~63);
+       constexpr int MBAR_BYTES_WIDE_HI = ((3 * NS_WIDE_HI * 8 + 63) & ~63);
+       constexpr int STAGE_128_K256 = 128 * (256 / 2) + 128 * (256 / 2) + 128 * (256 / 16) + 128 * (256 / 16);
+       constexpr int STAGE_256_K256 = 128 * (256 / 2) + 256 * (256 / 2) + 128 * (256 / 16) + 256 * (256 / 16);
+       constexpr int SMEM_128_HI_K256 = STAGE_128_K256 * NS_DEEP_HI + MBAR_BYTES_DEEP_HI;
+       constexpr int SMEM_128_LO_K256 = STAGE_128_K256 * NS_DEEP_LO + MBAR_BYTES_DEEP_LO;
+       constexpr int SMEM_256_HI_K256 = STAGE_256_K256 * NS_WIDE_HI + MBAR_BYTES_WIDE_HI;
+  
+       constexpr int THREADS = TMA_NUM_WARPS * WARP_SIZE;
+       CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+           &occ_128_hi, grouped_gemm_kernel_v4<false, 128, 128, NS_DEEP_HI>, THREADS, SMEM_128_HI_K256));
+       CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+           &occ_128_lo, grouped_gemm_kernel_v4<false, 128, 128, NS_DEEP_LO>, THREADS, SMEM_128_LO_K256));
+       CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+           &occ_256_hi, grouped_gemm_kernel_v4<false, 128, 256, NS_WIDE_HI, 4>, THREADS, SMEM_256_HI_K256));
+  
+       TORCH_CHECK(occ_128_hi > 0 && occ_128_lo > 0 && occ_256_hi > 0,
+                   "occupancy query returned zero blocks/SM");
+       occ_set = true;
+     }
+  
      std::vector<ProblemInfo> problem_infos(G);
      static thread_local std::vector<WorkItem> cached_work_items_128;
-     static thread_local std::vector<WorkItem> cached_work_items_64;
      static thread_local std::vector<WorkItem> cached_work_items_256;
      static thread_local uint64_t cached_work_hash = 0;
      static thread_local bool cached_work_valid = false;
   
      std::vector<uint8_t> active(G, 0);
-     std::vector<uint8_t> use_64(G, 0);
      std::vector<uint8_t> use_256(G, 0);
      std::vector<int> num_tiles_m(G, 0);
      std::vector<int> num_tiles_n(G, 0);
      std::vector<int64_t> Ms(G, 0), Ns(G, 0), Ks(G, 0);
   
-     uint64_t work_hash = 1469598103934665603ULL;
+     int64_t total_tiles = 0;
      for (int64_t i = 0; i < G; i++) {
          const int64_t M = sizes_accessor[i][0], N = sizes_accessor[i][1], K = sizes_accessor[i][2];
          Ms[(size_t)i] = M; Ns[(size_t)i] = N; Ks[(size_t)i] = K;
          if (A_list[i].stride(1) != 1 || B_list[i].stride(1) != 1) {
-           work_hash = hash_combine_u64(work_hash, 0);
            continue;
          }
          active[(size_t)i] = 1;
-         bool is_64 = (M <= 64) && (N <= 2048);
-         bool is_256 = (!is_64) && (N >= 4096) && (K >= 2048) && ((N & 255) == 0);
-         use_64[(size_t)i] = is_64;
+         bool is_256 = (N >= 4096) && (K >= 2048) && ((N & 255) == 0);
          use_256[(size_t)i] = is_256;
-         int block_m = is_64 ? 64 : 128;
+         int block_m = 128;
          int block_n = is_256 ? 256 : 128;
          num_tiles_m[(size_t)i] = ceil_div((int)M, block_m);
          num_tiles_n[(size_t)i] = ceil_div((int)N, block_n);
+  
+         total_tiles += (int64_t)num_tiles_m[(size_t)i] * (int64_t)num_tiles_n[(size_t)i];
+     }
+  
+     const int tma_block_k = 256;
+  
+     uint64_t work_hash = 1469598103934665603ULL;
+     work_hash = hash_combine_u64(work_hash, (uint64_t)tma_block_k);
+     for (int64_t i = 0; i < G; i++) {
+         const bool is_active = (active[(size_t)i] != 0);
+         if (!is_active) {
+           work_hash = hash_combine_u64(work_hash, 0);
+           continue;
+         }
+         const int64_t M = Ms[(size_t)i];
+         const int64_t N = Ns[(size_t)i];
+         const bool is_256 = (use_256[(size_t)i] != 0);
          
          work_hash = hash_combine_u64(work_hash, (uint64_t)M);
          work_hash = hash_combine_u64(work_hash, (uint64_t)N);
-         work_hash = hash_combine_u64(work_hash, (uint64_t)is_64);
          work_hash = hash_combine_u64(work_hash, (uint64_t)is_256);
      }
   
      if (!cached_work_valid || cached_work_hash != work_hash) {
        cached_work_items_128.clear();
-       cached_work_items_64.clear();
        cached_work_items_256.clear();
        cached_work_items_128.reserve(G * 32);
-       cached_work_items_64.reserve(G * 32);
        cached_work_items_256.reserve(G * 32);
+  
+       // Work scheduling.
+       // 128-wide path: iterate by tile_n across groups so a wave tends to touch
+       // the same B strip across different problems (better L2 locality).
+       int max_tn_128 = 0, max_tn_256 = 0;
        for (int64_t i = 0; i < G; i++) {
          if (!active[(size_t)i]) continue;
-         // M-tile-major ordering: consecutive work items share the same M-tile
-         // so A data stays hot in L2 while iterating N-tiles
+         if (use_256[(size_t)i]) {
+           max_tn_256 = std::max(max_tn_256, num_tiles_n[(size_t)i]);
+         } else {
+           max_tn_128 = std::max(max_tn_128, num_tiles_n[(size_t)i]);
+         }
+       }
+  
+       // 128-wide tiles: interleave by tile_n across groups
+       for (int tn = 0; tn < max_tn_128; tn++) {
+         for (int64_t i = 0; i < G; i++) {
+           if (!active[(size_t)i] || use_256[(size_t)i]) continue;
+           if (tn >= num_tiles_n[(size_t)i]) continue;
+           for (int tm = 0; tm < num_tiles_m[(size_t)i]; tm++) {
+             cached_work_items_128.push_back({(int)i, tm, tn});
+           }
+         }
+       }
+  
+       // 256-wide path: order by (problem, tile_m, tile_n) with tile_n innermost
+       // so each cluster of consecutive CTAs shares the same A tile (multicast).
+       constexpr int CLUSTER_SIZE_256 = 4;
+       for (int64_t i = 0; i < G; i++) {
+         if (!active[(size_t)i] || !use_256[(size_t)i]) continue;
          for (int tm = 0; tm < num_tiles_m[(size_t)i]; tm++) {
            for (int tn = 0; tn < num_tiles_n[(size_t)i]; tn++) {
-             if (use_64[(size_t)i]) {
-               cached_work_items_64.push_back({(int)i, tm, tn});
-             } else if (use_256[(size_t)i]) {
-               cached_work_items_256.push_back({(int)i, tm, tn});
-             } else {
-               cached_work_items_128.push_back({(int)i, tm, tn});
+             cached_work_items_256.push_back({(int)i, tm, tn});
+           }
+           // Pad to a multiple of CLUSTER_SIZE so every launched cluster is full.
+           int remainder = num_tiles_n[(size_t)i] % CLUSTER_SIZE_256;
+           if (remainder != 0) {
+             for (int p = 0; p < CLUSTER_SIZE_256 - remainder; p++) {
+               cached_work_items_256.push_back({(int)i, tm, 0});
              }
            }
          }
        }
+  
        cached_work_hash = work_hash;
        cached_work_valid = true;
      }
   
-     static thread_local std::vector<PadCacheEntry> A_pad_cache;
-     static thread_local std::vector<PadCacheEntry> B_pad_cache;
-     if ((int64_t)A_pad_cache.size() < G) A_pad_cache.resize((size_t)G);
-     if ((int64_t)B_pad_cache.size() < G) B_pad_cache.resize((size_t)G);
-  
      uint64_t probs_hash = 1469598103934665603ULL;
+     probs_hash = hash_combine_u64(probs_hash, (uint64_t)tma_block_k);
      for (int64_t i = 0; i < G; i++) {
          if (!active[(size_t)i]) continue;
          const int64_t M = Ms[(size_t)i], N = Ns[(size_t)i], K = Ks[(size_t)i];
   
-         int block_m = use_64[(size_t)i] ? 64 : 128;
-         int block_n = use_256[(size_t)i] ? 256 : 128;
-         int64_t padded_M = ((M + block_m - 1) / block_m) * block_m;
-         int64_t padded_N = ((N + block_n - 1) / block_n) * block_n;
-  
-         at::Tensor A = pad_u4_tensor_cached(A_list[i], padded_M, &A_pad_cache[(size_t)i]);
-         at::Tensor B = pad_u4_tensor_cached(B_list[i], padded_N, &B_pad_cache[(size_t)i]);
-  
          ProblemInfo& p = problem_infos[i];
          p.M = M; p.N = N; p.K = K;
-         p.Cs0 = C_list[i].stride(0); p.Cs1 = C_list[i].stride(1); p.Cs2 = C_list[i].stride(2);
+         p.Cs0 = C_list[i].stride(0); p.Cs1 = C_list[i].stride(1);
          p.C_ptr = (half*)C_list[i].data_ptr();
          p.SFA_ptr = (const char*)sfa_list[i].data_ptr();
          p.SFB_ptr = (const char*)sfb_list[i].data_ptr();
   
-         init_AB_tmap_u4(&p.A_tmap, A.data_ptr(), A.size(0), K, block_m, TMA_BLOCK_K);
-         init_AB_tmap_u4(&p.B_tmap, B.data_ptr(), B.size(0), K, 128, TMA_BLOCK_K);
+         init_AB_tmap_u4(&p.A_tmap, A_list[i].data_ptr(), A_list[i].size(0), K, 128, 256);
+         init_AB_tmap_u4(&p.B_tmap, B_list[i].data_ptr(), B_list[i].size(0), K, 128, 256);
          if (use_256[(size_t)i]) {
-           init_AB_tmap_u4(&p.B_tmap_256, B.data_ptr(), B.size(0), K, 256, TMA_BLOCK_K);
+           init_AB_tmap_u4(&p.B_tmap_256, B_list[i].data_ptr(), B_list[i].size(0), K, 256, 256);
          } else {
            p.B_tmap_256 = p.B_tmap;
          }
  ⋯ 2 unchanged lines
          probs_hash = hash_combine_u64(probs_hash, (uint64_t)M);
          probs_hash = hash_combine_u64(probs_hash, (uint64_t)N);
          probs_hash = hash_combine_u64(probs_hash, (uint64_t)K);
-         probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)A.data_ptr());
-         probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)B.data_ptr());
+         probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)A_list[i].data_ptr());
+         probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)B_list[i].data_ptr());
          probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)p.C_ptr);
          probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)p.SFA_ptr);
          probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)p.SFB_ptr);
          probs_hash = hash_combine_u64(probs_hash, (uint64_t)p.Cs0);
          probs_hash = hash_combine_u64(probs_hash, (uint64_t)p.Cs1);
-         probs_hash = hash_combine_u64(probs_hash, (uint64_t)p.Cs2);
      }
   
-     if (cached_work_items_128.empty() && cached_work_items_64.empty() && cached_work_items_256.empty()) return C_list;
+     if (cached_work_items_128.empty() && cached_work_items_256.empty()) return C_list;
   
      auto options = at::TensorOptions().dtype(at::kByte).device(dev);
      static thread_local at::Tensor d_probs_cache;
      static thread_local at::Tensor d_work_cache_128;
-     static thread_local at::Tensor d_work_cache_64;
      static thread_local at::Tensor d_work_cache_256;
-     static thread_local at::Tensor d_counter_cache;
      static thread_local uint64_t last_probs_hash = 0;
      static thread_local uint64_t last_work_hash = 0;
      static thread_local bool last_hash_valid = false;
   
      const int64_t probs_bytes = (int64_t)(G * sizeof(ProblemInfo));
      const int64_t work_bytes_128 = (int64_t)(cached_work_items_128.size() * sizeof(WorkItem));
-     const int64_t work_bytes_64 = (int64_t)(cached_work_items_64.size() * sizeof(WorkItem));
      const int64_t work_bytes_256 = (int64_t)(cached_work_items_256.size() * sizeof(WorkItem));
   
      bool probs_realloc = false;
  ⋯ 4 unchanged lines
      if (work_bytes_128 > 0 && (!d_work_cache_128.defined() || d_work_cache_128.device() != dev || d_work_cache_128.scalar_type() != at::kByte || d_work_cache_128.numel() < work_bytes_128)) {
        d_work_cache_128 = at::empty({work_bytes_128}, options);
      }
-     if (work_bytes_64 > 0 && (!d_work_cache_64.defined() || d_work_cache_64.device() != dev || d_work_cache_64.scalar_type() != at::kByte || d_work_cache_64.numel() < work_bytes_64)) {
-       d_work_cache_64 = at::empty({work_bytes_64}, options);
-     }
      if (work_bytes_256 > 0 && (!d_work_cache_256.defined() || d_work_cache_256.device() != dev || d_work_cache_256.scalar_type() != at::kByte || d_work_cache_256.numel() < work_bytes_256)) {
        d_work_cache_256 = at::empty({work_bytes_256}, options);
      }
  ⋯ 6 unchanged lines
        if (work_bytes_128 > 0) {
          CUDA_CHECK(cudaMemcpyAsync(d_work_cache_128.data_ptr(), cached_work_items_128.data(), work_bytes_128, cudaMemcpyHostToDevice));
        }
-       if (work_bytes_64 > 0) {
-         CUDA_CHECK(cudaMemcpyAsync(d_work_cache_64.data_ptr(), cached_work_items_64.data(), work_bytes_64, cudaMemcpyHostToDevice));
-       }
        if (work_bytes_256 > 0) {
          CUDA_CHECK(cudaMemcpyAsync(d_work_cache_256.data_ptr(), cached_work_items_256.data(), work_bytes_256, cudaMemcpyHostToDevice));
        }
  ⋯ 1 unchanged lines
      }
      last_hash_valid = true;
   
-     constexpr int MAX_CTAS = 264;
-     if (!d_counter_cache.defined() || d_counter_cache.device() != dev || d_counter_cache.scalar_type() != at::kInt || d_counter_cache.numel() != 1) {
-       d_counter_cache = at::empty({1}, options.dtype(at::kInt));
-     }
+     constexpr int NS_DEEP_HI = 6;
+     constexpr int NS_DEEP_LO = 3;
+     constexpr int NS_WIDE_HI = 4;
   
      if (!cached_work_items_128.empty()) {
        int num_items_128 = (int)cached_work_items_128.size();
-       constexpr int STAGE_SIZE_128 = 128 * 128 + 128 * 128 + 128 * 16 + 128 * 16;
-       constexpr int SMEM_SIZE_128 = STAGE_SIZE_128 * NUM_STAGES + MBAR_BYTES;
-       if (num_items_128 > MAX_CTAS) {
-         CUDA_CHECK(cudaMemsetAsync(d_counter_cache.data_ptr(), 0, sizeof(int)));
-         grouped_gemm_kernel_v4<true, 128, 128><<<MAX_CTAS, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_128>>>(
-             (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_128.data_ptr(), num_items_128, (int*)d_counter_cache.data_ptr());
-       } else {
-         grouped_gemm_kernel_v4<false, 128, 128><<<num_items_128, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_128>>>(
-             (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_128.data_ptr(), num_items_128, nullptr);
-       }
-     }
   
-     if (!cached_work_items_64.empty()) {
-       int num_items_64 = (int)cached_work_items_64.size();
-       constexpr int STAGE_SIZE_64 = 64 * 128 + 128 * 128 + 128 * 16 + 128 * 16;
-       constexpr int SMEM_SIZE_64 = STAGE_SIZE_64 * NUM_STAGES + MBAR_BYTES;
-       if (num_items_64 > MAX_CTAS) {
-         CUDA_CHECK(cudaMemsetAsync(d_counter_cache.data_ptr(), 0, sizeof(int)));
-         grouped_gemm_kernel_v4<true, 64, 128><<<MAX_CTAS, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_64>>>(
-             (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_64.data_ptr(), num_items_64, (int*)d_counter_cache.data_ptr());
-       } else {
-         grouped_gemm_kernel_v4<false, 64, 128><<<num_items_64, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_64>>>(
-             (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_64.data_ptr(), num_items_64, nullptr);
+       // Choose pipeline depth by expected waves.
+       const int wave_hi = sm_count * occ_128_hi;
+       const int wave_lo = sm_count * occ_128_lo;
+       // Only switch to low-smem variant when the grid is large enough that reducing
+       // waves is likely to outweigh reduced pipeline overlap.
+       const bool use_lo = (tma_block_k == 256) && (wave_lo > wave_hi) && (num_items_128 > 2 * wave_hi);
+  
+       const int ns = use_lo ? NS_DEEP_LO : NS_DEEP_HI;
+       const int occ = use_lo ? occ_128_lo : occ_128_hi;
+       const int wave_cap = sm_count * occ;
+  
+       constexpr int MBAR_HI = ((2 * NS_DEEP_HI * 8 + 63) & ~63);
+       constexpr int MBAR_LO = ((2 * NS_DEEP_LO * 8 + 63) & ~63);
+       const int stage_size_128 = 128 * (tma_block_k / 2) + 128 * (tma_block_k / 2) + 128 * (tma_block_k / 16) + 128 * (tma_block_k / 16);
+       const int SMEM_SIZE_128 = stage_size_128 * ns + (use_lo ? MBAR_LO : MBAR_HI);
+  
+       // Launch shaping: if more CTAs than one full wave, use persistent grid-stride.
+       const bool persistent_128 = (num_items_128 > wave_cap);
+       const int launch_ctas_128 = persistent_128 ? wave_cap : num_items_128;
+  
+       if (tma_block_k == 256) {
+         if (persistent_128) {
+           if (use_lo) {
+             grouped_gemm_kernel_v4<true, 128, 128, NS_DEEP_LO><<<launch_ctas_128, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_128>>>(
+                 (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_128.data_ptr(), num_items_128);
+           } else {
+             grouped_gemm_kernel_v4<true, 128, 128, NS_DEEP_HI><<<launch_ctas_128, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_128>>>(
+                 (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_128.data_ptr(), num_items_128);
+           }
+         } else {
+           if (use_lo) {
+             grouped_gemm_kernel_v4<false, 128, 128, NS_DEEP_LO><<<launch_ctas_128, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_128>>>(
+                 (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_128.data_ptr(), num_items_128);
+           } else {
+             grouped_gemm_kernel_v4<false, 128, 128, NS_DEEP_HI><<<launch_ctas_128, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_128>>>(
+                 (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_128.data_ptr(), num_items_128);
+           }
+         }
        }
      }
   
      if (!cached_work_items_256.empty()) {
        int num_items_256 = (int)cached_work_items_256.size();
-       constexpr int STAGE_SIZE_256 = 128 * 128 + 256 * 128 + 128 * 16 + 256 * 16;
-       constexpr int SMEM_SIZE_256 = STAGE_SIZE_256 * NUM_STAGES + MBAR_BYTES;
-       if (num_items_256 > MAX_CTAS) {
-         CUDA_CHECK(cudaMemsetAsync(d_counter_cache.data_ptr(), 0, sizeof(int)));
-         grouped_gemm_kernel_v4<true, 128, 256><<<MAX_CTAS, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_256>>>(
-             (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_256.data_ptr(), num_items_256, (int*)d_counter_cache.data_ptr());
-       } else {
-         grouped_gemm_kernel_v4<false, 128, 256><<<num_items_256, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_256>>>(
-             (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_256.data_ptr(), num_items_256, nullptr);
+       constexpr int CL256 = 4;
+  
+       const int wave_cap = (sm_count * occ_256_hi / CL256) * CL256;
+  
+       constexpr int MBAR_HI = ((3 * NS_WIDE_HI * 8 + 63) & ~63);
+       const int stage_size_256 = 128 * (tma_block_k / 2) + 256 * (tma_block_k / 2) + 128 * (tma_block_k / 16) + 256 * (tma_block_k / 16);
+       const int SMEM_SIZE_256 = stage_size_256 * NS_WIDE_HI + MBAR_HI;
+  
+       const bool persistent_256 = (num_items_256 > wave_cap);
+       int launch_ctas_256 = persistent_256 ? wave_cap : num_items_256;
+  
+       if (tma_block_k == 256) {
+         const ProblemInfo* d_probs_ptr = (const ProblemInfo*)d_probs_cache.data_ptr();
+         const WorkItem* d_work_ptr = (const WorkItem*)d_work_cache_256.data_ptr();
+  
+         cudaLaunchConfig_t config = {};
+         config.gridDim = dim3(launch_ctas_256);
+         config.blockDim = dim3(TMA_NUM_WARPS * WARP_SIZE);
+         config.dynamicSmemBytes = SMEM_SIZE_256;
+  
+         cudaLaunchAttribute launch_attrs[1];
+         launch_attrs[0].id = cudaLaunchAttributeClusterDimension;
+         launch_attrs[0].val.clusterDim = {CL256, 1, 1};
+         config.attrs = launch_attrs;
+         config.numAttrs = 1;
+  
+         if (persistent_256) {
+           CUDA_CHECK(cudaLaunchKernelEx(&config, grouped_gemm_kernel_v4<true, 128, 256, NS_WIDE_HI, CL256>,
+               d_probs_ptr, d_work_ptr, num_items_256));
+         } else {
+           CUDA_CHECK(cudaLaunchKernelEx(&config, grouped_gemm_kernel_v4<false, 128, 256, NS_WIDE_HI, CL256>,
+               d_probs_ptr, d_work_ptr, num_items_256));
+         }
        }
      }
+  
      CUDA_CHECK(cudaGetLastError());
      return C_list;
  }
  ⋯ 25 unchanged lines
  group_gemm = torch.ops.my_module.group_gemm
   
   
+ @lru_cache(maxsize=128)
+ def _sizes_cpu_cached(key: tuple[tuple[int, int, int], ...]) -> torch.Tensor:
+     return torch.tensor(key, dtype=torch.int64, device="cpu")
+  
+  
  def custom_kernel(data: input_t) -> output_t:
-     abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
+     abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data
      A_list = [t[0] for t in abc_tensors]
      B_list = [t[1] for t in abc_tensors]
  ⋯ diff truncated