submission 466012

  ⋯ 6 unchanged lines
  import cutlass.utils.blockscaled_layout as blockscaled_utils
  from cutlass.cute.runtime import make_ptr
   
- import functools
- from typing import Tuple, List, Dict, Any
-  
  import torch
+ import traceback
+ from typing import Dict, Any, Tuple
  from task import input_t, output_t
   
  # -----------------------------------------------------------------------------
- # Tunables
+ # Kernel configuration parameters
  # -----------------------------------------------------------------------------
  bytes_per_tensormap = 128
  num_tensormaps = 4
   
- # NVFP4 UMMA op requires M-mode=128 (cannot be 64)
  mma_tiler_mnk = (128, 128, 256)
  mma_inst_shape_k = 64
   
  ⋯ 2 unchanged lines
  c_dtype = cutlass.Float16
   
  sf_vec_size = 16
-  
  threads_per_cta = 128
   
  num_acc_stage = 1
- num_ab_stage = 2  # try 2/3/4 later
   
- # safe (can tune later)
+ # ✅ B200에서 stage=4는 SMEM 점유율/occupancy 박살로 느려지는 경우가 잦음
+ #    일단 stage=2로 두고, 나중에 (TMEM/SMEM footprint 줄인 뒤) stage 늘리는 게 맞음
+ num_ab_stage = 2
+  
+ # Must be power-of-two, multiple of 32, <= 512
  num_tmem_alloc_cols = 512
   
   
  ⋯ 1 unchanged lines
      return (a + b - 1) // b
   
   
+ def _raise_picklable(prefix: str):
+     tb = traceback.format_exc()
+     raise RuntimeError(f"{prefix}\n{tb}") from None
+  
+  
+ def _ps_key(problem_sizes):
+     return tuple(tuple(int(x) for x in ps) for ps in problem_sizes)
+  
+  
  # -----------------------------------------------------------------------------
- # 1) Init tensormaps kernel (one CTA per group)
- #    Writes tensormap descriptors into tensormaps[group, 0..3, :]
+ # Global caches
  # -----------------------------------------------------------------------------
+ _compiled_init_cache: Dict[int, Any] = {}
+ _compiled_gemm_cache: Dict[int, Any] = {}
+ _runtime_cache: Dict[Any, Dict[str, Any]] = {}
+  
+  
+ # -----------------------------------------------------------------------------
+ # Init tensormaps kernel (one CTA per group)
+ # -----------------------------------------------------------------------------
  @cute.kernel
  def init_tensormaps_kernel(
      tma_atom_a: cute.CopyAtom,
  ⋯ 16 unchanged lines
      k = tensor_of_problem_sizes[group_idx, 2]
      l = cutlass.Int32(1)
   
-     # Shared buffer for building descriptors
      size_tensormap_in_i64 = num_tensormaps * bytes_per_tensormap // 8
   
      @cute.struct
  ⋯ 9 unchanged lines
      tensormap_sfa_smem_ptr = tensormap_b_smem_ptr + bytes_per_tensormap // 8
      tensormap_sfb_smem_ptr = tensormap_sfa_smem_ptr + bytes_per_tensormap // 8
   
-     # Target gmem descriptor locations
      tensormap_manager = utils.TensorMapManager(utils.TensorMapUpdateMode.SMEM, 128)
   
      tensormap_a_gmem_ptr = tensormap_manager.get_tensormap_ptr(
  ⋯ 9 unchanged lines
          tensormaps[(group_idx, 3, None)].iterator
      )
   
-     # Real pointers
      mA_mkl_iter = cute.make_ptr(
          ab_dtype, tensor_of_abc_ptrs[group_idx, 0], cute.AddressSpace.gmem
      ).align(32)
  ⋯ 7 unchanged lines
          sf_dtype, tensor_of_sfasfb_ptrs[group_idx, 1], cute.AddressSpace.gmem
      ).align(32)
   
-     # Layouts
      mA_mkl_layout = cute.make_layout(
-         (m, k, l), stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32))
+         (m, k, l),
+         stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
      )
      mB_nkl_layout = cute.make_layout(
-         (n, k, l), stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32))
+         (n, k, l),
+         stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
      )
   
-     # SFA/SFB special layout (cublas doc)
+     # Scale factors layout (cublas doc)
      atom_shape = ((32, 4), (sf_vec_size, 4))
      atom_stride = ((16, 4), (0, 1))
      sfa_layout = cute.tile_to_shape(
  ⋯ 12 unchanged lines
      real_tensor_sfa = cute.make_tensor(sfa_mkl_iter, sfa_layout)
      real_tensor_sfb = cute.make_tensor(sfb_nkl_iter, sfb_layout)
   
-     # Only warp0 builds & writes descriptors
      if warp_idx == 0:
          tensormap_manager.init_tensormap_from_atom(tma_atom_a, tensormap_a_smem_ptr, 0)
          tensormap_manager.init_tensormap_from_atom(tma_atom_b, tensormap_b_smem_ptr, 0)
-         tensormap_manager.init_tensormap_from_atom(
-             tma_atom_sfa, tensormap_sfa_smem_ptr, 0
-         )
-         tensormap_manager.init_tensormap_from_atom(
-             tma_atom_sfb, tensormap_sfb_smem_ptr, 0
-         )
+         tensormap_manager.init_tensormap_from_atom(tma_atom_sfa, tensormap_sfa_smem_ptr, 0)
+         tensormap_manager.init_tensormap_from_atom(tma_atom_sfb, tensormap_sfb_smem_ptr, 0)
   
          tensormap_manager.update_tensormap(
              (real_tensor_a, real_tensor_b, real_tensor_sfa, real_tensor_sfb),
              (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb),
-             (
-                 tensormap_a_gmem_ptr,
-                 tensormap_b_gmem_ptr,
-                 tensormap_sfa_gmem_ptr,
-                 tensormap_sfb_gmem_ptr,
-             ),
+             (tensormap_a_gmem_ptr, tensormap_b_gmem_ptr, tensormap_sfa_gmem_ptr, tensormap_sfb_gmem_ptr),
              0,
-             (
-                 tensormap_a_smem_ptr,
-                 tensormap_b_smem_ptr,
-                 tensormap_sfa_smem_ptr,
-                 tensormap_sfb_smem_ptr,
-             ),
+             (tensormap_a_smem_ptr, tensormap_b_smem_ptr, tensormap_sfa_smem_ptr, tensormap_sfb_smem_ptr),
          )
   
          tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr)
  ⋯ 5 unchanged lines
   
   
  # -----------------------------------------------------------------------------
