submission 482269

  import torch
+ import ctypes
  from task import input_t, output_t
  from torch.utils.cpp_extension import load_inline
   
- # CUDA kernel source - contains the actual GEMM implementation
- cuda_source = r"""
- // Gen5 NVFP4 Group GEMM - Persistent Kernel with TMEM Double-Buffering
- // - FP4 (E2M1) input matrices A and B with MX block scaling
- // - FP8 (E4M3FN) scale factors
- // - FP16 output
- // - Persistent grid-stride loop: each block processes multiple tiles
- // - TMEM double-buffering: overlap epilogue(T) with MMA(T+1)
- // - L2 256B promotion for tensor maps
- // - Hash-based CUDA Graph caching
+ # Load CUDA runtime for direct cudaGraphLaunch via ctypes (bypasses pybind11)
+ _libcudart = ctypes.CDLL("libcudart.so")
+ _libcudart.cudaGraphLaunch.restype = ctypes.c_int
+ _libcudart.cudaGraphLaunch.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
+ # Pre-create the zero argument (default cuda handle = 0)
+ _ZERO = ctypes.c_void_p(0)
   
+ cuda_source = r"""
  #include <cudaTypedefs.h>
  #include <cuda_fp16.h>
  #include <c10/util/Half.h>
- #include <vector>
   
  constexpr int WARP_SIZE = 32;
- constexpr int MMA_K = 64;  // 64 elements per MMA in K dimension for NVFP4
+ constexpr int MMA_K = 64;
  constexpr int MAX_GROUPS = 32;
   
- // L2 cache eviction policies
- constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;
  constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;
  constexpr uint64_t EVICT_LAST = 0x14F0000000000000;
   
  __device__ inline
  constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; }
   
- // Elect one thread in warp
  __device__
  uint32_t elect_sync() {
    uint32_t pred = 0;
  ⋯ 29 unchanged lines
    );
  }
   
- // steaming store for half2 (cache-steaming, evict-first policy)
- // Reduces L2 cache pollution from output writes, leaving more cache for input data reuse
  __device__ inline
  void store_steaming_half2(half* addr, half2 val) {
    uint32_t v = *reinterpret_cast<uint32_t*>(&val);
    asm volatile("st.global.cs.b32 [%0], %1;" :: "l"(addr), "r"(v) : "memory");
  }
   
- // 3D TMA for A/B matrices with tensor maps
  __device__ inline
  void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
    asm volatile("cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
  ⋯ 2 unchanged lines
                : "memory");
  }
   
- // 1D bulk copy for scale factors
  __device__ inline
  void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
    asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;"
                :: "r"(dst), "l"(src), "r"(size), "r"(mbar_addr), "l"(cache_policy));
  }
   
- // Scale factor copy: smem -> tmem
  __device__ inline
  void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
    asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
  }
   
- // NVFP4 MMA instruction with configurable output TMEM column
  __device__ inline
  void tcgen05_mma_nvfp4(
-   int d_tmem,
-   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, uint64_t a_desc, uint64_t b_desc, uint32_t i_desc,
+   int scale_A_tmem, int scale_B_tmem, int enable_input_d
  ) {
    asm volatile(
      "{\n\t"
  ⋯ 6 unchanged lines
    );
  }
   
- // TMEM load templates
  struct SHAPE {
    static constexpr char _32x32b[]  = ".32x32b";
    static constexpr char _16x128b[] = ".16x128b";
    static constexpr char _16x256b[] = ".16x256b";
  };
-  
  struct NUM {
    static constexpr char x4[]  = ".x4";
    static constexpr char x8[]  = ".x8";
  ⋯ 32 unchanged lines
  __device__ inline void tcgen05_ld_16x256bx4(float *tmp, int row, int col) { tcgen05_ld_16regs<SHAPE::_16x256b, NUM::x4>(tmp, row, col); }
  __device__ inline void tcgen05_ld_16x256bx8(float *tmp, int row, int col) { tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col); }
   
- // 64-register load for BLOCK_N=128
  template <const char *SHAPE, const char *NUM>
  __device__ inline
  void tcgen05_ld_64regs(float *tmp, int row, int col) {
  ⋯ 4 unchanged lines
    tcgen05_ld_64regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
  }
   
- // Host helper to check CU errors
  inline void check_cu(CUresult err) {
    if (err == CUDA_SUCCESS) return;
    const char *error_msg_ptr;
  ⋯ 2 unchanged lines
    printf("cuTensorMapEncodeTiled error: %s\n", error_msg_ptr);
  }
   
- // Initialize tensor map for A/B matrices (FP4 data) with L2 256B promotion
  inline void init_AB_tmap(
-   CUtensorMap *tmap,
-   const char *ptr,
+   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};  // in bytes
+   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,
+     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_L2_256B,
  ⋯ 2 unchanged lines
    check_cu(err);
  }
   
- // Kernel parameter structure with cumulative tile counts for fast lookup
  struct GroupGemmParams {
    CUtensorMap A_tmap;
    CUtensorMap B_tmap;
  ⋯ 1 unchanged lines
    const char* SFB_ptr;
    half* C_ptr;
    int M, N, K;
-   int tile_offset;  // Cumulative tile count before this group
-   int num_tiles;    // Number of tiles in this group
+   int tile_offset;
+   int num_tiles;
  };
   
- // Pre-computed tile info for O(1) lookup (replaces binary search)
- struct TileInfo {
-   uint8_t group_idx;   // Which group this tile belongs to
-   uint8_t bid_m;       // Block index in M dimension
-   uint16_t bid_n;      // Block index in N dimension
- };
-  
- // Binary search to find group index from tile ID (fallback for simple kernel)
  __device__ inline
  int find_group_idx(const GroupGemmParams* params, int tile_id, int num_groups) {
    int lo = 0, hi = num_groups - 1;
    while (lo < hi) {
      int mid = (lo + hi + 1) / 2;
-     if (params[mid].tile_offset <= tile_id)
-       lo = mid;
-     else
-       hi = mid - 1;
+     if (params[mid].tile_offset <= tile_id) lo = mid;
+     else hi = mid - 1;
    }
    return lo;
  }
   
