submission 293471

  ⋯ 4 unchanged lines
  from task import input_t, output_t
  from torch.utils.cpp_extension import load_inline
   
- # 2-SM cluster variant of master_kernel
  cuda_src = """
  #include <cudaTypedefs.h>
  #include <cuda_fp16.h>
  #include <cuda_fp8.h>
  #include <cuda_runtime.h>
- #include <stdio.h>
-  
  #include <torch/library.h>
  #include <ATen/core/Tensor.h>
   
  ⋯ 1 unchanged lines
  constexpr int MMA_K = 64;
  constexpr int TMEM_COLS = 512;
   
-  
  constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;
  constexpr uint64_t EVICT_LAST = 0x14F0000000000000;
  constexpr uint32_t SM100_MMA_PEER_MASK = 0xFEFFFFFF;
   
- __device__ inline
- constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };
+ __device__ inline constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; }
   
- __device__ inline
- uint32_t elect_sync() {
+ __device__ inline uint32_t elect_sync() {
      uint32_t pred = 0;
      asm volatile(
          "{\\n\\t"
  ⋯ 1 unchanged lines
          "elect.sync _|%%px, %1;\\n\\t"
          "@%%px mov.s32 %0, 1;\\n\\t"
          "}"
-         : "+r"(pred)
-         : "r"(0xFFFFFFFF)
-     );
+         : "+r"(pred) : "r"(0xFFFFFFFF));
      return pred;
  }
   
- __device__ inline
- uint32_t get_cluster_ctarank() {
+ __device__ inline uint32_t get_cluster_ctarank() {
      uint32_t rank;
      asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank));
      return rank;
  }
   
- __device__ inline
- void barrier_cluster_arrive() {
-     asm volatile("barrier.cluster.arrive.aligned;");
- }
-  
- __device__ inline
- void barrier_cluster_wait() {
-     asm volatile("barrier.cluster.wait.aligned;");
- }
-  
- __device__ inline
- void mbarrier_init(int mbar_addr, int count) {
+ __device__ inline void mbarrier_init(int mbar_addr, int count) {
      asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
  }
   
- __device__
- void mbarrier_wait(int mbar_addr, int phase) {
+ __device__ void mbarrier_wait(int mbar_addr, int phase) {
      uint32_t ticks = 0x989680;
      asm volatile(
          "{\\n\\t"
  ⋯ 4 unchanged lines
          "bra.uni LAB_WAIT;\\n\\t"
          "DONE:\\n\\t"
          "}"
-         :: "r"(mbar_addr), "r"(phase), "r"(ticks)
-     );
+         :: "r"(mbar_addr), "r"(phase), "r"(ticks));
  }
   
- __device__ inline
- void mbarrier_arrive(int mbar_addr) {
-     asm volatile("mbarrier.arrive.release.cluster.shared::cta.b64 _, [%0];" :: "r"(mbar_addr) : "memory");
- }
-  
- __device__ inline
- void mbarrier_expect_tx(int mbar_addr, int size) {
+ __device__ inline void mbarrier_expect_tx(int mbar_addr, int size) {
      int mbar = mbar_addr & (int)SM100_MMA_PEER_MASK;
      asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;"
-                         :: "r"(mbar), "r"(size) : "memory");
+                  :: "r"(mbar), "r"(size) : "memory");
  }
   
- __device__ inline
- void tma_load_1d(int dst, const void *tmap_ptr, int x, int mbar_addr, uint64_t cache_policy) {
+ __device__ inline void tma_load_1d(int dst, const void *tmap_ptr, int x, int mbar_addr, uint64_t cache_policy) {
      uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(tmap_ptr);
      uint32_t smem_int_mbar = (uint32_t)mbar_addr & SM100_MMA_PEER_MASK;
      uint32_t smem_int_ptr  = (uint32_t)dst;
      asm volatile("cp.async.bulk.tensor.1d.cta_group::2.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint "
-                         "[%0], [%1, {%3}], [%2], %4;"
-                         :: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), "r"(x), "l"(cache_policy)
-                         : "memory");
+                  "[%0], [%1, {%3}], [%2], %4;"
+                  :: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), "r"(x), "l"(cache_policy) : "memory");
  }
   
- __device__ inline
- void tma_load_3d(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
+ __device__ inline void tma_load_3d(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
      uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(tmap_ptr);
      uint32_t smem_int_mbar = (uint32_t)mbar_addr & SM100_MMA_PEER_MASK;
      uint32_t smem_int_ptr  = (uint32_t)dst;
      asm volatile("cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint "
-                         "[%0], [%1, {%3, %4, %5}], [%2], %6;"
-                         :: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
-                            "r"(x), "r"(y), "r"(z), "l"(cache_policy)
-                         : "memory");
+                  "[%0], [%1, {%3, %4, %5}], [%2], %6;"
+                  :: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
+                     "r"(x), "r"(y), "r"(z), "l"(cache_policy) : "memory");
  }
   
+ __device__ inline void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
+     asm volatile("tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
+ }
+  
+ __device__ inline void tcgen05_mma_nvfp4(uint64_t a_desc, uint64_t b_desc, uint32_t i_desc,
+     int scale_A_tmem, int scale_B_tmem, int enable_input_d, int d_tmem) {
+     asm volatile(
+         "{\\n\\t"
+         ".reg .pred p;\\n\\t"
+         "setp.ne.b32 p, %6, 0;\\n\\t"
+         "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\\n\\t"
+         "}"
+         :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
+            "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d));
+ }
+  
+ __device__ inline void tcgen05_commit(int mbar_addr, uint16_t ctamask = 0x3) {
+     asm volatile("tcgen05.commit.cta_group::2.mbarrier::arrive::one.multicast::cluster.b64 [%0], %1;"
+                  :: "r"(mbar_addr), "h"(ctamask) : "memory");
+ }
+  
+ struct SHAPE { static constexpr char _16x256b[] = ".16x256b"; };
+ struct NUM { static constexpr char x8[] = ".x8"; static constexpr char x16[] = ".x16"; };
+  
+ template <const char *SHAPE, const char *NUM>
+ __device__ inline void tcgen05_ld_32regs(float *tmp, int row, int col) {
+     asm volatile("tcgen05.ld.sync.aligned%33%34.b32 "
+         "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
+         "  %8,  %9, %10, %11, %12, %13, %14, %15, "
+         " %16, %17, %18, %19, %20, %21, %22, %23, "
+         " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];"
+         : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]), "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7]),
+           "=f"(tmp[8]), "=f"(tmp[9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
+           "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
+           "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31])
+         : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
+ }
+  
+ template <const char *SHAPE, const char *NUM>
+ __device__ inline void tcgen05_ld_64regs(float *tmp, int row, int col) {
+     asm volatile("tcgen05.ld.sync.aligned%65%66.b32 "
+         "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
+         "  %8,  %9, %10, %11, %12, %13, %14, %15, "
+         " %16, %17, %18, %19, %20, %21, %22, %23, "
+         " %24, %25, %26, %27, %28, %29, %30, %31, "
+         " %32, %33, %34, %35, %36, %37, %38, %39, "
+         " %40, %41, %42, %43, %44, %45, %46, %47, "
+         " %48, %49, %50, %51, %52, %53, %54, %55, "
+         " %56, %57, %58, %59, %60, %61, %62, %63}, [%64];"
+         : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]), "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7]),
+           "=f"(tmp[8]), "=f"(tmp[9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
+           "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
+           "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31]),
+           "=f"(tmp[32]), "=f"(tmp[33]), "=f"(tmp[34]), "=f"(tmp[35]), "=f"(tmp[36]), "=f"(tmp[37]), "=f"(tmp[38]), "=f"(tmp[39]),
+           "=f"(tmp[40]), "=f"(tmp[41]), "=f"(tmp[42]), "=f"(tmp[43]), "=f"(tmp[44]), "=f"(tmp[45]), "=f"(tmp[46]), "=f"(tmp[47]),
+           "=f"(tmp[48]), "=f"(tmp[49]), "=f"(tmp[50]), "=f"(tmp[51]), "=f"(tmp[52]), "=f"(tmp[53]), "=f"(tmp[54]), "=f"(tmp[55]),
+           "=f"(tmp[56]), "=f"(tmp[57]), "=f"(tmp[58]), "=f"(tmp[59]), "=f"(tmp[60]), "=f"(tmp[61]), "=f"(tmp[62]), "=f"(tmp[63])
+         : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
+ }
+  
+ __device__ inline void tcgen05_ld_16x256bx8(float *tmp, int row, int col) { tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col); }
+ __device__ inline void tcgen05_ld_16x256bx16(float *tmp, int row, int col) { tcgen05_ld_64regs<SHAPE::_16x256b, NUM::x16>(tmp, row, col); }
+  
+ __device__ inline float silu(float x) { return x * __fdividef(1.0f, (1.0f + expf(-x))); }
+  
+ void check_cu(CUresult err) {
+     if (err == CUDA_SUCCESS) return;
+     const char *error_msg_ptr;
+     if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS) error_msg_ptr = "unable to get error string";
+     TORCH_CHECK(false, "cuTensorMapEncodeTiled error: ", error_msg_ptr);
+ }
+  
+ void init_AB_tmap(CUtensorMap *tmap, const char *ptr, uint64_t global_height, uint64_t global_width, uint32_t shared_height, uint32_t shared_width) {
+     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 elementStrides[rank] = {1, 1, 1};
+     check_cu(cuTensorMapEncodeTiled(tmap, 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_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
+ }
+  
+ void init_SF_tmap(CUtensorMap *tmap, const char *ptr, uint64_t total_bytes, uint32_t tile_bytes) {
+     TORCH_CHECK((total_bytes % 8) == 0 && (tile_bytes % 8) == 0, "SF bytes must be 8B aligned");
+     uint64_t globalDim[1] = {total_bytes / 8};
+     uint64_t globalStrides[1] = {globalDim[0] * 8};
+     uint32_t boxDim[1] = {tile_bytes / 8};
+     uint32_t elementStrides[1] = {1};
+     check_cu(cuTensorMapEncodeTiled(tmap, CU_TENSOR_MAP_DATA_TYPE_UINT64, 1, (void *)ptr,
+         globalDim, globalStrides, boxDim, elementStrides,
+         CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
+ }
+  
+ // ============== sub_test.py kernel (for M=512) ==============
+  
  template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
- __device__ inline
- void issue_tma_interleaved(
+ __device__ inline void issue_tma_sub(
      int smem, int stage_id, int iter_k,
      const CUtensorMap *A_tmap, const CUtensorMap *B1_tmap, const CUtensorMap *B2_tmap,
      const CUtensorMap *SFA_tmap, const CUtensorMap *SFB1_tmap, const CUtensorMap *SFB2_tmap,
-     int off_m, int off_n, int K, int ctarank,
-     int mbar_addr, uint64_t cache_A, uint64_t cache_B, bool dbg_print
+     int off_m, int off_n, int K, int ctarank, int mbar_addr, uint64_t cache_A, uint64_t cache_B
  ) {
-     constexpr int B_FRAG_N = BLOCK_N / 2;
+     constexpr int OUT_N = BLOCK_N / 2;
      constexpr int A_size = BLOCK_M * BLOCK_K / 2;
-     constexpr int B_size = B_FRAG_N * BLOCK_K / 2;
+     constexpr int B_size = OUT_N * BLOCK_K / 2;
      constexpr int SF_size = 128 * BLOCK_K / 16;
-     constexpr int STAGE_SIZE = A_size + B_size * 2 + SF_size * 3;
+     constexpr int STAGE_SIZE = A_size + B_size + SF_size * 3;
   
-  
      const int A_smem = smem + stage_id * STAGE_SIZE;
-     const int B1_smem = A_smem + A_size;
-     const int B2_smem = B1_smem + B_size;
-     const int SFA_smem = B2_smem + B_size;
+     const int B_smem = A_smem + A_size;
+     const int SFA_smem = B_smem + B_size;
      const int SFB1_smem = SFA_smem + SF_size;
      const int SFB2_smem = SFB1_smem + SF_size;
   
      const int off_k = iter_k * BLOCK_K;
-     const int off_n_frag = off_n + ctarank * (BLOCK_N / 2);
+     const CUtensorMap *B_tmap = (ctarank == 0) ? B1_tmap : B2_tmap;
   
-  
      tma_load_3d(A_smem, A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_B);
-     tma_load_3d(B1_smem, B1_tmap, 0, off_n_frag, off_k / 256, mbar_addr, cache_B);
-     tma_load_3d(B2_smem, B2_tmap, 0, off_n_frag, off_k / 256, mbar_addr, cache_B);
+     tma_load_3d(B_smem, B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
   
      const int rest_k = K / 16 / 4;
      constexpr int SF_ELEM_BYTES = 8;
  ⋯ 2 unchanged lines
      const int sfb_coord = ((off_n / 128) * rest_k + off_k / (16 * 4)) * SF_ELEMS_PER_512B;
   
      tma_load_1d(SFA_smem, SFA_tmap, sfa_coord, mbar_addr, cache_B);
-     tma_load_1d(SFB1_smem, SFB1_tmap, sfb_coord, mbar_addr, cache_B);
-     tma_load_1d(SFB2_smem, SFB2_tmap, sfb_coord, mbar_addr, cache_B);
   
+     const int SFB_smem = SFB1_smem;
+     #pragma unroll
+     for (int k = 0; k < BLOCK_K / MMA_K; k++) {
+         const int k_sfb_coord = sfb_coord + k * (512 / SF_ELEM_BYTES);
+         tma_load_1d(SFB_smem + k * 1024 + 0,   SFB1_tmap, k_sfb_coord, mbar_addr, cache_B);
+         tma_load_1d(SFB_smem + k * 1024 + 512, SFB2_tmap, k_sfb_coord, mbar_addr, cache_B);
+     }
+  
      mbarrier_expect_tx(mbar_addr, STAGE_SIZE);
  }
   
- __device__ inline
- void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
-     asm volatile("tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
- }
+ template <int BLOCK_N>
+ __device__ inline void epilogue_sub(
+     int warp_id, int lane_id, int off_m, int off_n,
+     int gemm_tmem, half *C_ptr, int N, int done_mbar_addr
+ ) {
+     constexpr int OUT_N = BLOCK_N / 2;
+     mbarrier_wait(done_mbar_addr, 0);
+     asm volatile("tcgen05.fence::after_thread_sync;");
+     #pragma unroll
+     for (int m = 0; m < 32 / 16; m++) {
+         float g1[OUT_N / 2];
+         float g2[OUT_N / 2];
   
- __device__ inline
- void tcgen05_mma_nvfp4(uint64_t a_desc, uint64_t b_desc, uint32_t i_desc,
-     int scale_A_tmem, int scale_B_tmem, int enable_input_d, int d_tmem) {
-     asm volatile(
-         "{\\n\\t"
-         ".reg .pred p;\\n\\t"
-         "setp.ne.b32 p, %6, 0;\\n\\t"
-         "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\\n\\t"
-         "}"
-         :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
-              "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
-     );
- }
+         if constexpr (OUT_N == 128) {
+             tcgen05_ld_16x256bx16(g1, warp_id * 32 + m * 16, gemm_tmem);
+             tcgen05_ld_16x256bx16(g2, warp_id * 32 + m * 16, gemm_tmem + OUT_N);
+         } else {
+             tcgen05_ld_16x256bx8(g1, warp_id * 32 + m * 16, gemm_tmem);
+             tcgen05_ld_16x256bx8(g2, warp_id * 32 + m * 16, gemm_tmem + OUT_N);
+         }
+         asm volatile("tcgen05.wait::ld.sync.aligned;");
   
- __device__ inline
- void tcgen05_commit(int mbar_addr, uint16_t ctamask = 0x3) {
-     asm volatile("tcgen05.commit.cta_group::2.mbarrier::arrive::one.multicast::cluster.b64 [%0], %1;"
-                         :: "r"(mbar_addr), "h"(ctamask) : "memory");
- }
+         #pragma unroll
+         for (int i = 0; i < OUT_N / 8; i++) {
+             const int row0 = warp_id * 32 + m * 16 + lane_id / 4;
+             const int row1 = row0 + 8;
+             const int col0 = i * 8 + (lane_id % 4) * 2;
   
- struct SHAPE {
-     static constexpr char _16x256b[] = ".16x256b";
- };
+             const float s00 = silu(g1[i * 4 + 0]);
+             const float s01 = silu(g1[i * 4 + 1]);
+             const float s10 = silu(g1[i * 4 + 2]);
+             const float s11 = silu(g1[i * 4 + 3]);
   
- struct NUM {
-     static constexpr char x8[]  = ".x8";
-     static constexpr char x16[] = ".x16";
- };
+             const float v00 = g2[i * 4 + 0] * s00;
+             const float v01 = g2[i * 4 + 1] * s01;
+             const float v10 = g2[i * 4 + 2] * s10;
+             const float v11 = g2[i * 4 + 3] * s11;
   
- template <const char *SHAPE, const char *NUM>
- __device__ inline
- void tcgen05_ld_32regs(float *tmp, int row, int col) {
-     asm volatile("tcgen05.ld.sync.aligned%33%34.b32 "
-         "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
-         "  %8,  %9, %10, %11, %12, %13, %14, %15, "
-         " %16, %17, %18, %19, %20, %21, %22, %23, "
-         " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];"
-         : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
-             "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
-             "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
-             "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31])
-         : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
+             const int out_row0 = off_m + row0;
+             const int out_row1 = off_m + row1;
+             const int out_col = off_n + col0;
+  
+             reinterpret_cast<half2 *>(C_ptr + out_row0 * N + out_col)[0] = __float22half2_rn({v00, v01});
+             reinterpret_cast<half2 *>(C_ptr + out_row1 * N + out_col)[0] = __float22half2_rn({v10, v11});
+         }
+     }
  }
   
- template <const char *SHAPE, const char *NUM>
- __device__ inline
- void tcgen05_ld_64regs(float *tmp, int row, int col) {
-     asm volatile("tcgen05.ld.sync.aligned%65%66.b32 "
-                 "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
-                 "  %8,  %9, %10, %11, %12, %13, %14, %15, "
-                 " %16, %17, %18, %19, %20, %21, %22, %23, "
-                 " %24, %25, %26, %27, %28, %29, %30, %31, "
-                 " %32, %33, %34, %35, %36, %37, %38, %39, "
-                 " %40, %41, %42, %43, %44, %45, %46, %47, "
-                 " %48, %49, %50, %51, %52, %53, %54, %55, "
-                 " %56, %57, %58, %59, %60, %61, %62, %63}, [%64];"
-                 : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
-                     "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
-                     "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
-                     "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31]),
-                     "=f"(tmp[32]), "=f"(tmp[33]), "=f"(tmp[34]), "=f"(tmp[35]), "=f"(tmp[36]), "=f"(tmp[37]), "=f"(tmp[38]), "=f"(tmp[39]),
-                     "=f"(tmp[40]), "=f"(tmp[41]), "=f"(tmp[42]), "=f"(tmp[43]), "=f"(tmp[44]), "=f"(tmp[45]), "=f"(tmp[46]), "=f"(tmp[47]),
-                     "=f"(tmp[48]), "=f"(tmp[49]), "=f"(tmp[50]), "=f"(tmp[51]), "=f"(tmp[52]), "=f"(tmp[53]), "=f"(tmp[54]), "=f"(tmp[55]),
-                     "=f"(tmp[56]), "=f"(tmp[57]), "=f"(tmp[58]), "=f"(tmp[59]), "=f"(tmp[60]), "=f"(tmp[61]), "=f"(tmp[62]), "=f"(tmp[63])
-                 : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
+ template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
+ __global__ __cluster_dims__(2) __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
+ void dual_gemm_sub_kernel(
+     const __grid_constant__ CUtensorMap A_tmap, const __grid_constant__ CUtensorMap B1_tmap,
+     const __grid_constant__ CUtensorMap B2_tmap, const __grid_constant__ CUtensorMap SFA_tmap,
+     const __grid_constant__ CUtensorMap SFB1_tmap, const __grid_constant__ CUtensorMap SFB2_tmap,
+     half *C_ptr, int M, int N, int K
+ ) {
+     const int tid = threadIdx.x, lane_id = tid % WARP_SIZE, warp_id = tid / WARP_SIZE;
+     const uint32_t ctarank = get_cluster_ctarank();
+     const bool is_cta0 = (ctarank == 0);
+     const int cluster_id = blockIdx.x / 2;
+  
+     constexpr int OUT_N = BLOCK_N / 2;
+     const int grid_n = N / OUT_N;
+     const int cluster_m = cluster_id / grid_n;
+     const int bid_n = cluster_id % grid_n;
+     const int base_m = cluster_m * (2 * BLOCK_M);
+     const int off_m = base_m + int(ctarank) * BLOCK_M;
+     const int off_n = bid_n * OUT_N;
+     const int bid_m = cluster_m * 2 + int(ctarank);
+  
+     constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;
+     const int num_iters = K / BLOCK_K;
+  
+     extern __shared__ __align__(1024) char smem_ptr[];
+     const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
+  
+     constexpr int A_size = BLOCK_M * BLOCK_K / 2;
+     constexpr int B_size = OUT_N * BLOCK_K / 2;
+     constexpr int SF_size = 128 * BLOCK_K / 16;
+     constexpr int STAGE_SIZE = A_size + B_size + SF_size * 3;
+  
+     #pragma nv_diag_suppress static_var_with_dynamic_init
+     __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
+     const int mbar_base = static_cast<int>(__cvta_generic_to_shared(mbars));
+     const int tma_mbar_addr = mbar_base;
+     const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
+     const int done_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
+  
+     constexpr int GEMM_D_TMEM = 0;
+     constexpr int SFA_tmem = BLOCK_N;
+     constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
+  
+     if (warp_id == 0 && elect_sync()) {
+         #pragma unroll
+         for (int i = 0; i < NUM_STAGES * 2 + 1; i++) {
+             const int count = (i < NUM_STAGES) ? 2 : 1;
+             mbarrier_init(mbar_base + i * 8, count);
+         }
+         asm volatile("fence.mbarrier_init.release.cluster;");
+     } else if (warp_id == 1) {
+         asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(TMEM_COLS));
+     }
+     __syncthreads();
+  
+     uint64_t cache_A = (M > N) ? EVICT_FIRST : EVICT_LAST;
+     uint64_t cache_B = (M > N) ? EVICT_LAST : EVICT_FIRST;
+  
+     auto make_desc_AB = [](int addr) -> uint64_t {
+         const int SBO = 8 * 128;
+         return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
+     };
+     auto make_desc_SF = [](int addr) -> uint64_t {
+         const int SBO = 8 * 16;
+         return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
+     };
+  
+     constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U) | ((uint32_t)BLOCK_N >> 3U << 17U) | ((uint32_t)(2 * BLOCK_M) >> 7U << 27U);
+  
+     if (warp_id == NUM_WARPS - 2 && elect_sync()) {
+         #pragma unroll
+         for (int iter_k = 0; iter_k < NUM_STAGES && iter_k < num_iters; iter_k++) {
+             const int tma_mbar = tma_mbar_addr + iter_k * 8;
+             issue_tma_sub<BLOCK_M, BLOCK_N, BLOCK_K>(
+                 smem, iter_k, iter_k,
+                 &A_tmap, &B1_tmap, &B2_tmap,
+                 &SFA_tmap, &SFB1_tmap, &SFB2_tmap,
+                 off_m, off_n, K, int(ctarank),
+                 tma_mbar, cache_A, cache_B);
+         }
+  
+         #pragma unroll
+         for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
+             const int stage_id = iter_k % NUM_STAGES;
+             const int tma_mbar = tma_mbar_addr + stage_id * 8;
+             mbarrier_wait(mma_mbar_addr + stage_id * 8, (iter_k / NUM_STAGES - 1) % 2);
+             issue_tma_sub<BLOCK_M, BLOCK_N, BLOCK_K>(
+                 smem, stage_id, iter_k,
+                 &A_tmap, &B1_tmap, &B2_tmap,
+                 &SFA_tmap, &SFB1_tmap, &SFB2_tmap,
+                 off_m, off_n, K, int(ctarank),
+                 tma_mbar, cache_A, cache_B);
+         }
+     }
+  
+     if (warp_id == NUM_WARPS - 1 && elect_sync() && is_cta0) {
+         #pragma unroll
+         for (int iter_k = 0; iter_k < num_iters; iter_k++) {
+             const int stage_id = iter_k % NUM_STAGES;
+             const int phase = (iter_k / NUM_STAGES) % 2;
+  
+             mbarrier_wait(tma_mbar_addr + stage_id * 8, phase);
+  
+             const int A_smem = smem + stage_id * STAGE_SIZE;
+             const int B_smem = A_smem + A_size;
+             const int SFA_smem = B_smem + B_size;
+             const int SFB_smem = SFA_smem + SF_size;
+  
+             constexpr uint64_t SF_desc_base = make_desc_SF(0);
+             const uint64_t SFA_desc = SF_desc_base + ((uint64_t)SFA_smem >> 4ULL);
+  
+             #pragma unroll
+             for (int k = 0; k < BLOCK_K / MMA_K; k++) {
+                 tcgen05_cp_nvfp4(SFA_tmem + k * 4, SFA_desc + (uint64_t)k * (512ULL >> 4ULL));
+                 const uint64_t SFB1_k_desc = SF_desc_base + ((uint64_t)(SFB_smem + k * 1024) >> 4ULL);
+                 const uint64_t SFB2_k_desc = SF_desc_base + ((uint64_t)(SFB_smem + k * 1024 + 512) >> 4ULL);
+                 tcgen05_cp_nvfp4(SFB_tmem + k * 8 + 0, SFB1_k_desc);
+                 tcgen05_cp_nvfp4(SFB_tmem + k * 8 + 4, SFB2_k_desc);
+             }
+  
+             #pragma unroll
+             for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
+                 #pragma unroll
+                 for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
+                     uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
+                     uint64_t b_desc = make_desc_AB(B_smem + k1 * OUT_N * 128 + k2 * 32);
+                     int k_sf = k1 * 4 + k2;
+                     const int scale_A_tmem = SFA_tmem + k_sf * 4 + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
+                     const int scale_B_tmem = SFB_tmem + k_sf * 8;
+                     const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
+                     tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d, GEMM_D_TMEM);
+                 }
+             }
+  
+             tcgen05_commit(mma_mbar_addr + stage_id * 8);
+         }
+         tcgen05_commit(done_mbar_addr);
+     }
+  
+     if (tid < BLOCK_M) {
+         epilogue_sub<BLOCK_N>(warp_id, lane_id, off_m, off_n, GEMM_D_TMEM, C_ptr, N, done_mbar_addr);
+     }
+  
+     __syncthreads();
+  
+     if (warp_id == 0)
+         asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(TMEM_COLS));
  }
   
- __device__ inline void tcgen05_ld_16x256bx8(float *tmp, int row, int col) { tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col); }
- __device__ inline void tcgen05_ld_16x256bx16(float *tmp, int row, int col) { tcgen05_ld_64regs<SHAPE::_16x256b, NUM::x16>(tmp, row, col); }
+ // ============== master_kernel2.py kernel (for M=256) ==============
   
- __device__ inline
- float silu(float x) {
-     //return x / (1.0f + expf(-x));
-     //return __fdividef(x, 1.0f + __expf(-x));
-     return x * __fdividef(1.0f, (1.0f + expf(-x)));
+ template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
+ __device__ inline void issue_tma_master(
+     int smem, int stage_id, int iter_k,
+     const CUtensorMap *A_tmap, const CUtensorMap *B1_tmap, const CUtensorMap *B2_tmap,
+     const CUtensorMap *SFA_tmap, const CUtensorMap *SFB1_tmap, const CUtensorMap *SFB2_tmap,
+     int off_m, int off_n, int K, int ctarank, int mbar_addr, uint64_t cache_A, uint64_t cache_B
+ ) {
+     constexpr int B_FRAG_N = BLOCK_N / 2;
+     constexpr int A_size = BLOCK_M * BLOCK_K / 2;
+     constexpr int B_size = B_FRAG_N * BLOCK_K / 2;
+     constexpr int SF_size = 128 * BLOCK_K / 16;
+     constexpr int STAGE_SIZE = A_size + B_size * 2 + SF_size * 3;
   
+     const int A_smem = smem + stage_id * STAGE_SIZE;
+     const int B1_smem = A_smem + A_size;
+     const int B2_smem = B1_smem + B_size;
+     const int SFA_smem = B2_smem + B_size;
+     const int SFB1_smem = SFA_smem + SF_size;
+     const int SFB2_smem = SFB1_smem + SF_size;
+  
+     const int off_k = iter_k * BLOCK_K;
+     const int off_n_frag = off_n + ctarank * (BLOCK_N / 2);
+  
+     tma_load_3d(A_smem, A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_B);
+     tma_load_3d(B1_smem, B1_tmap, 0, off_n_frag, off_k / 256, mbar_addr, cache_B);
+     tma_load_3d(B2_smem, B2_tmap, 0, off_n_frag, off_k / 256, mbar_addr, cache_B);
+  
+     const int rest_k = K / 16 / 4;
+     constexpr int SF_ELEM_BYTES = 8;
+     constexpr int SF_ELEMS_PER_512B = 512 / SF_ELEM_BYTES;
+     const int sfa_coord = ((off_m / 128) * rest_k + off_k / (16 * 4)) * SF_ELEMS_PER_512B;
+     const int sfb_coord = ((off_n / 128) * rest_k + off_k / (16 * 4)) * SF_ELEMS_PER_512B;
+  
+     tma_load_1d(SFA_smem, SFA_tmap, sfa_coord, mbar_addr, cache_B);
+     tma_load_1d(SFB1_smem, SFB1_tmap, sfb_coord, mbar_addr, cache_B);
+     tma_load_1d(SFB2_smem, SFB2_tmap, sfb_coord, mbar_addr, cache_B);
+  
+     mbarrier_expect_tx(mbar_addr, STAGE_SIZE);
  }
   
  template <int BLOCK_N>
- __device__ inline
- void epilogue_v1_baseline(
-     int warp_id, int lane_id,
-     int off_m, int off_n,
-     int gemm1_tmem, int gemm2_tmem,
-     half *C_ptr, int N,
-     int done_mbar_addr
+ __device__ inline void epilogue_master(
+     int warp_id, int lane_id, int off_m, int off_n,
+     int gemm1_tmem, int gemm2_tmem, half *C_ptr, int N, int done_mbar_addr
  ) {
      mbarrier_wait(done_mbar_addr, 0);
      asm volatile("tcgen05.fence::after_thread_sync;");
  ⋯ 37 unchanged lines
      }
  }
   
- void check_cu(CUresult err) {
-     if (err == CUDA_SUCCESS) return;
-     const char *error_msg_ptr;
-     if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS)
-         error_msg_ptr = "unable to get error string";
-     TORCH_CHECK(false, "cuTensorMapEncodeTiled error: ", error_msg_ptr);
- }
-  
- void init_AB_tmap(
-     CUtensorMap *tmap,
-     const char *ptr,
-     uint64_t global_height, uint64_t global_width,
-     uint32_t shared_height, uint32_t shared_width
- ) {
-     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 elementStrides[rank]  = {1, 1, 1};
-  
-     auto err = cuTensorMapEncodeTiled(
-         tmap,
-         CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
-         rank,
-         (void *)ptr,
-         globalDim,
-         globalStrides,
-         boxDim,
-         elementStrides,
-         CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
-         CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
-         CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
-         CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
-     );
-     check_cu(err);
- }
-  
- void init_SF_tmap(
-     CUtensorMap *tmap,
-     const char *ptr,
-     uint64_t total_bytes,
-     uint32_t tile_bytes
- ) {
-     constexpr uint32_t rank = 1;
-     TORCH_CHECK((total_bytes % 8) == 0, "SF total bytes must be 8B aligned");
-     TORCH_CHECK((tile_bytes % 8) == 0, "SF tile bytes must be 8B aligned");
-     uint64_t globalDim[rank]       = { total_bytes / 8 };
-     uint64_t globalStrides[1]      = { globalDim[0] * 8 };
-     uint32_t boxDim[rank]          = { tile_bytes / 8 };
-     uint32_t elementStrides[rank]  = { 1 };
-  
-     auto err = cuTensorMapEncodeTiled(
-         tmap,
-         CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT64,
-         rank,
-         (void *)ptr,
-         globalDim,
-         globalStrides,
-         boxDim,
-         elementStrides,
-         CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
-         CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE,
-         CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
-         CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
-     );
-     check_cu(err);
- }
-  
- // Interleaved dual GEMM kernel (2-SM cluster)
  template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
- __global__ __cluster_dims__(2)
- __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
- void dual_gemm_silu_kernel(
-     const __grid_constant__ CUtensorMap A_tmap,
-     const __grid_constant__ CUtensorMap B1_tmap,
-     const __grid_constant__ CUtensorMap B2_tmap,
-     const __grid_constant__ CUtensorMap SFA_tmap,
-     const __grid_constant__ CUtensorMap SFB1_tmap,
-     const __grid_constant__ CUtensorMap SFB2_tmap,
-     half *C_ptr,
-     int M, int N, int K
+ __global__ __cluster_dims__(2) __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
+ void dual_gemm_master_kernel(
+     const __grid_constant__ CUtensorMap A_tmap, const __grid_constant__ CUtensorMap B1_tmap,
+     const __grid_constant__ CUtensorMap B2_tmap, const __grid_constant__ CUtensorMap SFA_tmap,
+     const __grid_constant__ CUtensorMap SFB1_tmap, const __grid_constant__ CUtensorMap SFB2_tmap,
+     half *C_ptr, int M, int N, int K
  ) {
-     const int tid = threadIdx.x;
-     const int lane_id = tid % WARP_SIZE;
-     const int warp_id = tid / WARP_SIZE;
-  
+     const int tid = threadIdx.x, lane_id = tid % WARP_SIZE, warp_id = tid / WARP_SIZE;
      const uint32_t ctarank = get_cluster_ctarank();
      const bool is_cta0 = (ctarank == 0);
      const int cluster_id = blockIdx.x / 2;
-     const bool dbg = (cluster_id == 0 && ctarank == 0 && threadIdx.x == 0);
-     const bool dbg_cta1 = (cluster_id == 0 && ctarank == 1 && threadIdx.x == 0);
+  
      const int grid_n = N / BLOCK_N;
      const int cluster_m = cluster_id / grid_n;
      const int bid_n = cluster_id % grid_n;
  ⋯ 4 unchanged lines
   
      constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;
      const int num_iters = K / BLOCK_K;
-     const int tma_tid = (NUM_WARPS - 2) * WARP_SIZE;
-     const int mma_tid = (NUM_WARPS - 1) * WARP_SIZE;
-     const bool dbg_tma = (cluster_id == 0 && ctarank == 0 && threadIdx.x == tma_tid);
-     const bool dbg_tma_cta1 = (cluster_id == 0 && ctarank == 1 && threadIdx.x == tma_tid);
-     const bool dbg_mma = (cluster_id == 0 && ctarank == 0 && threadIdx.x == mma_tid);
-     const bool dbg_alloc = (cluster_id == 0 && ctarank == 0 && threadIdx.x == WARP_SIZE);
   
-  
      extern __shared__ __align__(1024) char smem_ptr[];
      const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
   
  ⋯ 16 unchanged lines
      constexpr int SFB1_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
      constexpr int SFB2_tmem = SFB1_tmem + 4 * (BLOCK_K / MMA_K);
   
-  
-     //barrier_cluster_arrive();
-     //barrier_cluster_wait();
-  
      if (warp_id == 0 && elect_sync()) {
          #pragma unroll
          for (int i = 0; i < NUM_STAGES * 2 + 1; i++) {
  ⋯ 1 unchanged lines
              mbarrier_init(mbar_base + i * 8, count);
          }
          asm volatile("fence.mbarrier_init.release.cluster;");
+     } else if (warp_id == 1) {
+         asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(TMEM_COLS));
      }
-     else if (warp_id == 1) {
-         asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;"
-             :: "r"(smem), "r"(TMEM_COLS));
-     }
      __syncthreads();
-    // barrier_cluster_arrive();
-    // barrier_cluster_wait();
   
      uint64_t cache_A = (M > N) ? EVICT_FIRST : EVICT_LAST;
      uint64_t cache_B = (M > N) ? EVICT_LAST : EVICT_FIRST;
  ⋯ 16 unchanged lines
          #pragma unroll
          for (int iter_k = 0; iter_k < NUM_STAGES && iter_k < num_iters; iter_k++) {
              const int tma_mbar = tma_mbar_addr + iter_k * 8;
-             const bool dbg_tma_iter = (dbg_tma || dbg_tma_cta1) && (iter_k == 0);
-             issue_tma_interleaved<BLOCK_M, BLOCK_N, BLOCK_K>(
+             issue_tma_master<BLOCK_M, BLOCK_N, BLOCK_K>(
                  smem, iter_k, iter_k,
                  &A_tmap, &B1_tmap, &B2_tmap,
                  &SFA_tmap, &SFB1_tmap, &SFB2_tmap,
                  off_m, off_n, K, int(ctarank),
-                 tma_mbar, cache_A, cache_B, dbg_tma_iter);
+                 tma_mbar, cache_A, cache_B);
          }
   
          #pragma unroll
  ⋯ 1 unchanged lines
              const int stage_id = iter_k % NUM_STAGES;
              const int tma_mbar = tma_mbar_addr + stage_id * 8;
              mbarrier_wait(mma_mbar_addr + stage_id * 8, (iter_k / NUM_STAGES - 1) % 2);
-             const bool dbg_tma_iter = dbg_tma && (iter_k == NUM_STAGES || iter_k == num_iters - 1);
-             issue_tma_interleaved<BLOCK_M, BLOCK_N, BLOCK_K>(
+             issue_tma_master<BLOCK_M, BLOCK_N, BLOCK_K>(
                  smem, stage_id, iter_k,
                  &A_tmap, &B1_tmap, &B2_tmap,
                  &SFA_tmap, &SFB1_tmap, &SFB2_tmap,
                  off_m, off_n, K, int(ctarank),
-                 tma_mbar, cache_A, cache_B, dbg_tma_iter);
+                 tma_mbar, cache_A, cache_B);
          }
      }
   
  ⋯ 24 unchanged lines
                  tcgen05_cp_nvfp4(SFB2_tmem + k * 4, SFB2_desc + (uint64_t)k * (512ULL >> 4ULL));
              }
   
-             // Interleaved GEMM1 and GEMM2 MMA
              #pragma unroll
              for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
                  #pragma unroll
  ⋯ 17 unchanged lines
      }
   
      if (tid < BLOCK_M) {
-         epilogue_v1_baseline<BLOCK_N>(
-             warp_id, lane_id,
-             off_m, off_n,
-             GEMM1_D_TMEM, GEMM2_D_TMEM,
-             C_ptr, N,
-             done_mbar_addr
-         );
+         epilogue_master<BLOCK_N>(warp_id, lane_id, off_m, off_n, GEMM1_D_TMEM, GEMM2_D_TMEM, C_ptr, N, done_mbar_addr);
      }
   
      __syncthreads();
-    // barrier_cluster_arrive();
-    // barrier_cluster_wait();
   
      if (warp_id == 0)
          asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(TMEM_COLS));
-  
-    // barrier_cluster_arrive();
-     //barrier_cluster_wait();
  }
   
- // BLOCK_N=64 kernel for m=256 cases
- at::Tensor dual_gemm_silu_n64(
-     const at::Tensor& A,
-     const at::Tensor& B1,
-     const at::Tensor& B2,
-     const at::Tensor& SFA,
-     const at::Tensor& SFB1,
-     const at::Tensor& SFB2,
-                 at::Tensor& C
+ // ============== Launcher for M=512 (sub_test.py n128) ==============
+  
+ at::Tensor dual_gemm_m512(
+     const at::Tensor& A, const at::Tensor& B1, const at::Tensor& B2,
+     const at::Tensor& SFA, const at::Tensor& SFB1, const at::Tensor& SFB2, at::Tensor& C
  ) {
-     const int M = A.size(0);
-     const int N = B1.size(0);
-     const int K = A.size(1) * 2;
+     const int M = A.size(0), N = B1.size(0), K = A.size(1) * 2;
   
      constexpr int BLOCK_M = 128;
-     constexpr int BLOCK_N = 64;
+     constexpr int BLOCK_N = 256;
      constexpr int BLOCK_K = 256;
-     constexpr int NUM_STAGES = 7;
+     constexpr int NUM_STAGES = 5;
   
-     auto A_ptr   = reinterpret_cast<const char *>(A.data_ptr());
-     auto B1_ptr  = reinterpret_cast<const char *>(B1.data_ptr());
-     auto B2_ptr  = reinterpret_cast<const char *>(B2.data_ptr());
-     auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
-     auto SFB1_ptr = reinterpret_cast<const char *>(SFB1.data_ptr());
-     auto SFB2_ptr = reinterpret_cast<const char *>(SFB2.data_ptr());
-     auto C_ptr   = reinterpret_cast<half *>(C.data_ptr());
+     CUtensorMap A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap;
+     init_AB_tmap(&A_tmap, (const char*)A.data_ptr(), M, K, BLOCK_M, BLOCK_K);
+     init_AB_tmap(&B1_tmap, (const char*)B1.data_ptr(), N, K, BLOCK_N / 2, BLOCK_K);
+     init_AB_tmap(&B2_tmap, (const char*)B2.data_ptr(), N, K, BLOCK_N / 2, BLOCK_K);
   
      int SF_sz = 128 * BLOCK_K / 16;
-     CUtensorMap A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap;
-     init_AB_tmap(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
-     init_AB_tmap(&B1_tmap, B1_ptr, N, K, BLOCK_N / 2, BLOCK_K);
-     init_AB_tmap(&B2_tmap, B2_ptr, N, K, BLOCK_N / 2, BLOCK_K);
+     int SFB_tile = 512;
      int rest_k = K / 64;
      uint64_t sfa_bytes = (uint64_t)(M / 128) * rest_k * 512;
      uint64_t sfb_bytes = (uint64_t)(N / 128) * rest_k * 512;
-     init_SF_tmap(&SFA_tmap, SFA_ptr, sfa_bytes, SF_sz);
-     init_SF_tmap(&SFB1_tmap, SFB1_ptr, sfb_bytes, SF_sz);
-     init_SF_tmap(&SFB2_tmap, SFB2_ptr, sfb_bytes, SF_sz);
+     init_SF_tmap(&SFA_tmap, (const char*)SFA.data_ptr(), sfa_bytes, SF_sz);
+     init_SF_tmap(&SFB1_tmap, (const char*)SFB1.data_ptr(), sfb_bytes, SFB_tile);
+     init_SF_tmap(&SFB2_tmap, (const char*)SFB2.data_ptr(), sfb_bytes, SFB_tile);
   
-     int num_tiles = (M / (2 * BLOCK_M)) * (N / BLOCK_N);
+     int num_tiles = (M / (2 * BLOCK_M)) * (N / (BLOCK_N / 2));
      int tb_size = BLOCK_M + 2 * WARP_SIZE;
   
      int A_sz = BLOCK_M * BLOCK_K / 2;
      int B_sz = (BLOCK_N / 2) * BLOCK_K / 2;
-     int stage_size = A_sz + B_sz * 2 + SF_sz * 3;
+     int stage_size = A_sz + B_sz + SF_sz * 3;
      int smem_size = stage_size * NUM_STAGES;
   
-     auto kernel = dual_gemm_silu_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
-  
+     auto kernel = dual_gemm_sub_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
      cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1);
      cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared);
  ⋯ 2 unchanged lines
      config.gridDim = dim3(num_tiles * 2, 1, 1);
      config.blockDim = dim3(tb_size, 1, 1);
      config.dynamicSmemBytes = smem_size;
-  
      cudaLaunchAttribute attrs[1];
      attrs[0].id = cudaLaunchAttributeClusterDimension;
      attrs[0].val.clusterDim.x = 2;
  ⋯ 2 unchanged lines
      config.attrs = attrs;
      config.numAttrs = 1;
   
-     cudaLaunchKernelEx(&config, kernel,
-         A_tmap, B1_tmap, B2_tmap,
-         SFA_tmap, SFB1_tmap, SFB2_tmap,
-         C_ptr, M, N, K
-     );
-  
+     cudaLaunchKernelEx(&config, kernel, A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap,
+                        (half*)C.data_ptr(), M, N, K);
      return C;
  }
   
- TORCH_LIBRARY(dual_gemm2_n64, m) {
-     m.def("dual_gemm_silu(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) C) -> Tensor");
-     m.impl("dual_gemm_silu", &dual_gemm_silu_n64);
- }
+ // ============== Launcher for M=256 (master_kernel2.py n64) ==============
   
- // BLOCK_N=128 kernel for m=512 cases
- at::Tensor dual_gemm_silu_n128(
-     const at::Tensor& A,
-     const at::Tensor& B1,
-     const at::Tensor& B2,
-     const at::Tensor& SFA,
-     const at::Tensor& SFB1,
-     const at::Tensor& SFB2,
-                 at::Tensor& C
+ at::Tensor dual_gemm_m256(
+     const at::Tensor& A, const at::Tensor& B1, const at::Tensor& B2,
+     const at::Tensor& SFA, const at::Tensor& SFB1, const at::Tensor& SFB2, at::Tensor& C
  ) {
-     const int M = A.size(0);
-     const int N = B1.size(0);
-     const int K = A.size(1) * 2;
+     const int M = A.size(0), N = B1.size(0), K = A.size(1) * 2;
   
      constexpr int BLOCK_M = 128;
-     constexpr int BLOCK_N = 128;
+     constexpr int BLOCK_N = 64;
      constexpr int BLOCK_K = 256;
-     constexpr int NUM_STAGES = 5;
+     constexpr int NUM_STAGES = 7;
   
-     auto A_ptr   = reinterpret_cast<const char *>(A.data_ptr());
-     auto B1_ptr  = reinterpret_cast<const char *>(B1.data_ptr());
-     auto B2_ptr  = reinterpret_cast<const char *>(B2.data_ptr());
-     auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
-     auto SFB1_ptr = reinterpret_cast<const char *>(SFB1.data_ptr());
-     auto SFB2_ptr = reinterpret_cast<const char *>(SFB2.data_ptr());
-     auto C_ptr   = reinterpret_cast<half *>(C.data_ptr());
+     CUtensorMap A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap;
+     init_AB_tmap(&A_tmap, (const char*)A.data_ptr(), M, K, BLOCK_M, BLOCK_K);
+     init_AB_tmap(&B1_tmap, (const char*)B1.data_ptr(), N, K, BLOCK_N / 2, BLOCK_K);
+     init_AB_tmap(&B2_tmap, (const char*)B2.data_ptr(), N, K, BLOCK_N / 2, BLOCK_K);
   
      int SF_sz = 128 * BLOCK_K / 16;
-     CUtensorMap A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap;
-     init_AB_tmap(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
-     init_AB_tmap(&B1_tmap, B1_ptr, N, K, BLOCK_N / 2, BLOCK_K);
-     init_AB_tmap(&B2_tmap, B2_ptr, N, K, BLOCK_N / 2, BLOCK_K);
      int rest_k = K / 64;
      uint64_t sfa_bytes = (uint64_t)(M / 128) * rest_k * 512;
      uint64_t sfb_bytes = (uint64_t)(N / 128) * rest_k * 512;
-     init_SF_tmap(&SFA_tmap, SFA_ptr, sfa_bytes, SF_sz);
-     init_SF_tmap(&SFB1_tmap, SFB1_ptr, sfb_bytes, SF_sz);
-     init_SF_tmap(&SFB2_tmap, SFB2_ptr, sfb_bytes, SF_sz);
+     init_SF_tmap(&SFA_tmap, (const char*)SFA.data_ptr(), sfa_bytes, SF_sz);
+     init_SF_tmap(&SFB1_tmap, (const char*)SFB1.data_ptr(), sfb_bytes, SF_sz);
+     init_SF_tmap(&SFB2_tmap, (const char*)SFB2.data_ptr(), sfb_bytes, SF_sz);
   
      int num_tiles = (M / (2 * BLOCK_M)) * (N / BLOCK_N);
      int tb_size = BLOCK_M + 2 * WARP_SIZE;
  ⋯ 3 unchanged lines
      int stage_size = A_sz + B_sz * 2 + SF_sz * 3;
      int smem_size = stage_size * NUM_STAGES;
   
-     auto kernel = dual_gemm_silu_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
-  
+     auto kernel = dual_gemm_master_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
      cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1);
      cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared);
  ⋯ 2 unchanged lines
      config.gridDim = dim3(num_tiles * 2, 1, 1);
      config.blockDim = dim3(tb_size, 1, 1);
      config.dynamicSmemBytes = smem_size;
-  
      cudaLaunchAttribute attrs[1];
      attrs[0].id = cudaLaunchAttributeClusterDimension;
      attrs[0].val.clusterDim.x = 2;
  ⋯ 2 unchanged lines
      config.attrs = attrs;
      config.numAttrs = 1;
   
-     cudaLaunchKernelEx(&config, kernel,
-         A_tmap, B1_tmap, B2_tmap,
-         SFA_tmap, SFB1_tmap, SFB2_tmap,
-         C_ptr, M, N, K
-     );
-  
+     cudaLaunchKernelEx(&config, kernel, A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap,
+                        (half*)C.data_ptr(), M, N, K);
      return C;
  }
   
- TORCH_LIBRARY(dual_gemm2_n128, m) {
+ TORCH_LIBRARY(dual_gemm_m512_lib, m) {
      m.def("dual_gemm_silu(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) C) -> Tensor");
-     m.impl("dual_gemm_silu", &dual_gemm_silu_n128);
+     m.impl("dual_gemm_silu", &dual_gemm_m512);
  }
+  
+ TORCH_LIBRARY(dual_gemm_m256_lib, m) {
+     m.def("dual_gemm_silu(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) C) -> Tensor");
+     m.impl("dual_gemm_silu", &dual_gemm_m256);
+ }
  """
   