- # 2) Compute kernel (NO tensormap update inside CTA)
+ # GEMM kernel (cluster_meta mapping)
  # -----------------------------------------------------------------------------
  @cute.kernel
  def gemm_kernel(
      tiled_mma: cute.TiledMma,
      tma_atom_a: cute.CopyAtom,
-     mA_mkl: cute.Tensor,  # proxy
+     mA_mkl: cute.Tensor,
      tma_atom_b: cute.CopyAtom,
-     mB_nkl: cute.Tensor,  # proxy
+     mB_nkl: cute.Tensor,
      tma_atom_sfa: cute.CopyAtom,
-     mSFA_mkl: cute.Tensor,  # proxy
+     mSFA_mkl: cute.Tensor,
      tma_atom_sfb: cute.CopyAtom,
-     mSFB_nkl: cute.Tensor,  # proxy
-     tensor_of_abc_ptrs: cute.Tensor,
-     tensor_of_sfasfb_ptrs: cute.Tensor,
-     tensormaps: cute.Tensor,              # [G,4,16] pre-initialized
-     tensor_of_problem_sizes: cute.Tensor,
-     tensor_of_cluster_meta: cute.Tensor,  # [total_clusters,3] -> (g, tx, ty)
+     mSFB_nkl: cute.Tensor,
+     tensor_of_abc_ptrs: cute.Tensor,          # [G,3] int64
+     tensor_of_sfasfb_ptrs: cute.Tensor,       # [G,2] int64
+     tensormaps: cute.Tensor,                  # [G,4,16] int64
+     tensor_of_problem_sizes: cute.Tensor,     # [G,4] int32
+     tensor_of_cluster_meta: cute.Tensor,      # [T,3] int32 -> (g, tx, ty)
      a_smem_layout_staged: cute.ComposedLayout,
      b_smem_layout_staged: cute.ComposedLayout,
      sfa_smem_layout_staged: cute.Layout,
  ⋯ 24 unchanged lines
          stride=(cute.assume(n, 32), 1, cute.assume(m * n, 32)),
      )
      mC_mnl = cute.make_tensor(mC_mnl_iter, mC_mnl_layout)
-  
      gC_mnl = cute.local_tile(
          mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (coord_x, coord_y, 0)
      )
   
-     # Shared storage (no tensormap buffer now)
      @cute.struct
      class SharedStorage:
          ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_ab_stage * 2]
  ⋯ 3 unchanged lines
      smem = utils.SmemAllocator()
      storage = smem.allocate(SharedStorage)
   
-     # SMEM tensors
      sA = smem.allocate_tensor(
          element_type=ab_dtype,
          layout=a_smem_layout_staged.outer,
  ⋯ 17 unchanged lines
          byte_alignment=128,
      )
   
-     # Pipelines
      ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
      ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
      ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
  ⋯ 11 unchanged lines
          consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, threads_per_cta),
      ).make_participants()
   
-     # Proxy partitioning
+     # Proxy tiles
      gA_mkl = cute.local_tile(
          mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
      )
  ⋯ 14 unchanged lines
      tCgSFB = thr_mma.partition_B(gSFB_nkl)
      tCgC = thr_mma.partition_C(gC_mnl)
   
-     # Read prebuilt tensormaps for this group
      tensormap_manager = utils.TensorMapManager(utils.TensorMapUpdateMode.SMEM, 128)
   
      tensormap_a_gmem_ptr = tensormap_manager.get_tensormap_ptr(
  ⋯ 9 unchanged lines
          tensormaps[(group_idx, 3, None)].iterator
      )
   
-     tma_desc_a = tensormap_manager.get_tensormap_ptr(
-         tensormap_a_gmem_ptr, cute.AddressSpace.generic
-     )
-     tma_desc_b = tensormap_manager.get_tensormap_ptr(
-         tensormap_b_gmem_ptr, cute.AddressSpace.generic
-     )
-     tma_desc_sfa = tensormap_manager.get_tensormap_ptr(
-         tensormap_sfa_gmem_ptr, cute.AddressSpace.generic
-     )
-     tma_desc_sfb = tensormap_manager.get_tensormap_ptr(
-         tensormap_sfb_gmem_ptr, cute.AddressSpace.generic
-     )
+     tma_desc_a = tensormap_manager.get_tensormap_ptr(tensormap_a_gmem_ptr, cute.AddressSpace.generic)
+     tma_desc_b = tensormap_manager.get_tensormap_ptr(tensormap_b_gmem_ptr, cute.AddressSpace.generic)
+     tma_desc_sfa = tensormap_manager.get_tensormap_ptr(tensormap_sfa_gmem_ptr, cute.AddressSpace.generic)
+     tma_desc_sfb = tensormap_manager.get_tensormap_ptr(tensormap_sfb_gmem_ptr, cute.AddressSpace.generic)
   
-     # TMA partition
+     # TMA partitions
      tAsA, tAgA = cpasync.tma_partition(
          tma_atom_a,
          0,
  ⋯ 35 unchanged lines
      acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
      tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
   
+     # Allocate TMEM
      tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=threads_per_cta)
      tmem = utils.TmemAllocator(storage.tmem_holding_buf, barrier_for_retrieve=tmem_alloc_barrier)
      tmem.allocate(num_tmem_alloc_cols)
      tmem.wait_for_alloc()