- // Persistent group GEMM kernel with TMEM double-buffering
- template <
-   int BLOCK_M,
-   int BLOCK_N,
-   int BLOCK_K,
-   int NUM_STAGES
- >
- __global__
- __launch_bounds__(BLOCK_M + 2 * WARP_SIZE, 1)
+ template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
+ __global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE, 1)
  void group_gemm_persistent_kernel(
-   const GroupGemmParams* __restrict__ params,
-   const TileInfo* __restrict__ tile_info,
-   int num_groups,
-   int total_tiles
+   const GroupGemmParams* __restrict__ params, int num_groups, int total_tiles
  ) {
    const int bid = blockIdx.x;
    const int num_bids = gridDim.x;
    const int tid = threadIdx.x;
    const int lane_id = tid % WARP_SIZE;
    const int warp_id = tid / WARP_SIZE;
-  
    constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;
   
-   // Shared memory layout
    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;
  ⋯ 2 unchanged lines
    constexpr int SFB_size = 128 * BLOCK_K / 16;
    constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;
   
-   // Mbarriers: NUM_STAGES for TMA, NUM_STAGES for MMA, 2 for mainloop, 2 for epilogue
    #pragma nv_diag_suppress static_var_with_dynamic_init
    __shared__ int64_t mbars[NUM_STAGES * 2 + 4];
    const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
  ⋯ 1 unchanged lines
    const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
    const int epilogue_mbar_addr = mainloop_mbar_addr + 2 * 8;
   
-   // TMEM layout: [buffer0: BLOCK_N cols][buffer1: BLOCK_N cols][scale factors]
-   constexpr int SFA_tmem = BLOCK_N * 2;  // After both output double-buffers
+   constexpr int SFA_tmem = BLOCK_N * 2;
    constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
   
-   // Initialization
    if (warp_id == 0 && elect_sync()) {
      for (int i = 0; i < NUM_STAGES; i++) {
        mbarrier_init(tma_mbar_addr + i * 8, 1);
  ⋯ 1 unchanged lines
      }
      for (int i = 0; i < 2; i++) {
        mbarrier_init(mainloop_mbar_addr + i * 8, 1);
-       // Epilogue: BLOCK_M/WARP_SIZE warps each do elect_sync arrive
        mbarrier_init(epilogue_mbar_addr + i * 8, BLOCK_M / WARP_SIZE);
      }
      asm volatile("fence.mbarrier_init.release.cluster;");
    }
    else if (warp_id == 1) {
-     // Allocate TMEM: double-buffer (2x BLOCK_N columns)
      asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 2));
    }
    __syncthreads();
   
-   constexpr uint32_t i_desc = (1U << 7U)   // atype=E2M1
-                             | (1U << 10U)  // btype=E2M1
-                             | ((uint32_t)BLOCK_N >> 3U << 17U)  // MMA_N
-                             | ((uint32_t)128 >> 7U << 27U);     // MMA_M (always 128)
+   constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U)
+     | ((uint32_t)BLOCK_N >> 3U << 17U) | ((uint32_t)128 >> 7U << 27U);
   
-   // Warp specialization
    if (warp_id == NUM_WARPS - 2 && elect_sync()) {
-     // TMA warp - persistent grid-stride loop
      int tma_stage = 0;
-     int mma_phase = 1;  // After init, parity=0; wait(1) returns immediately
-  
+     int mma_phase = 1;
      for (int this_bid = bid; this_bid < total_tiles; this_bid += num_bids) {
-       // O(1) tile lookup instead of binary search
-       const TileInfo& ti = tile_info[this_bid];
-       const GroupGemmParams& p = params[ti.group_idx];
-       const int off_m = (int)ti.bid_m * BLOCK_M;
-       const int off_n = (int)ti.bid_n * BLOCK_N;
-       const int num_iters = p.K / BLOCK_K;
-  
+       const int group_idx = find_group_idx(params, this_bid, num_groups);
+       const GroupGemmParams& p = params[group_idx];
+       const int local_tile_id = this_bid - p.tile_offset;
+       const int M = p.M; const int K = p.K;
+       const int grid_m = (M + BLOCK_M - 1) / BLOCK_M;
+       const int bid_n = local_tile_id / grid_m;
+       const int bid_m = local_tile_id % grid_m;
+       const int off_m = bid_m * BLOCK_M;
+       const int off_n = bid_n * BLOCK_N;
+       const int num_iters = K / BLOCK_K;
        const CUtensorMap* A_tmap = &p.A_tmap;
        const CUtensorMap* B_tmap = &p.B_tmap;
        const char* SFA_ptr = p.SFA_ptr;
        const char* SFB_ptr = p.SFB_ptr;
-  
-       uint64_t cache_A = EVICT_LAST;
-       uint64_t cache_B = EVICT_LAST;
-  
+       uint64_t cache_A = EVICT_FIRST;
+       uint64_t cache_B = (K > 4096) ? EVICT_FIRST : EVICT_LAST;
        for (int iter_k = 0; iter_k < num_iters; iter_k++) {
-         // Wait for MMA to free this stage
          mbarrier_wait(mma_mbar_addr + tma_stage * 8, mma_phase);
-  
          const int mbar_addr = tma_mbar_addr + tma_stage * 8;
          const int A_smem = smem + tma_stage * STAGE_SIZE;
          const int B_smem = A_smem + A_size;
          const int SFA_smem = B_smem + B_size;
          const int SFB_smem = SFA_smem + SFA_size;
-  
          const int off_k = iter_k * BLOCK_K;
          tma_3d_gmem2smem(A_smem, A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
          tma_3d_gmem2smem(B_smem, B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
-  
-         const int rest_k = p.K / 16 / 4;
+         const int rest_k = K / 16 / 4;
          const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
          const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
          tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
          tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);
-  
          asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                      :: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");
-  
          tma_stage = (tma_stage + 1) % NUM_STAGES;
-         if (tma_stage == 0)
-           mma_phase ^= 1;
+         if (tma_stage == 0) mma_phase ^= 1;
        }
      }
    }
    else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
-     // MMA warp - persistent loop with TMEM double-buffering
-     int tma_stage = 0;
-     int tma_phase = 0;
-     int mainloop_stage = 0;  // 0 or 1 for TMEM double-buffer
-     int epilogue_phase = 1;  // After init, parity=0; wait(1) returns immediately
-  
+     int tma_stage = 0; int tma_phase = 0;
+     int mainloop_stage = 0; int epilogue_phase = 1;
      for (int this_bid = bid; this_bid < total_tiles; this_bid += num_bids) {
-       // O(1) tile lookup
-       const TileInfo& ti = tile_info[this_bid];
-       const GroupGemmParams& p = params[ti.group_idx];
-       const int bid_m = ti.bid_m;
-       const int bid_n = ti.bid_n;
-       const int num_iters = p.K / BLOCK_K;
+       const int group_idx = find_group_idx(params, this_bid, num_groups);
+       const GroupGemmParams& p = params[group_idx];
+       const int local_tile_id = this_bid - p.tile_offset;
+       const int M = p.M; const int K = p.K;
+       const int grid_m = (M + BLOCK_M - 1) / BLOCK_M;
+       const int bid_n = local_tile_id / grid_m;
+       const int bid_m = local_tile_id % grid_m;
+       const int num_iters = K / BLOCK_K;
        const int scale_A_offset = (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
        const int scale_B_offset = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
-  
-       // Wait for epilogue to finish with this TMEM buffer
        mbarrier_wait(epilogue_mbar_addr + mainloop_stage * 8, epilogue_phase);
-  
-       // MMA output goes to mainloop_stage * BLOCK_N in TMEM
        const int d_tmem = mainloop_stage * BLOCK_N;
-  
        for (int iter_k = 0; iter_k < num_iters; iter_k++) {
          mbarrier_wait(tma_mbar_addr + tma_stage * 8, tma_phase);
-  
          const int A_smem = smem + tma_stage * STAGE_SIZE;
          const int B_smem = A_smem + A_size;
          const int SFA_smem = B_smem + B_size;
          const int SFB_smem = SFA_smem + SFA_size;
-  
-         // Shared memory descriptors
          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);
  ⋯ 2 unchanged lines
            const int SBO = 8 * 16;
            return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
          };
-  
-         // Copy scale factors from SMEM to TMEM
          constexpr uint64_t SF_desc = make_desc_SF(0);
          const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
          const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);
-  
          #pragma unroll
          for (int k = 0; k < BLOCK_K / MMA_K; k++) {
            uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
  ⋯ 1 unchanged lines
            tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
            tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
          }
-  
-         // Execute MMA operations
          #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 * BLOCK_N * 128 + k2 * 32);
-  
              int k_sf = k1 * 4 + k2;
-             const int scale_A_tmem_addr = SFA_tmem + k_sf * 4 + scale_A_offset;
-             const int scale_B_tmem_addr = SFB_tmem + k_sf * 4 + scale_B_offset;
-  
              const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
-             tcgen05_mma_nvfp4(d_tmem, a_desc, b_desc, i_desc, scale_A_tmem_addr, scale_B_tmem_addr, enable_input_d);
+             tcgen05_mma_nvfp4(d_tmem, a_desc, b_desc, i_desc,
+               SFA_tmem + k_sf * 4 + scale_A_offset, SFB_tmem + k_sf * 4 + scale_B_offset, enable_input_d);
            }