- # Compile both kernels together so they share common code.
  load_inline(
-     "dual_gemm_kernels2",
+     "dual_gemm_v2",
      cpp_sources="",
      cuda_sources=cuda_src,
-   #  verbose=True,
      is_python_module=False,
      no_implicit_headers=True,
      extra_cuda_cflags=[
-         "-O3",
-         "-gencode=arch=compute_100a,code=sm_100a",
-         "--use_fast_math",
-         "--expt-relaxed-constexpr",
-         "--relocatable-device-code=false",
-         "-lineinfo",
-      #   "-Xptxas=-v",
+         "-O3", "-gencode=arch=compute_100a,code=sm_100a",
+         "--use_fast_math", "--expt-relaxed-constexpr",
+         "--relocatable-device-code=false", "-lineinfo",
      ],
      extra_ldflags=["-lcuda"],
  )
   
- dual_gemm_silu_n64 = torch.ops.dual_gemm2_n64.dual_gemm_silu
- dual_gemm_silu_n128 = torch.ops.dual_gemm2_n128.dual_gemm_silu
+ dual_gemm_m512 = torch.ops.dual_gemm_m512_lib.dual_gemm_silu
+ dual_gemm_m256 = torch.ops.dual_gemm_m256_lib.dual_gemm_silu
   
   
  def custom_kernel(data: input_t) -> output_t:
  ⋯ 2 unchanged lines
      c = data[9]
   
      M = a.shape[0]
-     if M == 256:
-         # Use 256x64 (2-CTA cluster) for M=256
-         return dual_gemm_silu_n64(a, b1, b2, sfa_perm, sfb1_perm, sfb2_perm, c)
+     if M == 512:
+         # M=512: sub_test.py n128 pathway
+         return dual_gemm_m512(a, b1, b2, sfa_perm, sfb1_perm, sfb2_perm, c)
      else:
-         # Use 256x128 (2-CTA cluster) for M=512
-         return dual_gemm_silu_n128(a, b1, b2, sfa_perm, sfb1_perm, sfb2_perm, c)
+         # M=256: master_kernel2.py n64 pathway
+         return dual_gemm_m256(a, b1, b2, sfa_perm, sfb1_perm, sfb2_perm, c)