+  
      acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
      tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
   
-     # SFA/SFB in TMEM
+     # SFA/SFB TMEM tensors
      sfa_tmem_ptr = cute.recast_ptr(
          acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
          dtype=sf_dtype,
  ⋯ 20 unchanged lines
      )
      tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
   
-     # S2T copy for SFA/SFB (keep original pattern)
+     # S2T copy for SFs
      copy_atom_s2t = cute.make_copy_atom(
          tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
          sf_dtype,
  ⋯ 4 unchanged lines
      tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact)
      thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0)
      tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
-     tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t_sfa, tCsSFA_compact_s2t_)
+     tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+         tiled_copy_s2t_sfa, tCsSFA_compact_s2t_
+     )
      tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)
   
      tCsSFB_compact = cute.filter_zeros(sSFB)
  ⋯ 1 unchanged lines
      tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact)
      thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0)
      tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact)
-     tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t_sfb, tCsSFB_compact_s2t_)
+     tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+         tiled_copy_s2t_sfb, tCsSFB_compact_s2t_
+     )
      tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)
   
-     # k_tile_cnt
+     # K tiles
      k_tile_cnt = k // cutlass.Int32(mma_tiler_mnk[2])
   
-     # Slice tile coords
+     # Fix tile coords (so that later we index only by k_tile)
      tAgA = tAgA[(None, coord_x, None, 0)]
      tBgB = tBgB[(None, coord_y, None, 0)]
      tAgSFA = tAgSFA[(None, coord_x, None, 0)]
      tBgSFB = tBgSFB[(None, coord_y, None, 0)]
   
-     # Main loop
+     # -------------------------------------------------------------------------
+     # Warp0 mainloop: N-stage prefetch (여기선 num_ab_stage=2)
+     # -------------------------------------------------------------------------
      if warp_idx == 0:
          acc_empty = acc_producer.acquire_and_advance()
          tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
   
-         # prime tile0
-         ab_empty0 = ab_producer.acquire_and_advance()
-         cute.copy(
-             tma_atom_a,
-             tAgA[(None, 0)],
-             tAsA[(None, ab_empty0.index)],
-             tma_bar_ptr=ab_empty0.barrier,
-             tma_desc_ptr=tma_desc_a,
-         )
-         cute.copy(
-             tma_atom_b,
-             tBgB[(None, 0)],
-             tBsB[(None, ab_empty0.index)],
-             tma_bar_ptr=ab_empty0.barrier,
-             tma_desc_ptr=tma_desc_b,
-         )
-         cute.copy(
-             tma_atom_sfa,
-             tAgSFA[(None, 0)],
-             tAsSFA[(None, ab_empty0.index)],
-             tma_bar_ptr=ab_empty0.barrier,
-             tma_desc_ptr=tma_desc_sfa,
-         )
-         cute.copy(
-             tma_atom_sfb,
-             tBgSFB[(None, 0)],
-             tBsSFB[(None, ab_empty0.index)],
-             tma_bar_ptr=ab_empty0.barrier,
-             tma_desc_ptr=tma_desc_sfb,
-         )
-  
-         for k_tile in range(k_tile_cnt):
-             # prefetch next
-             if k_tile + 1 < k_tile_cnt:
+         # Prologue: issue first up to num_ab_stage tiles
+         for pre_k in range(num_ab_stage):
+             if pre_k < k_tile_cnt:
                  ab_empty = ab_producer.acquire_and_advance()
-                 kt = k_tile + 1
-                 cute.copy(
-                     tma_atom_a,
-                     tAgA[(None, kt)],
-                     tAsA[(None, ab_empty.index)],
-                     tma_bar_ptr=ab_empty.barrier,
-                     tma_desc_ptr=tma_desc_a,
-                 )
-                 cute.copy(
-                     tma_atom_b,
-                     tBgB[(None, kt)],
-                     tBsB[(None, ab_empty.index)],
-                     tma_bar_ptr=ab_empty.barrier,
-                     tma_desc_ptr=tma_desc_b,
-                 )
-                 cute.copy(
-                     tma_atom_sfa,
-                     tAgSFA[(None, kt)],
-                     tAsSFA[(None, ab_empty.index)],
-                     tma_bar_ptr=ab_empty.barrier,
-                     tma_desc_ptr=tma_desc_sfa,
-                 )
-                 cute.copy(
-                     tma_atom_sfb,
-                     tBgSFB[(None, kt)],
-                     tBsSFB[(None, ab_empty.index)],
-                     tma_bar_ptr=ab_empty.barrier,
-                     tma_desc_ptr=tma_desc_sfb,
-                 )
+                 st = ab_empty.index
   
+                 cute.copy(tma_atom_a, tAgA[(None, pre_k)], tAsA[(None, st)],
+                           tma_bar_ptr=ab_empty.barrier, tma_desc_ptr=tma_desc_a)
+                 cute.copy(tma_atom_b, tBgB[(None, pre_k)], tBsB[(None, st)],
+                           tma_bar_ptr=ab_empty.barrier, tma_desc_ptr=tma_desc_b)
+                 cute.copy(tma_atom_sfa, tAgSFA[(None, pre_k)], tAsSFA[(None, st)],
+                           tma_bar_ptr=ab_empty.barrier, tma_desc_ptr=tma_desc_sfa)
+                 cute.copy(tma_atom_sfb, tBgSFB[(None, pre_k)], tBsSFB[(None, st)],
+                           tma_bar_ptr=ab_empty.barrier, tma_desc_ptr=tma_desc_sfb)
+  
+         # Steady-state
+         for k_tile in range(k_tile_cnt):
              ab_full = ab_consumer.wait_and_advance()
+             st = ab_full.index
   