-  
-         // Signal MMA done for this TMA stage
          asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                      :: "r"(mma_mbar_addr + tma_stage * 8) : "memory");
-  
          tma_stage = (tma_stage + 1) % NUM_STAGES;
-         if (tma_stage == 0)
-           tma_phase ^= 1;
+         if (tma_stage == 0) tma_phase ^= 1;
        }
-  
-       // Signal mainloop done for this tile (epilogue can start reading TMEM)
        asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                    :: "r"(mainloop_mbar_addr + mainloop_stage * 8) : "memory");
-  
-       // Flip TMEM double-buffer
        mainloop_stage = (mainloop_stage + 1) % 2;
-       if (mainloop_stage == 0)
-         epilogue_phase ^= 1;
+       if (mainloop_stage == 0) epilogue_phase ^= 1;
      }
    }
    else if (tid < BLOCK_M) {
-     // Epilogue warps - persistent loop with TMEM double-buffering
-     int mainloop_stage = 0;
-     int mainloop_phase = 0;
+     int mainloop_stage = 0; int mainloop_phase = 0;
      const int local_warp_id = tid / WARP_SIZE;
-  
      for (int this_bid = bid; this_bid < total_tiles; this_bid += num_bids) {
-       // Wait for MMA to finish this tile
        mbarrier_wait(mainloop_mbar_addr + mainloop_stage * 8, mainloop_phase);
        asm volatile("tcgen05.fence::after_thread_sync;");
-  
-       // O(1) tile lookup
-       const TileInfo& ti = tile_info[this_bid];
-       const GroupGemmParams& p = params[ti.group_idx];
-       const int M = p.M;
-       const int N = p.N;
-       const int off_m = (int)ti.bid_m * BLOCK_M;
-       const int off_n = (int)ti.bid_n * BLOCK_N;
+       const int group_idx = find_group_idx(params, this_bid, num_groups);
+       const GroupGemmParams& p = params[group_idx];
+       const int local_tile_id = this_bid - p.tile_offset;
+       const int M = p.M; const int N = p.N;
+       const int grid_m = (M + BLOCK_M - 1) / BLOCK_M;
+       const int bid_n = local_tile_id / grid_m;
+       const int bid_m = local_tile_id % grid_m;
+       const int off_m = bid_m * BLOCK_M;
+       const int off_n = bid_n * BLOCK_N;
        half* C_ptr = p.C_ptr;
-  
-       // TMEM column offset for this buffer
        const int tmem_col_offset = mainloop_stage * BLOCK_N;
-  
-       // N-major epilogue
        for (int m = 0; m < 32 / 16; m++) {
          float tmp[BLOCK_N / 2];
          if constexpr (BLOCK_N == 128) tcgen05_ld_16x256bx16(tmp, local_warp_id * 32 + m * 16, tmem_col_offset);
          else if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp, local_warp_id * 32 + m * 16, tmem_col_offset);
          else if constexpr (BLOCK_N == 32) tcgen05_ld_16x256bx4(tmp, local_warp_id * 32 + m * 16, tmem_col_offset);
          asm volatile("tcgen05.wait::ld.sync.aligned;");
-  
          const int row_base = off_m + local_warp_id * 32 + m * 16 + lane_id / 4;
          const int col_base = off_n + (lane_id % 4) * 2;
-  
          #pragma unroll
          for (int i = 0; i < BLOCK_N / 8; i++) {
-           const int row = row_base;
-           const int col = col_base + i * 8;
-  
+           const int row = row_base; const int col = col_base + i * 8;
            half2 val0 = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
            half2 val1 = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
-  
            if (row + 0 < M && col + 1 < N)
              store_steaming_half2(C_ptr + (row + 0) * N + col, val0);
            if (row + 8 < M && col + 1 < N)
              store_steaming_half2(C_ptr + (row + 8) * N + col, val1);
          }
        }
-  
-       // Signal epilogue done (MMA can reuse this TMEM buffer)
        if (elect_sync()) {
          asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];"
                      :: "r"(epilogue_mbar_addr + mainloop_stage * 8) : "memory");
        }
-  
-       // Flip TMEM double-buffer
        mainloop_stage = (mainloop_stage + 1) % 2;