-             # S2T
-             s2t_stage_coord = (None, None, None, None, ab_full.index)
-             cute.copy(
-                 tiled_copy_s2t_sfa,
-                 tCsSFA_compact_s2t[s2t_stage_coord],
-                 tCtSFA_compact_s2t,
-             )
-             cute.copy(
-                 tiled_copy_s2t_sfb,
-                 tCsSFB_compact_s2t[s2t_stage_coord],
-                 tCtSFB_compact_s2t,
-             )
+             # S2T SFs for this stage
+             s2t_stage_coord = (None, None, None, None, st)
+             cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t[s2t_stage_coord], tCtSFA_compact_s2t)
+             cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t[s2t_stage_coord], tCtSFB_compact_s2t)
   
-             # GEMM
+             # MMA (UMMA)
              num_kblocks = cute.size(tCrA, mode=[2])
              for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
-                 kblock_coord = (None, None, kblock_idx, ab_full.index)
-  
+                 kblock_coord = (None, None, kblock_idx, st)
                  sf_kblock_coord = (None, None, kblock_idx)
+  
                  tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator)
                  tiled_mma.set(tcgen05.Field.SFB, tCtSFB[sf_kblock_coord].iterator)
   
  ⋯ 8 unchanged lines
   
              ab_full.release()
   
+             # refill: issue tile (k_tile + num_ab_stage)
+             next_k = k_tile + cutlass.Int32(num_ab_stage)
+             if next_k < k_tile_cnt:
+                 ab_empty = ab_producer.acquire_and_advance()
+                 stp = ab_empty.index
+  
+                 cute.copy(tma_atom_a, tAgA[(None, next_k)], tAsA[(None, stp)],
+                           tma_bar_ptr=ab_empty.barrier, tma_desc_ptr=tma_desc_a)
+                 cute.copy(tma_atom_b, tBgB[(None, next_k)], tBsB[(None, stp)],
+                           tma_bar_ptr=ab_empty.barrier, tma_desc_ptr=tma_desc_b)
+                 cute.copy(tma_atom_sfa, tAgSFA[(None, next_k)], tAsSFA[(None, stp)],
+                           tma_bar_ptr=ab_empty.barrier, tma_desc_ptr=tma_desc_sfa)
+                 cute.copy(tma_atom_sfb, tBgSFB[(None, next_k)], tBsSFB[(None, stp)],
+                           tma_bar_ptr=ab_empty.barrier, tma_desc_ptr=tma_desc_sfb)
+  
          acc_empty.commit()
   
-     # Epilogue
+     # -------------------------------------------------------------------------
+     # Epilogue: TMEM -> R -> GMEM
+     #   - full tile이면 pred 없이 store (M,N 둘 다 full인 타일이 대부분)
+     # -------------------------------------------------------------------------
      op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
      copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32)
      tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc[None, 0, 0])
      thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+  
      tDtAcc = thr_copy_t2r.partition_S(tCtAcc[None, 0, 0])
      tDgC = thr_copy_t2r.partition_D(tCgC[None, 0, 0])
   
  ⋯ 4 unchanged lines
      acc_full = acc_consumer.wait_and_advance()
   
      cute.copy(tiled_copy_t2r, tDtAcc, tDrAcc)
-     acc_vec = tDrAcc.load()
-     tDrC.store(acc_vec.to(c_dtype))
+     tDrC.store(tDrAcc.load().to(c_dtype))
   
-     # Store: N is multiple of 128, so only M-tail
-     simt_atom = cute.make_copy_atom(
-         cute.nvgpu.CopyUniversalOp(), c_dtype, num_bits_per_copy=16
-     )
-  
+     simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype, num_bits_per_copy=16)
      thread_layout = cute.make_layout((1, threads_per_cta), stride=(threads_per_cta, 1))
      value_layout = cute.make_layout((1, 1))
      tiled_copy_r2g = cute.make_tiled_copy_tv(simt_atom, thread_layout, value_layout)
  ⋯ 3 unchanged lines
      tDcC = thr_copy_r2g.partition_D(cC)
   
      residue_m = mC_mnl.shape[0] - cutlass.Int32(coord_x) * cutlass.Int32(mma_tiler_mnk[0])
+     residue_n = mC_mnl.shape[1] - cutlass.Int32(coord_y) * cutlass.Int32(mma_tiler_mnk[1])
   
+     # Fast path: full tile (대부분 여기)
      if residue_m >= cutlass.Int32(mma_tiler_mnk[0]):
-         cute.copy(simt_atom, cute.flatten(tDrC), cute.flatten(tDgC))
+         if residue_n >= cutlass.Int32(mma_tiler_mnk[1]):
+             cute.copy(simt_atom, cute.flatten(tDrC), cute.flatten(tDgC))
+         else:
+             tDpC = cute.make_rmem_tensor(tDrC.shape, cutlass.Boolean)
+             for i in range(cute.size(tDrC.shape)):
+                 tDpC[i] = cute.elem_less(tDcC[i], (residue_n, residue_m))
+             cute.copy(simt_atom, cute.flatten(tDrC), cute.flatten(tDgC), pred=cute.flatten(tDpC))
      else:
          tDpC = cute.make_rmem_tensor(tDrC.shape, cutlass.Boolean)
-         residue_n = cutlass.Int32(mma_tiler_mnk[1])
          for i in range(cute.size(tDrC.shape)):
              tDpC[i] = cute.elem_less(tDcC[i], (residue_n, residue_m))
-         cute.copy(
-             simt_atom,
-             cute.flatten(tDrC),
-             cute.flatten(tDgC),
-             pred=cute.flatten(tDpC),
-         )
+         cute.copy(simt_atom, cute.flatten(tDrC), cute.flatten(tDgC), pred=cute.flatten(tDpC))
   
      acc_full.release()
-  
      cute.arch.barrier()
      tmem.free(acc_tmem_ptr)
      pass
   
   
  # -----------------------------------------------------------------------------
- # Host-side JIT wrappers
+ # JIT wrappers
  # -----------------------------------------------------------------------------
  @cute.jit