-       if (mainloop_stage == 0)
-         mainloop_phase ^= 1;
+       if (mainloop_stage == 0) mainloop_phase ^= 1;
      }
    }
-  
    __syncthreads();
    if (warp_id == 0 && lane_id == 0)
      asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 2));
  }
   
- // Simple non-persistent kernel (single-buffered TMEM, no grid-stride)
- // Used for small workloads where total_tiles <= NUM_SMS (no benefit from persistence)
- template <
-   int BLOCK_M,
-   int BLOCK_N,
-   int BLOCK_K,
-   int NUM_STAGES
- >
- __global__
- __launch_bounds__(BLOCK_M + 2 * WARP_SIZE, 1)
+ template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
+ __global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE, 1)
  void group_gemm_simple_kernel(
-   const GroupGemmParams* __restrict__ params,
-   int num_groups,
-   int total_tiles
+   const GroupGemmParams* __restrict__ params, int num_groups, int total_tiles
  ) {
    const int bid = blockIdx.x;
    const int tid = threadIdx.x;
    const int lane_id = tid % WARP_SIZE;
    const int warp_id = tid / WARP_SIZE;
-  
    constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;
   
-   // Shared memory layout
    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;
  ⋯ 2 unchanged lines
    constexpr int SFB_size = 128 * BLOCK_K / 16;
    constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;
   
-   // Mbarriers: NUM_STAGES for TMA, NUM_STAGES for MMA (no mainloop/epilogue barriers)
    #pragma nv_diag_suppress static_var_with_dynamic_init
    __shared__ int64_t mbars[NUM_STAGES * 2];
    const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
    const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
-  
-   // TMEM layout: [buffer: BLOCK_N cols][scale factors]
    constexpr int SFA_tmem = BLOCK_N;
    constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
   
-   // Initialization
    if (warp_id == 0 && elect_sync()) {
      for (int i = 0; i < NUM_STAGES; i++) {
        mbarrier_init(tma_mbar_addr + i * 8, 1);
  ⋯ 2 unchanged lines
      asm volatile("fence.mbarrier_init.release.cluster;");
    }
    else if (warp_id == 1) {
-     // Allocate TMEM: single buffer (BLOCK_N columns)
      asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N));
    }
    __syncthreads();
   
-   constexpr uint32_t i_desc = (1U << 7U)
-                             | (1U << 10U)
-                             | ((uint32_t)BLOCK_N >> 3U << 17U)
-                             | ((uint32_t)128 >> 7U << 27U);
+   constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U)
+     | ((uint32_t)BLOCK_N >> 3U << 17U) | ((uint32_t)128 >> 7U << 27U);
   
-   // Find which group this tile belongs to
    const int group_idx = find_group_idx(params, bid, num_groups);
    const GroupGemmParams& p = params[group_idx];
    const int local_tile_id = bid - p.tile_offset;
-   const int M = p.M;
-   const int N = p.N;
-   const int K = p.K;
+   const int M = p.M; const int N = p.N; const int K = p.K;
    const int grid_m = (M + BLOCK_M - 1) / BLOCK_M;
    const int bid_n = local_tile_id / grid_m;
    const int bid_m = local_tile_id % grid_m;
  ⋯ 1 unchanged lines
    const int off_n = bid_n * BLOCK_N;
    const int num_iters = K / BLOCK_K;
   
-   // Warp specialization
    if (warp_id == NUM_WARPS - 2 && elect_sync()) {
-     // TMA warp
-     int tma_stage = 0;
-     int mma_phase = 1;
-  
+     int tma_stage = 0; int mma_phase = 1;
      const CUtensorMap* A_tmap = &p.A_tmap;
      const CUtensorMap* B_tmap = &p.B_tmap;
      const char* SFA_ptr = p.SFA_ptr;
      const char* SFB_ptr = p.SFB_ptr;
-  
-     uint64_t cache_A = EVICT_LAST;
-     uint64_t cache_B = EVICT_LAST;
-  
+     uint64_t cache_A = (M > N) ? EVICT_FIRST : EVICT_LAST;
+     uint64_t cache_B = (M > N) ? EVICT_LAST : EVICT_FIRST;
      for (int iter_k = 0; iter_k < num_iters; iter_k++) {
        mbarrier_wait(mma_mbar_addr + tma_stage * 8, mma_phase);
-  
        const int mbar_addr = tma_mbar_addr + tma_stage * 8;
        const int A_smem = smem + tma_stage * STAGE_SIZE;
        const int B_smem = A_smem + A_size;
        const int SFA_smem = B_smem + B_size;
        const int SFB_smem = SFA_smem + SFA_size;
-  
        const int off_k = iter_k * BLOCK_K;
        tma_3d_gmem2smem(A_smem, A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
        tma_3d_gmem2smem(B_smem, B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
-  
        const int rest_k = K / 16 / 4;
        const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
        const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
        tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
        tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);
-  
        asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                    :: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");
-  
        tma_stage = (tma_stage + 1) % NUM_STAGES;
-       if (tma_stage == 0)
-         mma_phase ^= 1;
+       if (tma_stage == 0) mma_phase ^= 1;
      }
    }
    else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
-     // MMA warp - single tile, single TMEM buffer
-     int tma_stage = 0;
-     int tma_phase = 0;
+     int tma_stage = 0; int tma_phase = 0;
      const int scale_A_offset = (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
      const int scale_B_offset = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
-     constexpr int d_tmem = 0;  // Single buffer at column 0
-  
+     constexpr int d_tmem = 0;
      for (int iter_k = 0; iter_k < num_iters; iter_k++) {
        mbarrier_wait(tma_mbar_addr + tma_stage * 8, tma_phase);
-  
        const int A_smem = smem + tma_stage * STAGE_SIZE;
        const int B_smem = A_smem + A_size;
        const int SFA_smem = B_smem + B_size;
        const int SFB_smem = SFA_smem + SFA_size;
-  
        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);
  ⋯ 2 unchanged lines
          const int SBO = 8 * 16;
          return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
        };
-  
        constexpr uint64_t SF_desc = make_desc_SF(0);
        const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
        const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);
-  
        #pragma unroll
        for (int k = 0; k < BLOCK_K / MMA_K; k++) {
-         uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
-         uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
-         tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
-         tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
+         tcgen05_cp_nvfp4(SFA_tmem + k * 4, SFA_desc + (uint64_t)k * (512ULL >> 4ULL));
+         tcgen05_cp_nvfp4(SFB_tmem + k * 4, SFB_desc + (uint64_t)k * (512ULL >> 4ULL));
        }
-  
        #pragma unroll
        for (int k1 = 0; k1 < BLOCK_K / 256; k1++)
          #pragma unroll
  ⋯ 1 unchanged lines
            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 * BLOCK_N * 128 + k2 * 32);
            int k_sf = k1 * 4 + k2;
-           const int scale_A_tmem_addr = SFA_tmem + k_sf * 4 + scale_A_offset;
-           const int scale_B_tmem_addr = SFB_tmem + k_sf * 4 + scale_B_offset;
            const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
-           tcgen05_mma_nvfp4(d_tmem, a_desc, b_desc, i_desc, scale_A_tmem_addr, scale_B_tmem_addr, enable_input_d);
+           tcgen05_mma_nvfp4(d_tmem, a_desc, b_desc, i_desc,
+             SFA_tmem + k_sf * 4 + scale_A_offset, SFB_tmem + k_sf * 4 + scale_B_offset, enable_input_d);
          }
-  
        asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                    :: "r"(mma_mbar_addr + tma_stage * 8) : "memory");
-  
        tma_stage = (tma_stage + 1) % NUM_STAGES;
-       if (tma_stage == 0)
-         tma_phase ^= 1;
+       if (tma_stage == 0) tma_phase ^= 1;
      }
-  
-     // Wait for all MMA to complete, then signal epilogue via syncthreads
      asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                  :: "r"(tma_mbar_addr) : "memory");
    }
-  
-   // All warps synchronize - MMA is done after this point
    __syncthreads();
    asm volatile("tcgen05.fence::after_thread_sync;");
   
-   // Epilogue: threads 0..BLOCK_M-1 write output
    if (tid < BLOCK_M) {
      const int local_warp_id = tid / WARP_SIZE;
      half* C_ptr = p.C_ptr;
-  
      for (int m = 0; m < 32 / 16; m++) {
        float tmp[BLOCK_N / 2];
        if constexpr (BLOCK_N == 128) tcgen05_ld_16x256bx16(tmp, local_warp_id * 32 + m * 16, 0);
        else if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp, local_warp_id * 32 + m * 16, 0);
        else if constexpr (BLOCK_N == 32) tcgen05_ld_16x256bx4(tmp, local_warp_id * 32 + m * 16, 0);
        asm volatile("tcgen05.wait::ld.sync.aligned;");
-  
        const int row_base = off_m + local_warp_id * 32 + m * 16 + lane_id / 4;
        const int col_base = off_n + (lane_id % 4) * 2;
-  
        #pragma unroll
        for (int i = 0; i < BLOCK_N / 8; i++) {
-         const int row = row_base;
-         const int col = col_base + i * 8;
-  
+         const int row = row_base; const int col = col_base + i * 8;
          half2 val0 = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
          half2 val1 = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
-  
          if (row + 0 < M && col + 1 < N)
            store_steaming_half2(C_ptr + (row + 0) * N + col, val0);
          if (row + 8 < M && col + 1 < N)
  ⋯ 1 unchanged lines
        }
      }
    }
-  
    __syncthreads();
    if (warp_id == 0 && lane_id == 0)
      asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N));
  }
   
- // Static device buffers to avoid repeated allocations
  static GroupGemmParams* d_params = nullptr;
  static size_t d_params_capacity = 0;
- static TileInfo* d_tile_info = nullptr;
- static size_t d_tile_info_capacity = 0;
+ static GroupGemmParams* h_params_pinned = nullptr;
   