- def init_tensormaps_jit(
+ def init_jit(
      ptr_of_tensor_of_problem_sizes: cute.Pointer,
      ptr_of_tensor_of_abc_ptrs: cute.Pointer,
      ptr_of_tensor_of_sfasfb_ptrs: cute.Pointer,
  ⋯ 35 unchanged lines
   
      sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_a.shape, sf_vec_size)
      sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_b.shape, sf_vec_size)
+     initial_sfa = cute.make_tensor(cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfa_layout)
+     initial_sfb = cute.make_tensor(cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfb_layout)
   
-     initial_sfa = cute.make_tensor(
-         cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfa_layout
-     )
-     initial_sfb = cute.make_tensor(
-         cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfb_layout
-     )
-  
      mma_op = tcgen05.MmaMXF4NVF4Op(
          sf_dtype,
          (mma_tiler_mnk[0], mma_tiler_mnk[1], mma_inst_shape_k),
  ⋯ 2 unchanged lines
      )
      tiled_mma = cute.make_tiled_mma(mma_op)
   
-     cluster_layout_vmnk = cute.tiled_divide(
-         cute.make_layout((1, 1, 1)),
-         (tiled_mma.thr_id.shape,),
-     )
+     cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,))
   
-     # SMEM layouts for atoms
-     a_smem_layout_staged = sm100_utils.make_smem_layout_a(
-         tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage
-     )
-     b_smem_layout_staged = sm100_utils.make_smem_layout_b(
-         tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage
-     )
-     sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
-         tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage
-     )
-     sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
-         tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage
-     )
+     a_smem_layout_staged = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage)
+     b_smem_layout_staged = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage)
+     sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage)
+     sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage)
   
      a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0))
      b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0))
  ⋯ 2 unchanged lines
   
      tma_atom_a, _ = cute.nvgpu.make_tiled_tma_atom_A(
          cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
-         initial_a,
-         a_smem_layout,
-         mma_tiler_mnk,
-         tiled_mma,
-         cluster_layout_vmnk.shape,
+         initial_a, a_smem_layout,
+         mma_tiler_mnk, tiled_mma, cluster_layout_vmnk.shape,
      )
      tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B(
          cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
-         initial_b,
-         b_smem_layout,
-         mma_tiler_mnk,
-         tiled_mma,
-         cluster_layout_vmnk.shape,
+         initial_b, b_smem_layout,
+         mma_tiler_mnk, tiled_mma, cluster_layout_vmnk.shape,
      )
      tma_atom_sfa, _ = cute.nvgpu.make_tiled_tma_atom_A(
          cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
-         initial_sfa,
-         sfa_smem_layout,
-         mma_tiler_mnk,
-         tiled_mma,
-         cluster_layout_vmnk.shape,
+         initial_sfa, sfa_smem_layout,
+         mma_tiler_mnk, tiled_mma, cluster_layout_vmnk.shape,
          internal_type=cutlass.Int16,
      )
      tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B(
          cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
-         initial_sfb,
-         sfb_smem_layout,
-         mma_tiler_mnk,
-         tiled_mma,
-         cluster_layout_vmnk.shape,
+         initial_sfb, sfb_smem_layout,
+         mma_tiler_mnk, tiled_mma, cluster_layout_vmnk.shape,
          internal_type=cutlass.Int16,
      )
   
      init_tensormaps_kernel(
-         tma_atom_a,
-         tma_atom_b,
-         tma_atom_sfa,
-         tma_atom_sfb,
-         tensor_of_abc_ptrs,
-         tensor_of_sfasfb_ptrs,
-         tensor_of_problem_sizes,
-         tensormaps,
+         tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb,
+         tensor_of_abc_ptrs, tensor_of_sfasfb_ptrs, tensor_of_problem_sizes, tensormaps,
      ).launch(
          grid=(1, 1, num_groups),
          block=[threads_per_cta, 1, 1],
  ⋯ 12 unchanged lines
      total_num_clusters: cutlass.Int32,
      num_groups: cutlass.Int32,
  ):
-     tensor_of_abc_ptrs = cute.make_tensor(
-         ptr_of_tensor_of_abc_ptrs, cute.make_layout((num_groups, 3), stride=(3, 1))
-     )
-     tensor_of_sfasfb_ptrs = cute.make_tensor(
-         ptr_of_tensor_of_sfasfb_ptrs, cute.make_layout((num_groups, 2), stride=(2, 1))
-     )
-     tensor_of_problem_sizes = cute.make_tensor(
-         ptr_of_tensor_of_problem_sizes, cute.make_layout((num_groups, 4), stride=(4, 1))
-     )
-     tensormaps = cute.make_tensor(
-         ptr_of_tensor_of_tensormap, cute.make_layout((num_groups, 4, 16), stride=(64, 16, 1))
-     )
-     tensor_of_cluster_meta = cute.make_tensor(
-         ptr_of_tensor_of_cluster_meta, cute.make_layout((total_num_clusters, 3), stride=(3, 1))
-     )
+     tensor_of_abc_ptrs = cute.make_tensor(ptr_of_tensor_of_abc_ptrs, cute.make_layout((num_groups, 3), stride=(3, 1)))
+     tensor_of_sfasfb_ptrs = cute.make_tensor(ptr_of_tensor_of_sfasfb_ptrs, cute.make_layout((num_groups, 2), stride=(2, 1)))
+     tensor_of_problem_sizes = cute.make_tensor(ptr_of_tensor_of_problem_sizes, cute.make_layout((num_groups, 4), stride=(4, 1)))
+     tensormaps = cute.make_tensor(ptr_of_tensor_of_tensormap, cute.make_layout((num_groups, 4, 16), stride=(64, 16, 1)))
+     tensor_of_cluster_meta = cute.make_tensor(ptr_of_tensor_of_cluster_meta, cute.make_layout((total_num_clusters, 3), stride=(3, 1)))
   
      # Fake tensors for atom creation
      min_a_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
  ⋯ 16 unchanged lines
   
      sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_a.shape, sf_vec_size)
      sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_b.shape, sf_vec_size)