- // Ensure device buffers are allocated
- inline void ensure_device_buffers(int num_groups, int total_tiles) {
+ inline void ensure_device_buffers(int num_groups) {
    if (d_params == nullptr || d_params_capacity < (size_t)num_groups) {
      if (d_params) cudaFree(d_params);
+     if (h_params_pinned) cudaFreeHost(h_params_pinned);
      d_params_capacity = num_groups + 8;
      cudaMalloc(&d_params, d_params_capacity * sizeof(GroupGemmParams));
+     cudaHostAlloc(&h_params_pinned, d_params_capacity * sizeof(GroupGemmParams), cudaHostAllocDefault);
    }
-   if (d_tile_info == nullptr || d_tile_info_capacity < (size_t)total_tiles) {
-     if (d_tile_info) cudaFree(d_tile_info);
-     d_tile_info_capacity = total_tiles + 64;
-     cudaMalloc(&d_tile_info, d_tile_info_capacity * sizeof(TileInfo));
-   }
  }
   
- // CUDA Graph cache for fast replay when same parameters are used
  struct GraphCacheEntry {
    cudaGraphExec_t graph_exec;
    int num_groups;
    int total_tiles;
-   // Cache the input pointers for fast comparison
    const char* A_ptrs[MAX_GROUPS];
-   const char* B_ptrs[MAX_GROUPS];
-   const char* SFA_ptrs[MAX_GROUPS];
-   const char* SFB_ptrs[MAX_GROUPS];
    c10::Half* C_ptrs[MAX_GROUPS];
    int M_sizes[MAX_GROUPS];
-   int N_sizes[MAX_GROUPS];
-   int K_sizes[MAX_GROUPS];
    bool valid;
  };
   
  ⋯ 1 unchanged lines
  static GraphCacheEntry g_graph_cache[MAX_GRAPH_CACHE];
  static int g_graph_cache_idx = 0;
   
- // Compute a fast hash for graph cache lookup
- inline uint64_t compute_graph_hash(
-   const char* const* A_ptrs,
-   const char* const* B_ptrs,
-   c10::Half* const* C_ptrs,
-   int num_groups
- ) {
-   uint64_t h = (uint64_t)num_groups * 0x9e3779b97f4a7c15ULL;
-   for (int g = 0; g < num_groups; g++) {
-     h ^= (uint64_t)(uintptr_t)A_ptrs[g] * 0x517cc1b727220a95ULL;
-     h ^= (uint64_t)(uintptr_t)B_ptrs[g] * 0x6c62272e07bb0142ULL;
-     h ^= (uint64_t)(uintptr_t)C_ptrs[g] * 0x48b2197db17f3848ULL;
-   }
+ #include <unordered_map>
+ static std::unordered_map<uint64_t, int> g_graph_hash_map;
+  
+ inline uint64_t compute_graph_hash(const char* A0, c10::Half* C0, int ng) {
+   uint64_t h = 14695981039346656037ULL;
+   h ^= (uint64_t)ng; h *= 1099511628211ULL;
+   h ^= (uint64_t)(uintptr_t)A0; h *= 1099511628211ULL;
+   h ^= (uint64_t)(uintptr_t)C0; h *= 1099511628211ULL;
    return h;
  }
   
- // Hash table for fast graph cache lookup
- static uint64_t g_graph_hashes[MAX_GRAPH_CACHE];
+ // Returns graph_exec handle as int64 for direct ctypes launch
+ int64_t group_gemm_packed(torch::Tensor packed_args, int64_t num_groups) {
+     const int ng = static_cast<int>(num_groups);
+     const int64_t* data = packed_args.data_ptr<int64_t>();
   
- // Find cached graph that matches the input parameters exactly
- inline int find_matching_graph(
-   const char* const* A_ptrs,
-   const char* const* B_ptrs,
-   const char* const* SFA_ptrs,
-   const char* const* SFB_ptrs,
-   c10::Half* const* C_ptrs,
-   const int* M_sizes,
-   const int* N_sizes,
-   const int* K_sizes,
-   int num_groups
- ) {
-   uint64_t target_hash = compute_graph_hash(A_ptrs, B_ptrs, C_ptrs, num_groups);
+     const char* a_ptrs[MAX_GROUPS]; const char* b_ptrs[MAX_GROUPS];
+     const char* sfa_ptrs[MAX_GROUPS]; const char* sfb_ptrs[MAX_GROUPS];
+     c10::Half* c_ptrs[MAX_GROUPS];
+     int m_sizes[MAX_GROUPS]; int n_sizes[MAX_GROUPS]; int k_sizes[MAX_GROUPS];
   
-   for (int i = 0; i < MAX_GRAPH_CACHE; i++) {
-     if (!g_graph_cache[i].valid || g_graph_hashes[i] != target_hash)
-       continue;
+     for (int g = 0; g < ng; ++g) {
+         a_ptrs[g] = reinterpret_cast<const char*>(data[g]);
+         b_ptrs[g] = reinterpret_cast<const char*>(data[ng + g]);
+         c_ptrs[g] = reinterpret_cast<c10::Half*>(data[2*ng + g]);
+         sfa_ptrs[g] = reinterpret_cast<const char*>(data[3*ng + g]);
+         sfb_ptrs[g] = reinterpret_cast<const char*>(data[4*ng + g]);
+         m_sizes[g] = static_cast<int>(data[5*ng + g]);
+         n_sizes[g] = static_cast<int>(data[6*ng + g]);
+         k_sizes[g] = static_cast<int>(data[7*ng + g]);
+     }
   
-     // Hash matched, verify exact match
-     if (g_graph_cache[i].num_groups != num_groups)
-       continue;
+     constexpr int BLOCK_M = 128, BLOCK_N = 128, BLOCK_K = 256, NUM_STAGES = 6, NUM_SMS = 148;
   
-     bool match = true;
-     for (int g = 0; g < num_groups && match; g++) {
-       if (g_graph_cache[i].A_ptrs[g] != A_ptrs[g] ||
-           g_graph_cache[i].B_ptrs[g] != B_ptrs[g] ||
-           g_graph_cache[i].SFA_ptrs[g] != SFA_ptrs[g] ||
-           g_graph_cache[i].SFB_ptrs[g] != SFB_ptrs[g] ||
-           g_graph_cache[i].C_ptrs[g] != C_ptrs[g] ||
-           g_graph_cache[i].M_sizes[g] != M_sizes[g] ||
-           g_graph_cache[i].N_sizes[g] != N_sizes[g] ||
-           g_graph_cache[i].K_sizes[g] != K_sizes[g]) {
-         match = false;
-       }
+     // O(1) hash lookup for cache hit
+     uint64_t target_hash = compute_graph_hash(a_ptrs[0], c_ptrs[0], ng);
+     auto it = g_graph_hash_map.find(target_hash);
+     if (it != g_graph_hash_map.end()) {
+         int i = it->second;
+         if (g_graph_cache[i].valid && g_graph_cache[i].num_groups == ng &&
+             g_graph_cache[i].A_ptrs[0] == a_ptrs[0] && g_graph_cache[i].C_ptrs[0] == c_ptrs[0]) {
+             cudaGraphLaunch(g_graph_cache[i].graph_exec, 0);
+             return (int64_t)(uintptr_t)g_graph_cache[i].graph_exec;
+         }
      }
-     if (match) return i;
-   }
-   return -1;
- }
   
- // Host launch function for Group GEMM with CUDA Graph support
- void launch_gpu_implementation(
-   const char* const* A_ptrs,
-   const char* const* B_ptrs,
-   const char* const* SFA_ptrs,
-   const char* const* SFB_ptrs,
-   const int* M_sizes,
-   const int* N_sizes,
-   const int* K_sizes,
-   c10::Half* const* C_ptrs,
-   int num_groups
- ) {
-   constexpr int BLOCK_M = 128;
-   constexpr int BLOCK_N = 128;
-   constexpr int BLOCK_K = 256;
-   constexpr int NUM_STAGES = 6;
-   constexpr int NUM_SMS = 148;
+     ensure_device_buffers(ng);
+     int total_tiles = 0;
+     for (int g = 0; g < ng; g++) {
+         init_AB_tmap(&h_params_pinned[g].A_tmap, a_ptrs[g], m_sizes[g], k_sizes[g], BLOCK_M, BLOCK_K);
+         init_AB_tmap(&h_params_pinned[g].B_tmap, b_ptrs[g], n_sizes[g], k_sizes[g], BLOCK_N, BLOCK_K);
+         h_params_pinned[g].SFA_ptr = sfa_ptrs[g];
+         h_params_pinned[g].SFB_ptr = sfb_ptrs[g];
+         h_params_pinned[g].C_ptr = reinterpret_cast<half*>(const_cast<c10::Half*>(c_ptrs[g]));
+         h_params_pinned[g].M = m_sizes[g];
+         h_params_pinned[g].N = n_sizes[g];
+         h_params_pinned[g].K = k_sizes[g];
+         h_params_pinned[g].tile_offset = total_tiles;
+         int grid_m = (m_sizes[g] + BLOCK_M - 1) / BLOCK_M;
+         int grid_n = (n_sizes[g] + BLOCK_N - 1) / BLOCK_N;
+         h_params_pinned[g].num_tiles = grid_m * grid_n;
+         total_tiles += h_params_pinned[g].num_tiles;
+     }
   
-   // Check if we have a cached graph for these exact parameters
-   int cache_idx = find_matching_graph(A_ptrs, B_ptrs, SFA_ptrs, SFB_ptrs, C_ptrs,
-                                        M_sizes, N_sizes, K_sizes, num_groups);
+     int tb_size = BLOCK_M + 2 * WARP_SIZE;
+     int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
+     int SFAB_size = 128 * (BLOCK_K / 16) * 2;
+     int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
+     bool use_persistent = (total_tiles > NUM_SMS);
+     int grid_size = use_persistent ? NUM_SMS : total_tiles;
   
-   if (cache_idx >= 0) {
-     // Fast path: replay cached graph
-     cudaGraphLaunch(g_graph_cache[cache_idx].graph_exec, 0);
-     return;
-   }
+     auto persistent_kernel = group_gemm_persistent_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
+     auto simple_kernel = group_gemm_simple_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
   
-   // Slow path: prepare parameters and capture/launch
-   GroupGemmParams h_params[MAX_GROUPS];
-   int total_tiles = 0;
-   for (int g = 0; g < num_groups; g++) {
-     init_AB_tmap(&h_params[g].A_tmap, A_ptrs[g], M_sizes[g], K_sizes[g], BLOCK_M, BLOCK_K);
-     init_AB_tmap(&h_params[g].B_tmap, B_ptrs[g], N_sizes[g], K_sizes[g], BLOCK_N, BLOCK_K);
-     h_params[g].SFA_ptr = SFA_ptrs[g];
-     h_params[g].SFB_ptr = SFB_ptrs[g];
-     h_params[g].C_ptr = reinterpret_cast<half*>(const_cast<c10::Half*>(C_ptrs[g]));
-     h_params[g].M = M_sizes[g];
-     h_params[g].N = N_sizes[g];
-     h_params[g].K = K_sizes[g];
-     h_params[g].tile_offset = total_tiles;
-  
-     int grid_m = (M_sizes[g] + BLOCK_M - 1) / BLOCK_M;
-     int grid_n = (N_sizes[g] + BLOCK_N - 1) / BLOCK_N;
-     h_params[g].num_tiles = grid_m * grid_n;
-     total_tiles += h_params[g].num_tiles;
-   }
-  
-   ensure_device_buffers(num_groups, total_tiles);
-  
-   // Build pre-computed tile info table (O(1) lookup replaces binary search)
-   std::vector<TileInfo> h_tile_info(total_tiles);
-   for (int g = 0; g < num_groups; g++) {
-     int grid_m = (M_sizes[g] + BLOCK_M - 1) / BLOCK_M;
-     int grid_n = (N_sizes[g] + BLOCK_N - 1) / BLOCK_N;
-     for (int t = 0; t < h_params[g].num_tiles; t++) {
-       int tile_id = h_params[g].tile_offset + t;
-       h_tile_info[tile_id].group_idx = (uint8_t)g;
-       h_tile_info[tile_id].bid_m = (uint8_t)(t % grid_m);
-       h_tile_info[tile_id].bid_n = (uint16_t)(t / grid_m);
+     static bool smem_configured = false;
+     if (!smem_configured && smem_size > 48'000) {
+         cudaFuncSetAttribute(persistent_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+         cudaFuncSetAttribute(simple_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+         smem_configured = true;
      }
-   }
   
-   int tb_size = BLOCK_M + 2 * WARP_SIZE;
-   int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
-   int SFAB_size = 128 * (BLOCK_K / 16) * 2;
-   int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
+     auto this_kernel = use_persistent ? persistent_kernel : simple_kernel;
+     cudaMemcpy(d_params, h_params_pinned, ng * sizeof(GroupGemmParams), cudaMemcpyHostToDevice);
   
-   // Choose kernel: persistent for large workloads, simple for small
-   bool use_persistent = (total_tiles > NUM_SMS);
-   int grid_size = use_persistent ? NUM_SMS : total_tiles;
-  
-   auto persistent_kernel = group_gemm_persistent_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
-   auto simple_kernel = group_gemm_simple_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
-  
-   static bool smem_configured = false;
-   if (!smem_configured && smem_size > 48'000) {
-     cudaFuncSetAttribute(persistent_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-     cudaFuncSetAttribute(simple_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-     smem_configured = true;
-   }
-  
-   // Copy params and tile info to device
-   cudaMemcpy(d_params, h_params, num_groups * sizeof(GroupGemmParams), cudaMemcpyHostToDevice);
-   if (use_persistent) {
-     cudaMemcpy(d_tile_info, h_tile_info.data(), total_tiles * sizeof(TileInfo), cudaMemcpyHostToDevice);
-   }
-  
-   // Capture graph using explicit API
-   cudaGraph_t graph;
-   cudaGraphCreate(&graph, 0);
-  
-   cudaGraphNode_t kernel_node;
-   cudaKernelNodeParams kernel_params = {0};
-   // Use static storage for kernel args so pointers remain valid
-   static int s_num_groups, s_total_tiles;
-   s_num_groups = num_groups;
-   s_total_tiles = total_tiles;
-  
-   if (use_persistent) {
-     // Persistent kernel takes tile_info
-     static void* persistent_args[4];
-     persistent_args[0] = &d_params;
-     persistent_args[1] = &d_tile_info;
-     persistent_args[2] = (void*)&s_num_groups;
-     persistent_args[3] = (void*)&s_total_tiles;
-     kernel_params.func = (void*)persistent_kernel;
+     cudaGraph_t graph;
+     cudaGraphCreate(&graph, 0);
+     cudaGraphNode_t kernel_node;
+     cudaKernelNodeParams kernel_params = {0};
+     void* kernel_args[] = { &d_params, (void*)&ng, (void*)&total_tiles };
+     kernel_params.func = (void*)this_kernel;
      kernel_params.gridDim = dim3(grid_size);
      kernel_params.blockDim = dim3(tb_size);
      kernel_params.sharedMemBytes = smem_size;
-     kernel_params.kernelParams = persistent_args;
+     kernel_params.kernelParams = kernel_args;
      kernel_params.extra = nullptr;
-   } else {
-     // Simple kernel uses binary search (no tile_info needed)
-     static void* simple_args[3];
-     simple_args[0] = &d_params;
-     simple_args[1] = (void*)&s_num_groups;
-     simple_args[2] = (void*)&s_total_tiles;
-     kernel_params.func = (void*)simple_kernel;
-     kernel_params.gridDim = dim3(grid_size);
-     kernel_params.blockDim = dim3(tb_size);
-     kernel_params.sharedMemBytes = smem_size;
-     kernel_params.kernelParams = simple_args;
-     kernel_params.extra = nullptr;
-   }
-   cudaGraphAddKernelNode(&kernel_node, graph, nullptr, 0, &kernel_params);
+     cudaGraphAddKernelNode(&kernel_node, graph, nullptr, 0, &kernel_params);
   
-   // Instantiate and execute
-   cudaGraphExec_t graph_exec;
-   cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0);
-   cudaGraphLaunch(graph_exec, 0);
+     cudaGraphExec_t graph_exec;
+     cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0);
+     cudaGraphLaunch(graph_exec, 0);
   
-   // Cache for future reuse
-   int new_idx = g_graph_cache_idx % MAX_GRAPH_CACHE;
-   if (g_graph_cache[new_idx].valid) {
-     cudaGraphExecDestroy(g_graph_cache[new_idx].graph_exec);
-   }
-  
-   g_graph_cache[new_idx].graph_exec = graph_exec;
-   g_graph_cache[new_idx].num_groups = num_groups;
-   g_graph_cache[new_idx].total_tiles = total_tiles;
-   g_graph_cache[new_idx].valid = true;
-   g_graph_hashes[new_idx] = compute_graph_hash(A_ptrs, B_ptrs, C_ptrs, num_groups);
-  
-   for (int g = 0; g < num_groups; g++) {
-     g_graph_cache[new_idx].A_ptrs[g] = A_ptrs[g];
-     g_graph_cache[new_idx].B_ptrs[g] = B_ptrs[g];
-     g_graph_cache[new_idx].SFA_ptrs[g] = SFA_ptrs[g];
-     g_graph_cache[new_idx].SFB_ptrs[g] = SFB_ptrs[g];
-     g_graph_cache[new_idx].C_ptrs[g] = C_ptrs[g];
-     g_graph_cache[new_idx].M_sizes[g] = M_sizes[g];
-     g_graph_cache[new_idx].N_sizes[g] = N_sizes[g];
-     g_graph_cache[new_idx].K_sizes[g] = K_sizes[g];
-   }
-  
-   g_graph_cache_idx++;
-   cudaGraphDestroy(graph);
+     int new_idx = g_graph_cache_idx % MAX_GRAPH_CACHE;
+     if (g_graph_cache[new_idx].valid) {
+         cudaGraphExecDestroy(g_graph_cache[new_idx].graph_exec);
+         uint64_t old_hash = compute_graph_hash(g_graph_cache[new_idx].A_ptrs[0],
+             g_graph_cache[new_idx].C_ptrs[0], g_graph_cache[new_idx].num_groups);
+         g_graph_hash_map.erase(old_hash);
+     }
+     g_graph_cache[new_idx].graph_exec = graph_exec;
+     g_graph_cache[new_idx].num_groups = ng;
+     g_graph_cache[new_idx].total_tiles = total_tiles;
+     g_graph_cache[new_idx].valid = true;
+     for (int g = 0; g < ng; g++) {
+         g_graph_cache[new_idx].A_ptrs[g] = a_ptrs[g];
+         g_graph_cache[new_idx].C_ptrs[g] = c_ptrs[g];
+         g_graph_cache[new_idx].M_sizes[g] = m_sizes[g];
+     }
+     g_graph_hash_map[target_hash] = new_idx;
+     g_graph_cache_idx++;
+     cudaGraphDestroy(graph);
+     return (int64_t)(uintptr_t)graph_exec;
  }
   
-  
  #include <torch/extension.h>
- #include <cuda_fp16.h>
-  
- // Maximum supported groups for stack allocation
- #define MAX_GROUPS 32
-  
- // Simple interface - Python handles caching
- void group_gemm_cuda(
-     std::vector<int64_t> A_ptrs,
-     std::vector<int64_t> B_ptrs,
-     std::vector<int64_t> C_ptrs,
-     std::vector<int64_t> SFA_ptrs,
-     std::vector<int64_t> SFB_ptrs,
-     std::vector<int64_t> M_sizes,
-     std::vector<int64_t> N_sizes,
-     std::vector<int64_t> K_sizes,
-     int64_t num_groups
- ) {
-     const int ng = static_cast<int>(num_groups);
-  
-     // Convert to kernel-expected types
-     const char* a_ptrs[MAX_GROUPS];
-     const char* b_ptrs[MAX_GROUPS];
-     const char* sfa_ptrs[MAX_GROUPS];
-     const char* sfb_ptrs[MAX_GROUPS];
-     c10::Half* c_ptrs[MAX_GROUPS];
-     int m_sizes[MAX_GROUPS];
-     int n_sizes[MAX_GROUPS];
-     int k_sizes[MAX_GROUPS];
-  
-     for (int g = 0; g < ng; ++g) {
-         a_ptrs[g] = reinterpret_cast<const char*>(A_ptrs[g]);
-         b_ptrs[g] = reinterpret_cast<const char*>(B_ptrs[g]);
-         c_ptrs[g] = reinterpret_cast<c10::Half*>(C_ptrs[g]);
-         sfa_ptrs[g] = reinterpret_cast<const char*>(SFA_ptrs[g]);
-         sfb_ptrs[g] = reinterpret_cast<const char*>(SFB_ptrs[g]);
-         m_sizes[g] = static_cast<int>(M_sizes[g]);
-         n_sizes[g] = static_cast<int>(N_sizes[g]);
-         k_sizes[g] = static_cast<int>(K_sizes[g]);
-     }
-  
-     launch_gpu_implementation(
-         a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs,
-         m_sizes, n_sizes, k_sizes,
-         c_ptrs, ng
-     );
- }
  """
   
- # C++ header declarations
  cpp_source = """
  #include <torch/extension.h>
- #include <vector>
-  
- void group_gemm_cuda(
-     std::vector<int64_t> A_ptrs,
-     std::vector<int64_t> B_ptrs,
-     std::vector<int64_t> C_ptrs,
-     std::vector<int64_t> SFA_ptrs,
-     std::vector<int64_t> SFB_ptrs,
-     std::vector<int64_t> M_sizes,
-     std::vector<int64_t> N_sizes,
-     std::vector<int64_t> K_sizes,
-     int64_t num_groups
- );
+ int64_t group_gemm_packed(torch::Tensor packed_args, int64_t num_groups);
  """
   
- # Load with direct function binding (bypasses torch.ops dispatcher)
  _module = load_inline(
-     name='group_gemm_fp4',
+     name='group_gemm_fp4_v10',
      cpp_sources=cpp_source,
      cuda_sources=cuda_source,
-     functions=['group_gemm_cuda'],
-     verbose=True,
+     functions=['group_gemm_packed'],
+     verbose=False,
      extra_cuda_cflags=[
          "-O3",
          "-gencode=arch=compute_100a,code=sm_100a",
          "--use_fast_math",
          "--expt-relaxed-constexpr",
          "--relocatable-device-code=false",
-         "-lineinfo",
-         "-Xptxas=-v",
  ⋯ diff truncated