+     initial_sfa = cute.make_tensor(cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfa_layout)
+     initial_sfb = cute.make_tensor(cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfb_layout)
   
-     initial_sfa = cute.make_tensor(
-         cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfa_layout
-     )
-     initial_sfb = cute.make_tensor(
-         cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfb_layout
-     )
-  
      mma_op = tcgen05.MmaMXF4NVF4Op(
          sf_dtype,
          (mma_tiler_mnk[0], mma_tiler_mnk[1], mma_inst_shape_k),
  ⋯ 2 unchanged lines
      )
      tiled_mma = cute.make_tiled_mma(mma_op)
   
-     cluster_layout_vmnk = cute.tiled_divide(
-         cute.make_layout((1, 1, 1)),
-         (tiled_mma.thr_id.shape,),
-     )
+     cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,))
   
-     # SMEM layouts
-     a_smem_layout_staged = sm100_utils.make_smem_layout_a(
-         tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage
-     )
-     b_smem_layout_staged = sm100_utils.make_smem_layout_b(
-         tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage
-     )
-     sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
-         tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage
-     )
-     sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
-         tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage
-     )
+     a_smem_layout_staged = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage)
+     b_smem_layout_staged = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage)
+     sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage)
+     sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage)
   
      a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0))
      b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0))
  ⋯ 2 unchanged lines
   
      tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
          cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
-         initial_a,
-         a_smem_layout,
-         mma_tiler_mnk,
-         tiled_mma,
-         cluster_layout_vmnk.shape,
+         initial_a, a_smem_layout,
+         mma_tiler_mnk, tiled_mma, cluster_layout_vmnk.shape,
      )
      tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
          cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
-         initial_b,
-         b_smem_layout,
-         mma_tiler_mnk,
-         tiled_mma,
-         cluster_layout_vmnk.shape,
+         initial_b, b_smem_layout,
+         mma_tiler_mnk, tiled_mma, cluster_layout_vmnk.shape,
      )
      tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
          cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
-         initial_sfa,
-         sfa_smem_layout,
-         mma_tiler_mnk,
-         tiled_mma,
-         cluster_layout_vmnk.shape,
+         initial_sfa, sfa_smem_layout,
+         mma_tiler_mnk, tiled_mma, cluster_layout_vmnk.shape,
          internal_type=cutlass.Int16,
      )
      tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
          cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
-         initial_sfb,
-         sfb_smem_layout,
-         mma_tiler_mnk,
-         tiled_mma,
-         cluster_layout_vmnk.shape,
+         initial_sfb, sfb_smem_layout,
+         mma_tiler_mnk, tiled_mma, cluster_layout_vmnk.shape,
          internal_type=cutlass.Int16,
      )
   
  ⋯ 6 unchanged lines
   
      gemm_kernel(
          tiled_mma,
-         tma_atom_a,
-         tma_tensor_a,
-         tma_atom_b,
-         tma_tensor_b,
-         tma_atom_sfa,
-         tma_tensor_sfa,
-         tma_atom_sfb,
-         tma_tensor_sfb,
+         tma_atom_a, tma_tensor_a,
+         tma_atom_b, tma_tensor_b,
+         tma_atom_sfa, tma_tensor_sfa,
+         tma_atom_sfb, tma_tensor_sfb,
          tensor_of_abc_ptrs,
          tensor_of_sfasfb_ptrs,
          tensormaps,
  ⋯ 12 unchanged lines
      return
   
   
- # -----------------------------------------------------------------------------
- # Compile caches
- # -----------------------------------------------------------------------------
- _compiled_init_cache: Dict[str, Any] = {}
- _compiled_gemm_cache: Dict[str, Any] = {}
+ def compile_init(num_groups: int):
+     if num_groups in _compiled_init_cache:
+         return _compiled_init_cache[num_groups]
+     try:
+         cute_ptr_ps = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
+         cute_ptr_abc = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
+         cute_ptr_sfs = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
+         cute_ptr_tm = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
+         ng = cutlass.Int32(num_groups)
   
- _runtime_cache: Dict[Any, Any] = {}
+         fn = cute.compile(init_jit, cute_ptr_ps, cute_ptr_abc, cute_ptr_sfs, cute_ptr_tm, ng)
+         _compiled_init_cache[num_groups] = fn
+         return fn
+     except Exception:
+         _raise_picklable("compile_init failed")
   
   
- def compile_init(num_groups: int):
-     key = f"ng={num_groups}"
-     if key in _compiled_init_cache:
-         return _compiled_init_cache[key]
+ def compile_gemm(num_groups: int):
+     if num_groups in _compiled_gemm_cache:
+         return _compiled_gemm_cache[num_groups]
+     try:
+         cute_ptr_ps = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
+         cute_ptr_abc = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
+         cute_ptr_sfs = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
+         cute_ptr_tm = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
+         cute_ptr_meta = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
   
-     cute_ptr_ps = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
-     cute_ptr_abc = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
-     cute_ptr_sfs = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
-     cute_ptr_tm = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
+         total_clusters = cutlass.Int32(1)
+         ng = cutlass.Int32(num_groups)
   
-     ng = cutlass.Int32(num_groups)
+         fn = cute.compile(gemm_jit, cute_ptr_ps, cute_ptr_abc, cute_ptr_sfs, cute_ptr_tm, cute_ptr_meta, total_clusters, ng)
+         _compiled_gemm_cache[num_groups] = fn
+         return fn
+     except Exception:
+         _raise_picklable("compile_gemm failed")
   
-     fn = cute.compile(
-         init_tensormaps_jit,
-         cute_ptr_ps,
-         cute_ptr_abc,
-         cute_ptr_sfs,
-         cute_ptr_tm,
-         ng,
-     )
-     _compiled_init_cache[key] = fn
-     return fn
   
+ def _get_or_create_runtime(problem_sizes):
+     num_groups = len(problem_sizes)
+     key = (num_groups, _ps_key(problem_sizes))
+     if key in _runtime_cache:
+         return _runtime_cache[key]
   
- def compile_gemm(num_groups: int):
-     key = f"ng={num_groups}"
-     if key in _compiled_gemm_cache:
-         return _compiled_gemm_cache[key]
+     # cluster_meta: ✅ tx-major로 만들어서 같은 A 타일(=tx)이 연속되게 배치 (L2 reuse 도움)
+     cluster_meta = []
+     total_num_clusters = 0
+     for g, (m, n, k, l) in enumerate(problem_sizes):
+         tiles_m = ceil_div(m, mma_tiler_mnk[0])
+         tiles_n = ceil_div(n, mma_tiler_mnk[1])
+         total_num_clusters += tiles_m * tiles_n
+         for tx in range(tiles_m):
+             for ty in range(tiles_n):
+                 cluster_meta.append((g, tx, ty))
   
-     cute_ptr_ps = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
-     cute_ptr_abc = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
-     cute_ptr_sfs = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
-     cute_ptr_tm = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
-     cute_ptr_meta = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
+     tensor_of_problem_sizes = torch.tensor(problem_sizes, dtype=torch.int32, device="cuda")
+     tensor_of_cluster_meta = torch.tensor(cluster_meta, dtype=torch.int32, device="cuda")
   
-     total_clusters = cutlass.Int32(1)
-     ng = cutlass.Int32(num_groups)
+     tensor_of_abc_ptrs = torch.empty((num_groups, 3), dtype=torch.int64, device="cuda")
+     tensor_of_sfs_ptrs = torch.empty((num_groups, 2), dtype=torch.int64, device="cuda")
+     tensor_of_tensormap = torch.empty((num_groups, 4, 16), dtype=torch.int64, device="cuda")
   
-     fn = cute.compile(
-         gemm_jit,
-         cute_ptr_ps,
-         cute_ptr_abc,
-         cute_ptr_sfs,
-         cute_ptr_tm,
-         cute_ptr_meta,
-         total_clusters,
-         ng,
-     )
-     _compiled_gemm_cache[key] = fn
-     return fn
+     # ✅ pinned host buffers: 매 호출마다 torch.tensor(list) 생성하지 않기
+     host_abc_ptrs = torch.empty((num_groups, 3), dtype=torch.int64, pin_memory=True)
+     host_sfs_ptrs = torch.empty((num_groups, 2), dtype=torch.int64, pin_memory=True)
   
+     # ✅ cached make_ptr
+     cute_ptr_ps = make_ptr(cutlass.Int32, tensor_of_problem_sizes.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     cute_ptr_abc = make_ptr(cutlass.Int64, tensor_of_abc_ptrs.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     cute_ptr_sfs = make_ptr(cutlass.Int64, tensor_of_sfs_ptrs.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     cute_ptr_tm = make_ptr(cutlass.Int64, tensor_of_tensormap.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+     cute_ptr_meta = make_ptr(cutlass.Int32, tensor_of_cluster_meta.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
   
- # -----------------------------------------------------------------------------
- # Entry point
- # -----------------------------------------------------------------------------
- def custom_kernel(data: input_t) -> output_t:
-     abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
-     num_groups = len(problem_sizes)
+     _runtime_cache[key] = {
+         "num_groups": num_groups,
+         "total_num_clusters": total_num_clusters,
+         "tensor_of_problem_sizes": tensor_of_problem_sizes,
+         "tensor_of_cluster_meta": tensor_of_cluster_meta,
+         "tensor_of_abc_ptrs": tensor_of_abc_ptrs,
+         "tensor_of_sfs_ptrs": tensor_of_sfs_ptrs,
+         "tensor_of_tensormap": tensor_of_tensormap,
+         "host_abc_ptrs": host_abc_ptrs,
+         "host_sfs_ptrs": host_sfs_ptrs,
+         "cute_ptr_ps": cute_ptr_ps,
+         "cute_ptr_abc": cute_ptr_abc,
+         "cute_ptr_sfs": cute_ptr_sfs,
+         "cute_ptr_tm": cute_ptr_tm,
+         "cute_ptr_meta": cute_ptr_meta,
+         # init 커널 스킵용 (A/B/SF만 체크: C 바뀌어도 init 필요 없음)
+         "last_ab_sfs_sig": None,
+     }
+     return _runtime_cache[key]
   
-     init_fn = compile_init(num_groups)
-     gemm_fn = compile_gemm(num_groups)
   
-     # Build host pointer tuples
-     abc_ptrs = []
-     sfs_ptrs = []
-     for (a, b, c), (sfa_r, sfb_r), _sz in zip(abc_tensors, sfasfb_reordered_tensors, problem_sizes):
-         abc_ptrs.append((a.data_ptr(), b.data_ptr(), c.data_ptr()))
-         sfs_ptrs.append((sfa_r.data_ptr(), sfb_r.data_ptr()))
+ def custom_kernel(data: input_t) -> output_t:
+     try:
+         abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
+         num_groups = len(problem_sizes)
   
-     # Runtime cache per (problem_sizes)
-     ps_key = tuple((int(m), int(n), int(k), int(l)) for (m, n, k, l) in problem_sizes)
-     cache_key = (num_groups, ps_key)
+         init_fn = compile_init(num_groups)
+         gemm_fn = compile_gemm(num_groups)
+         st = _get_or_create_runtime(problem_sizes)
   
-     if cache_key not in _runtime_cache:
-         # problem_sizes tensor (device)
-         tensor_of_problem_sizes = torch.tensor(problem_sizes, dtype=torch.int32, device="cuda")
+         host_abc = st["host_abc_ptrs"]
+         host_sfs = st["host_sfs_ptrs"]
   
-         # cluster_meta (device)
-         cluster_meta = []
-         cta_m = mma_tiler_mnk[0]
-         cta_n = mma_tiler_mnk[1]
-         for g, (m, n, k, l) in enumerate(problem_sizes):
-             tiles_m = ceil_div(m, cta_m)
-             tiles_n = ceil_div(n, cta_n)
-             for ty in range(tiles_n):
-                 for tx in range(tiles_m):
-                     cluster_meta.append((g, tx, ty))
-         total_num_clusters = len(cluster_meta)
-         tensor_of_cluster_meta = torch.tensor(cluster_meta, dtype=torch.int32, device="cuda")
+         # A/B/C pointers + SFA/SFB pointers 채우기
+         # 그리고 init 스킵을 위해 (A,B,SFA,SFB)만 signature 구성
+         sig = []
+         for i, ((a, b, c), (sfa_r, sfb_r)) in enumerate(zip(abc_tensors, sfasfb_reordered_tensors)):
+             a_ptr = a.data_ptr()
+             b_ptr = b.data_ptr()
+             c_ptr = c.data_ptr()
+             sfa_ptr = sfa_r.data_ptr()
+             sfb_ptr = sfb_r.data_ptr()
   
-         # persistent device pointer arrays
-         tensor_of_abc_ptrs = torch.empty((num_groups, 3), dtype=torch.int64, device="cuda")
-         tensor_of_sfs_ptrs = torch.empty((num_groups, 2), dtype=torch.int64, device="cuda")
+             host_abc[i, 0] = a_ptr
+             host_abc[i, 1] = b_ptr
+             host_abc[i, 2] = c_ptr
+             host_sfs[i, 0] = sfa_ptr
+             host_sfs[i, 1] = sfb_ptr
   
-         # tensormaps per group
-         tensor_of_tensormap = torch.empty((num_groups, 4, bytes_per_tensormap // 8), dtype=torch.int64, device="cuda")
+             sig.extend((a_ptr, b_ptr, sfa_ptr, sfb_ptr))
   
-         _runtime_cache[cache_key] = {
-             "tensor_of_problem_sizes": tensor_of_problem_sizes,
-             "tensor_of_cluster_meta": tensor_of_cluster_meta,
-             "total_num_clusters": total_num_clusters,
-             "tensor_of_abc_ptrs": tensor_of_abc_ptrs,
-             "tensor_of_sfs_ptrs": tensor_of_sfs_ptrs,
-             "tensor_of_tensormap": tensor_of_tensormap,
-             "last_abc_ptrs": None,
-             "last_sfs_ptrs": None,
-         }
+         ab_sfs_sig = tuple(sig)
   
-     st = _runtime_cache[cache_key]
+         # 포인터 배열은 작으니 매번 async copy (pinned -> cuda)
+         st["tensor_of_abc_ptrs"].copy_(host_abc, non_blocking=True)
+         st["tensor_of_sfs_ptrs"].copy_(host_sfs, non_blocking=True)
   
-     # Update pointer arrays if changed
-     need_init = False
-     if st["last_abc_ptrs"] != tuple(abc_ptrs) or st["last_sfs_ptrs"] != tuple(sfs_ptrs):
-         # small host tensors -> copy into persistent device tensors
-         cpu_abc = torch.tensor(abc_ptrs, dtype=torch.int64, device="cpu")
-         cpu_sfs = torch.tensor(sfs_ptrs, dtype=torch.int64, device="cpu")
-         st["tensor_of_abc_ptrs"].copy_(cpu_abc, non_blocking=False)
-         st["tensor_of_sfs_ptrs"].copy_(cpu_sfs, non_blocking=False)
-         st["last_abc_ptrs"] = tuple(abc_ptrs)
-         st["last_sfs_ptrs"] = tuple(sfs_ptrs)
-         need_init = True
+         # ✅ A/B/SF 포인터가 바뀐 경우에만 tensormap init
+         if st["last_ab_sfs_sig"] != ab_sfs_sig:
+             init_fn(st["cute_ptr_ps"], st["cute_ptr_abc"], st["cute_ptr_sfs"], st["cute_ptr_tm"], num_groups)
+             st["last_ab_sfs_sig"] = ab_sfs_sig
   
-     # CuTe pointers
-     cute_ptr_ps = make_ptr(
-         cutlass.Int32, st["tensor_of_problem_sizes"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
-     )
-     cute_ptr_abc = make_ptr(
-         cutlass.Int64, st["tensor_of_abc_ptrs"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
-     )
-     cute_ptr_sfs = make_ptr(
-         cutlass.Int64, st["tensor_of_sfs_ptrs"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
-     )
-     cute_ptr_tm = make_ptr(
-         cutlass.Int64, st["tensor_of_tensormap"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
-     )
-     cute_ptr_meta = make_ptr(
-         cutlass.Int32, st["tensor_of_cluster_meta"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
-     )
-  
-     # Init tensormaps once per (problem_sizes, pointers)
-     if need_init:
-         init_fn(
-             cute_ptr_ps,
-             cute_ptr_abc,
-             cute_ptr_sfs,
-             cute_ptr_tm,
+         # GEMM
+         gemm_fn(
+             st["cute_ptr_ps"],
+             st["cute_ptr_abc"],
+             st["cute_ptr_sfs"],
+             st["cute_ptr_tm"],
+             st["cute_ptr_meta"],
+             st["total_num_clusters"],
              num_groups,
          )
   
-     # GEMM
-     gemm_fn(
-         cute_ptr_ps,
-         cute_ptr_abc,
-         cute_ptr_sfs,
-         cute_ptr_tm,
-         cute_ptr_meta,
-         st["total_num_clusters"],
-         num_groups,
-     )
-  
-     return [abc_tensors[i][2] for i in range(num_groups)]
+         return [abc_tensors[i][2] for i in range(num_groups)]
+     except Exception:
+         _raise_picklable("custom_kernel failed")