submission 496386

  ⋯ 2 unchanged lines
  from task import input_t, output_t
  from torch.utils.cpp_extension import load_inline
   
- _libcudart = ctypes.CDLL("libcudart.so")
- _libcudart.cudaGraphLaunch.restype = ctypes.c_int
- _libcudart.cudaGraphLaunch.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
- _ZERO = ctypes.c_void_p(0)
-  
  cuda_source = r"""
  #include <cudaTypedefs.h>
  #include <cuda_fp16.h>
  #include <c10/util/Half.h>
+ #include <cstring>
   
  constexpr int WARP_SIZE = 32;
  constexpr int MMA_K = 64;
- constexpr int MAX_GROUPS = 32;
- constexpr int MAX_TILES = 1024;
- constexpr int MAX_BATCH_SLOTS = 32;
+ constexpr int MAX_GROUPS = 8;
+ constexpr int MAX_TILES = 768;
   
  constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;
  constexpr uint64_t EVICT_LAST = 0x14F0000000000000;
  ⋯ 44 unchanged lines
  }
   
  __device__ inline
+ void store_cg_u64(void* addr, uint64_t val) {
+   asm volatile("st.global.cg.b64 [%0], %1;" :: "l"(addr), "l"(val) : "memory");
+ }
+ __device__ inline
+ void store_cg_u32(void* addr, uint32_t val) {
+   asm volatile("st.global.cg.b32 [%0], %1;" :: "l"(addr), "r"(val) : "memory");
+ }
+  
+ __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 "
                "[%0], [%1, {%2, %3, %4}], [%5], %6;"
  ⋯ 110 unchanged lines
    check_cu(err);
  }
   
- struct GroupGemmParams {
-   CUtensorMap A_tmap;
-   CUtensorMap B_tmap;
+ constexpr int TMAP_ADDR_BYTE = 0;
+ constexpr int TMAP_DIM1_BYTE = 36;
+  
+ __device__ CUtensorMap d_A_tmaps[MAX_GROUPS];
+ __device__ CUtensorMap d_B_tmaps[MAX_GROUPS];
+  
+ struct GroupDynArgs {
+   uint64_t a_addr;
+   uint64_t b_addr;
    const char* SFA_ptr;
    const char* SFB_ptr;
    half* C_ptr;
-   int M, N, K;
-   int num_tiles;
+   int M;
+   int N;
    uint64_t cache_A;
    uint64_t cache_B;
  };
   
- struct __align__(4) TileInfo {
-   int8_t group_idx;
-   int8_t bid_m;
-   int8_t bid_n;
-   int8_t _pad;
+ // Packed 2-byte TileInfo: group_idx(3) + bid_m(2) + bid_n(6) + num_k_iters(5) = 16 bits
+ // Halves tile table size from 3072 to 1536 bytes for 768 tiles
+ __device__ __host__ inline
+ void pack_tile(uint16_t &packed, int group_idx, int bid_m, int bid_n, int num_k_iters) {
+   packed = (uint16_t)((group_idx & 0x7) | ((bid_m & 0x3) << 3) | ((bid_n & 0x3F) << 5) | ((num_k_iters & 0x1F) << 11));
+ }
+  
+ __device__ inline int tile_group(uint16_t t) { return t & 0x7; }
+ __device__ inline int tile_bid_m(uint16_t t) { return (t >> 3) & 0x3; }
+ __device__ inline int tile_bid_n(uint16_t t) { return (t >> 5) & 0x3F; }
+ __device__ inline int tile_num_k_iters(uint16_t t) { return (t >> 11) & 0x1F; }
+  
+ struct KernelParams {
+   GroupDynArgs groups[MAX_GROUPS];
+   uint16_t tiles[MAX_TILES];
+   int ng;
+   int total_tiles;
  };
   
- struct TileLookup {
-   TileInfo tiles[MAX_TILES];
+ // Compact params for simple kernel with small tile table
+ // Max 148 tiles (since simple kernel means total_tiles <= NUM_SMS=148)
+ constexpr int MAX_SIMPLE_TILES = 148;
+ struct SimpleKernelParams {
+   GroupDynArgs groups[MAX_GROUPS];
+   int ng;
+   int total_tiles;
+   uint16_t tiles[MAX_SIMPLE_TILES];
  };
   
  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 TileLookup* __restrict__ lookup, int total_tiles
- ) {
+ void group_gemm_persistent_kernel(KernelParams kp) {
    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;
-  
    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;
  ⋯ 1 unchanged lines
    constexpr int SFA_size = 128 * BLOCK_K / 16;
    constexpr int SFB_size = 128 * BLOCK_K / 16;
    constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;
-  
    #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));
    const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
    const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
    const int epilogue_mbar_addr = mainloop_mbar_addr + 2 * 8;
-  
    constexpr int SFA_tmem = BLOCK_N * 2;
    constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
-  
    if (warp_id == 0 && elect_sync()) {
      for (int i = 0; i < NUM_STAGES; i++) {
        mbarrier_init(tma_mbar_addr + i * 8, 1);
  ⋯ 10 unchanged lines
    }
    __syncthreads();
   
+   if (warp_id == 0 && elect_sync()) {
+     for (int g = 0; g < kp.ng; g++) {
+       const GroupDynArgs& gp = kp.groups[g];
+       store_cg_u64((char*)&d_A_tmaps[g] + TMAP_ADDR_BYTE, gp.a_addr);
+       store_cg_u32((char*)&d_A_tmaps[g] + TMAP_DIM1_BYTE, (uint32_t)(gp.M - 1));
+       store_cg_u64((char*)&d_B_tmaps[g] + TMAP_ADDR_BYTE, gp.b_addr);
+     }
+   }
+   __syncthreads();
+  
    constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U)
      | ((uint32_t)BLOCK_N >> 3U << 17U) | ((uint32_t)128 >> 7U << 27U);
   
    if (warp_id == NUM_WARPS - 2 && elect_sync()) {
-     int tma_stage = 0;
-     int mma_phase = 1;
-     for (int this_bid = bid; this_bid < total_tiles; this_bid += num_bids) {
-       const TileInfo ti = lookup->tiles[this_bid];
-       const GroupGemmParams& p = params[ti.group_idx];
-       const int K = p.K;
-       const int off_m = (int)ti.bid_m * BLOCK_M;
-       const int off_n = (int)ti.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 = p.cache_A;
-       uint64_t cache_B = p.cache_B;
+     int tma_stage = 0; int mma_phase = 1;
+     for (int this_bid = bid; this_bid < kp.total_tiles; this_bid += num_bids) {
+       const uint16_t ti = kp.tiles[this_bid];
+       const int g = tile_group(ti);
+       const GroupDynArgs& gp = kp.groups[g];
+       const int bid_m = tile_bid_m(ti);
+       const int bid_n = tile_bid_n(ti);
+       const int num_iters = tile_num_k_iters(ti);
+       const int off_m = bid_m * BLOCK_M;
+       const int off_n = bid_n * BLOCK_N;
+       const CUtensorMap* A_tmap = &d_A_tmaps[g];
+       const CUtensorMap* B_tmap = &d_B_tmaps[g];
+       const char* SFA_ptr = gp.SFA_ptr;
+       const char* SFB_ptr = gp.SFB_ptr;
+       uint64_t cache_A = gp.cache_A;
+       uint64_t cache_B = gp.cache_B;
+       const int rest_k = num_iters * (BLOCK_K / 64);
        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;
  ⋯ 4 unchanged lines
          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);
+         const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / 64) * 512;
+         const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / 64) * 512;
+         // Scale factors are tiny (~2KB each), always keep in L2
+         tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, EVICT_LAST);
+         tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, EVICT_LAST);
          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;
  ⋯ 4 unchanged lines
    else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
      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) {
-       const TileInfo ti = lookup->tiles[this_bid];
-       const GroupGemmParams& p = params[ti.group_idx];
-       const int K = p.K;
-       const int num_iters = K / BLOCK_K;
-       const int scale_A_offset = ((int)ti.bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
-       const int scale_B_offset = ((int)ti.bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
+     for (int this_bid = bid; this_bid < kp.total_tiles; this_bid += num_bids) {
+       const uint16_t ti = kp.tiles[this_bid];
+       const int g = tile_group(ti);
+       const GroupDynArgs& gp = kp.groups[g];
+       const int bid_m = tile_bid_m(ti);
+       const int bid_n = tile_bid_n(ti);
+       const int num_iters = tile_num_k_iters(ti);
+       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);
        mbarrier_wait(epilogue_mbar_addr + mainloop_stage * 8, epilogue_phase);
        const int d_tmem = mainloop_stage * BLOCK_N;
        for (int iter_k = 0; iter_k < num_iters; iter_k++) {
  ⋯ 15 unchanged lines
          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++)
  ⋯ 20 unchanged lines
    else if (tid < BLOCK_M) {
      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) {
+     for (int this_bid = bid; this_bid < kp.total_tiles; this_bid += num_bids) {
        mbarrier_wait(mainloop_mbar_addr + mainloop_stage * 8, mainloop_phase);
        asm volatile("tcgen05.fence::after_thread_sync;");
-       const TileInfo ti = lookup->tiles[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;
-       half* C_ptr = p.C_ptr;
+       const uint16_t ti = kp.tiles[this_bid];
+       const int g = tile_group(ti);
+       const GroupDynArgs& gp = kp.groups[g];
+       const int bid_m = tile_bid_m(ti);
+       const int bid_n = tile_bid_n(ti);
+       const int M = gp.M; const int N = gp.N;
+       const int off_m = bid_m * BLOCK_M;
+       const int off_n = bid_n * BLOCK_N;
+       half* C_ptr = gp.C_ptr;
        const int tmem_col_offset = mainloop_stage * BLOCK_N;
        for (int m = 0; m < 32 / 16; m++) {
          float tmp[BLOCK_N / 2];
  ⋯ 27 unchanged lines
      asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 2));
  }
   
+ // ========== Simple kernel with compact SimpleKernelParams ==========
  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, const TileLookup* __restrict__ lookup, int total_tiles
- ) {
+ void group_gemm_simple_kernel(SimpleKernelParams kp) {
    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;
-  
    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;
  ⋯ 1 unchanged lines
    constexpr int SFA_size = 128 * BLOCK_K / 16;
    constexpr int SFB_size = 128 * BLOCK_K / 16;
    constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;
-  
    #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;
    constexpr int SFA_tmem = BLOCK_N;
    constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
-  
    if (warp_id == 0 && elect_sync()) {
      for (int i = 0; i < NUM_STAGES; i++) {
        mbarrier_init(tma_mbar_addr + i * 8, 1);
  ⋯ 6 unchanged lines
    }
    __syncthreads();
   
+   if (warp_id == 0 && elect_sync()) {
+     for (int gg = 0; gg < kp.ng; gg++) {
+       const GroupDynArgs& gp2 = kp.groups[gg];
+       store_cg_u64((char*)&d_A_tmaps[gg] + TMAP_ADDR_BYTE, gp2.a_addr);
+       store_cg_u32((char*)&d_A_tmaps[gg] + TMAP_DIM1_BYTE, (uint32_t)(gp2.M - 1));
+       store_cg_u64((char*)&d_B_tmaps[gg] + TMAP_ADDR_BYTE, gp2.b_addr);
+     }
+   }
+   __syncthreads();
+  
    constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U)
      | ((uint32_t)BLOCK_N >> 3U << 17U) | ((uint32_t)128 >> 7U << 27U);
   
-   const TileInfo ti = lookup->tiles[bid];
-   const GroupGemmParams& p = params[ti.group_idx];
-   const int M = p.M; const int N = p.N; const int K = p.K;
-   const int bid_m = (int)ti.bid_m;
-   const int bid_n = (int)ti.bid_n;
-   const int off_m = bid_m * BLOCK_M;
-   const int off_n = bid_n * BLOCK_N;
-   const int num_iters = K / BLOCK_K;
+   // Direct tile lookup from compact table (no division needed)
+   const uint16_t ti = kp.tiles[bid];
+   const int g = tile_group(ti);
+   const int bid_m = tile_bid_m(ti);
+   const int bid_n = tile_bid_n(ti);
   
+   const GroupDynArgs& gp = kp.groups[g];
+   const int M = gp.M; const int N = gp.N;
+   const int off_m = bid_m * BLOCK_M; const int off_n = bid_n * BLOCK_N;
+   const int num_iters = tile_num_k_iters(ti);
+  
+   // TMA warp
    if (warp_id == NUM_WARPS - 2 && elect_sync()) {
      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 = p.cache_A;
-     uint64_t cache_B = p.cache_B;
+     const CUtensorMap* A_tmap = &d_A_tmaps[g];
+     const CUtensorMap* B_tmap = &d_B_tmaps[g];
+     const char* SFA_ptr = gp.SFA_ptr;
+     const char* SFB_ptr = gp.SFB_ptr;
+     uint64_t cache_A = gp.cache_A;
+     uint64_t cache_B = gp.cache_B;
+     const int rest_k = num_iters * (BLOCK_K / 64);
      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;
  ⋯ 4 unchanged lines
        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;
-       tma_gmem2smem(SFA_smem, SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512, SFA_size, mbar_addr, cache_A);
-       tma_gmem2smem(SFB_smem, SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512, SFB_size, mbar_addr, cache_B);
+       const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / 64) * 512;
+       const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / 64) * 512;
+       // Scale factors are tiny (~2KB each), always keep in L2
+       tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, EVICT_LAST);
+       tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, EVICT_LAST);
        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;
      }
    }
+   // MMA warp
    else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
      int tma_stage = 0; int tma_phase = 0;
      const int scale_A_offset = (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
  ⋯ 42 unchanged lines
    }
    __syncthreads();
    asm volatile("tcgen05.fence::after_thread_sync;");
-  
+   // Epilogue
    if (tid < BLOCK_M) {
      const int local_warp_id = tid / WARP_SIZE;
-     half* C_ptr = p.C_ptr;
+     half* C_ptr = gp.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);
  ⋯ 19 unchanged lines
      asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N));
  }
   
- struct SlotBuffers {
-   GroupGemmParams* d_params;
-   TileLookup* d_lookup;
- };
+ // ========== Host-side launch infrastructure ==========
   
- static SlotBuffers g_slots[MAX_BATCH_SLOTS];
- static int g_num_slots = 0;
- static size_t g_slot_groups_capacity = 0;
- static GroupGemmParams* h_params_pinned = nullptr;
+ constexpr int BLOCK_M = 128, BLOCK_N = 128, BLOCK_K = 256, NUM_SMS = 148;
+ constexpr int NUM_STAGES_P = 6, NUM_STAGES_S = 6;
+ constexpr int tb_size = BLOCK_M + 2 * WARP_SIZE;
+ constexpr int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
+ constexpr int SFAB_size = 128 * (BLOCK_K / 16) * 2;
   
- inline void ensure_slots(int num_slots, int max_groups) {
-   if (g_num_slots >= num_slots && g_slot_groups_capacity >= (size_t)max_groups) return;
-   for (int i = 0; i < g_num_slots; i++) {
-     cudaFree(g_slots[i].d_params);
-     cudaFree(g_slots[i].d_lookup);
-   }
-   if (h_params_pinned) cudaFreeHost(h_params_pinned);
-   g_num_slots = num_slots;
-   g_slot_groups_capacity = max_groups + 8;
-   for (int i = 0; i < num_slots; i++) {
-     cudaMalloc(&g_slots[i].d_params, g_slot_groups_capacity * sizeof(GroupGemmParams));
-     cudaMalloc(&g_slots[i].d_lookup, sizeof(TileLookup));
-   }
-   cudaHostAlloc(&h_params_pinned, g_slot_groups_capacity * sizeof(GroupGemmParams), cudaHostAllocDefault);
- }
+ static int g_cached_ng = 0;
+ static int g_cached_n[MAX_GROUPS];
+ static int g_cached_k[MAX_GROUPS];
+ static int g_grid_n[MAX_GROUPS];
+ static uint64_t g_cache_B_policy[MAX_GROUPS];
+ static uint64_t g_cached_cA[MAX_GROUPS];
+ static uint64_t g_cached_cB[MAX_GROUPS];
+ static bool g_templates_init = false;
   
- struct GraphCacheEntry {
-   cudaGraphExec_t graph_exec;
-   int num_groups;
-   int total_tiles;
-   int grid_size;
-   bool use_persistent;
-   int slot_idx;
-   const char* A_ptrs[MAX_GROUPS];
-   c10::Half* C_ptrs[MAX_GROUPS];
-   bool valid;
- };
+ static CUfunction g_cu_persistent = nullptr;
+ static CUfunction g_cu_simple = nullptr;
+ static int g_smem_persistent = 0;
+ static int g_smem_simple = 0;
+ static bool g_init = false;
   
- static constexpr int MAX_GRAPH_CACHE = 128;
- static GraphCacheEntry g_graph_cache[MAX_GRAPH_CACHE];
- static int g_graph_cache_idx = 0;
+ static KernelParams g_kp;
+ static SimpleKernelParams g_skp;
+ static void* g_kp_args[1] = { &g_kp };
+ static void* g_skp_args[1] = { &g_skp };
   
- #include <unordered_map>
- static std::unordered_map<uint64_t, int> g_graph_hash_map;
+ static CUlaunchConfig g_launch_config;
+ static CUlaunchAttribute g_launch_attrs[1];
   
- 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;
- }
+ static CUtensorMap* h_d_A_tmaps = nullptr;
+ static CUtensorMap* h_d_B_tmaps = nullptr;
   
- torch::Tensor group_gemm_packed(torch::Tensor packed_args, int64_t num_groups, int64_t slot_idx) {
-     const int ng = static_cast<int>(num_groups);
-     const int si = static_cast<int>(slot_idx);
-     const int64_t* data = packed_args.data_ptr<int64_t>();
+ // N, K, cache_A, cache_B are now passed in GroupDynArgs (no device globals needed)
   
-     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];
+ static void ensure_init() {
+   if (!g_init) {
+     auto pk = group_gemm_persistent_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES_P>;
+     auto sk = group_gemm_simple_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES_S>;
+     g_smem_persistent = (AB_size + SFAB_size) * NUM_STAGES_P;
+     g_smem_simple = (AB_size + SFAB_size) * NUM_STAGES_S;
+     if (g_smem_persistent > 48000) cudaFuncSetAttribute(pk, cudaFuncAttributeMaxDynamicSharedMemorySize, g_smem_persistent);
+     if (g_smem_simple > 48000) cudaFuncSetAttribute(sk, cudaFuncAttributeMaxDynamicSharedMemorySize, g_smem_simple);
+     cudaGetFuncBySymbol(&g_cu_persistent, (const void*)pk);
+     cudaGetFuncBySymbol(&g_cu_simple, (const void*)sk);
   
-     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]);
-     }
+     memset(&g_launch_config, 0, sizeof(g_launch_config));
+     g_launch_config.gridDimY = 1;
+     g_launch_config.gridDimZ = 1;
+     g_launch_config.blockDimX = tb_size;
+     g_launch_config.blockDimY = 1;
+     g_launch_config.blockDimZ = 1;
+     g_launch_attrs[0].id = (CUlaunchAttributeID)6;
+     *(int*)&g_launch_attrs[0].value = 1;
+     g_launch_config.numAttrs = 1;
+     g_launch_config.attrs = g_launch_attrs;
   
-     constexpr int BLOCK_M = 128, BLOCK_N = 128, BLOCK_K = 256, NUM_STAGES = 6, NUM_SMS = 148;
+     cudaGetSymbolAddress((void**)&h_d_A_tmaps, d_A_tmaps);
+     cudaGetSymbolAddress((void**)&h_d_B_tmaps, d_B_tmaps);
+     // N, K, cache_A, cache_B are passed in GroupDynArgs
   
-     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);
-             auto result = torch::tensor({(int64_t)(uintptr_t)g_graph_cache[i].graph_exec, (int64_t)i}, torch::kInt64);
-             return result;
-         }
-     }
+     cudaDeviceSynchronize();
+     g_init = true;
+   }
+ }
   
-     GroupGemmParams* d_params = g_slots[si].d_params;
-     TileLookup* d_lookup = g_slots[si].d_lookup;
+ // Packing + launch in a single C call via ctypes
+ // packed layout: [a_ptrs(ng), b_ptrs(ng), c_ptrs(ng), sfa_ptrs(ng), sfb_ptrs(ng), M(ng), N(ng), K(ng)]
+ extern "C" void launch_group_gemm(const int64_t* packed, int ng) {
+   ensure_init();
   
-     int total_tiles = 0;
-     TileLookup h_lookup;
+   bool config_changed = !g_templates_init || g_cached_ng != ng;
+   if (!config_changed) {
      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];
-         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;
-         int64_t B_bytes = (int64_t)n_sizes[g] * k_sizes[g] / 2;
-         h_params_pinned[g].cache_A = EVICT_LAST;
-         if (B_bytes > 8*1024*1024)
-           h_params_pinned[g].cache_B = EVICT_FIRST;
-         else if (B_bytes > 4*1024*1024)
-           h_params_pinned[g].cache_B = EVICT_NORMAL;
-         else
-           h_params_pinned[g].cache_B = EVICT_LAST;
-         for (int t = 0; t < h_params_pinned[g].num_tiles; t++) {
-           int idx = total_tiles + t;
-           h_lookup.tiles[idx].group_idx = (int8_t)g;
-           h_lookup.tiles[idx].bid_m = (int8_t)(t / grid_n);
-           h_lookup.tiles[idx].bid_n = (int8_t)(t % grid_n);
-           h_lookup.tiles[idx]._pad = 0;
-         }
-         total_tiles += h_params_pinned[g].num_tiles;
+       if (g_cached_n[g] != (int)packed[6*ng + g] || g_cached_k[g] != (int)packed[7*ng + g]) {
+         config_changed = true; break;
+       }
      }
+   }
   
-     bool use_persistent = (total_tiles > NUM_SMS);
-     int grid_size = use_persistent ? NUM_SMS : total_tiles;
+   if (config_changed) {
+     CUtensorMap h_tmap;
+     for (int g = 0; g < ng; g++) {
+       int n = (int)packed[6*ng + g];
+       int k = (int)packed[7*ng + g];
+       g_cached_n[g] = n;
+       g_cached_k[g] = k;
+       g_grid_n[g] = (n + BLOCK_N - 1) / BLOCK_N;
+       int64_t B_bytes = (int64_t)n * k / 2;
+       uint64_t cb = (B_bytes > 8*1024*1024) ? EVICT_FIRST :
+                      (B_bytes > 4*1024*1024) ? EVICT_NORMAL : EVICT_LAST;
+       g_cache_B_policy[g] = cb;
+       // A is small (max ~4MB total across groups), always keep in L2
+       g_cached_cA[g] = EVICT_LAST;
+       // B is large for K>4096 (>100MB total), evict first; for smaller K, use size-based policy
+       g_cached_cB[g] = (k > 4096) ? EVICT_FIRST : cb;
   
-     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>;
+       init_AB_tmap(&h_tmap, (const char*)0x1000000, 512, k, BLOCK_M, BLOCK_K);
+       cudaMemcpy(&h_d_A_tmaps[g], &h_tmap, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
   
-     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;
-  
-     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;
+       init_AB_tmap(&h_tmap, (const char*)0x1000000, n, k, BLOCK_N, BLOCK_K);
+       cudaMemcpy(&h_d_B_tmaps[g], &h_tmap, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
      }
   
-     auto this_kernel = use_persistent ? persistent_kernel : simple_kernel;
-     cudaMemcpy(d_params, h_params_pinned, ng * sizeof(GroupGemmParams), cudaMemcpyHostToDevice);
-     cudaMemcpy(d_lookup, &h_lookup, total_tiles * sizeof(TileInfo), cudaMemcpyHostToDevice);
+     g_cached_ng = ng;
+     g_templates_init = true;
+     cudaDeviceSynchronize();
+   }
   
-     cudaGraph_t graph;
-     cudaGraphCreate(&graph, 0);
-     cudaGraphNode_t kernel_node;
-     cudaKernelNodeParams kernel_params = {0};
-     void* kernel_args[] = { &d_params, &d_lookup, (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 = kernel_args;
-     kernel_params.extra = nullptr;
-     cudaGraphAddKernelNode(&kernel_node, graph, nullptr, 0, &kernel_params);
+   // Compute per-group grid_m
+   int m_arr[MAX_GROUPS];
+   int gm_arr[MAX_GROUPS];
+   int total_tiles = 0;
+   for (int g = 0; g < ng; g++) {
+     int m = (int)packed[5*ng + g];
+     m_arr[g] = m;
+     gm_arr[g] = (m + BLOCK_M - 1) / BLOCK_M;
+     total_tiles += gm_arr[g] * g_grid_n[g];
+   }
   
-     cudaGraphExec_t graph_exec;
-     cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0);
-     cudaGraphLaunch(graph_exec, 0);
+   bool use_persistent = (total_tiles > NUM_SMS);
   
-     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);
+   if (use_persistent) {
+     // Pack persistent kernel args (with tile table)
+     for (int g = 0; g < ng; g++) {
+       GroupDynArgs& gp = g_kp.groups[g];
+       gp.a_addr = (uint64_t)packed[g];
+       gp.b_addr = (uint64_t)packed[ng + g];
+       gp.SFA_ptr = (const char*)packed[3*ng + g];
+       gp.SFB_ptr = (const char*)packed[4*ng + g];
+       gp.C_ptr = (half*)packed[2*ng + g];
+       gp.M = m_arr[g];
+       gp.N = g_cached_n[g];
+       gp.cache_A = g_cached_cA[g];
+       gp.cache_B = g_cached_cB[g];
      }
-     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].grid_size = grid_size;
-     g_graph_cache[new_idx].use_persistent = use_persistent;
-     g_graph_cache[new_idx].slot_idx = si;
-     g_graph_cache[new_idx].valid = true;
+     g_kp.ng = ng;
+  
+     // Build tile table with group-interleaved ordering for L2 locality
+     bool same_n = true;
+     for (int g = 1; g < ng; g++) {
+       if (g_cached_n[g] != g_cached_n[0]) { same_n = false; break; }
+     }
+  
+     int tt = 0;
+     if (same_n && ng > 1) {
+       int grid_n = g_grid_n[0];
+       for (int bn = 0; bn < grid_n; bn++) {
+         for (int g = 0; g < ng; g++) {
+           int k_iters = g_cached_k[g] / BLOCK_K;
+           for (int bm = 0; bm < gm_arr[g]; bm++) {
+             pack_tile(g_kp.tiles[tt], g, bm, bn, k_iters);
+             tt++;
+           }
+         }
+       }
+     } else {
+       for (int g = 0; g < ng; g++) {
+         int grid_n = g_grid_n[g];
+         int k_iters = g_cached_k[g] / BLOCK_K;
+         for (int bn = 0; bn < grid_n; bn++) {
+           for (int bm = 0; bm < gm_arr[g]; bm++) {
+             pack_tile(g_kp.tiles[tt], g, bm, bn, k_iters);
+             tt++;
+           }
+         }
+       }
+     }
+     g_kp.total_tiles = total_tiles;
+  
+     g_launch_config.gridDimX = NUM_SMS;
+     g_launch_config.sharedMemBytes = g_smem_persistent;
+     cuLaunchKernelEx(&g_launch_config, g_cu_persistent, g_kp_args, nullptr);
+   } else {
+     // Pack simple kernel args with compact tile table
      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];
+       GroupDynArgs& gp = g_skp.groups[g];
+       gp.a_addr = (uint64_t)packed[g];
+       gp.b_addr = (uint64_t)packed[ng + g];
+       gp.SFA_ptr = (const char*)packed[3*ng + g];
+       gp.SFB_ptr = (const char*)packed[4*ng + g];
+       gp.C_ptr = (half*)packed[2*ng + g];
+       gp.M = m_arr[g];
+       gp.N = g_cached_n[g];
+       gp.cache_A = g_cached_cA[g];
+       gp.cache_B = g_cached_cB[g];
      }
-     g_graph_hash_map[target_hash] = new_idx;
-     g_graph_cache_idx++;
-     cudaGraphDestroy(graph);
-     auto result = torch::tensor({(int64_t)(uintptr_t)graph_exec, (int64_t)new_idx}, torch::kInt64);
-     return result;
+     g_skp.ng = ng;
+     g_skp.total_tiles = total_tiles;
+     // Build compact tile table (max 148 entries)
+     int ti = 0;
+     for (int g = 0; g < ng; g++) {
+       int gn = g_grid_n[g];
+       int k_iters = g_cached_k[g] / BLOCK_K;
+       for (int bn = 0; bn < gn; bn++) {
+         for (int bm = 0; bm < gm_arr[g]; bm++) {
+           pack_tile(g_skp.tiles[ti], g, bm, bn, k_iters);
+           ti++;
+         }
+       }
+     }
+  
+     g_launch_config.gridDimX = total_tiles;
+     g_launch_config.sharedMemBytes = g_smem_simple;
+     cuLaunchKernelEx(&g_launch_config, g_cu_simple, g_skp_args, nullptr);
+   }
  }
   
- int64_t create_batch_graph(torch::Tensor cache_indices) {
-     const int64_t* indices = cache_indices.data_ptr<int64_t>();
-     int n = cache_indices.size(0);
+ // _dummy_init is defined in cpp_source
   
-     constexpr int BLOCK_M = 128, BLOCK_N = 128, BLOCK_K = 256, NUM_STAGES = 6;
-     auto persistent_fn = (void*)group_gemm_persistent_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
-     auto simple_fn = (void*)group_gemm_simple_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
+ #include <torch/extension.h>
+ """
   
-     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;
+ cpp_source = r"""
+ #include <torch/extension.h>
+ #include <Python.h>
+ #include <torch/csrc/autograd/python_variable.h>
   
-     cudaGraph_t batch_graph;
-     cudaGraphCreate(&batch_graph, 0);
+ int64_t _dummy_init() { return 0; }
   
-     static GroupGemmParams** batch_d_params = nullptr;
-     static TileLookup** batch_d_lookup = nullptr;
-     static int* batch_total_tiles = nullptr;
-     static int batch_alloc = 0;
-     if (batch_alloc < n) {
-         free(batch_d_params);
-         free(batch_d_lookup);
-         free(batch_total_tiles);
-         batch_alloc = n + 8;
-         batch_d_params = (GroupGemmParams**)malloc(batch_alloc * sizeof(GroupGemmParams*));
-         batch_d_lookup = (TileLookup**)malloc(batch_alloc * sizeof(TileLookup*));
-         batch_total_tiles = (int*)malloc(batch_alloc * sizeof(int));
-     }
+ // Forward-declare the CUDA launch function
+ extern "C" void launch_group_gemm(const int64_t* packed, int ng);
   
-     for (int i = 0; i < n; i++) {
-         int ci = (int)indices[i];
-         auto& entry = g_graph_cache[ci];
-         batch_d_params[i] = g_slots[entry.slot_idx].d_params;
-         batch_d_lookup[i] = g_slots[entry.slot_idx].d_lookup;
-         batch_total_tiles[i] = entry.total_tiles;
-     }
+ static int64_t g_buf[8 * 8];
   
-     for (int i = 0; i < n; i++) {
-         int ci = (int)indices[i];
-         auto& entry = g_graph_cache[ci];
+ // Raw CPython function: fast_launch(data) -> list[Tensor]
+ static PyObject* fast_launch_impl(PyObject* self, PyObject* arg) {
+     // arg is data tuple: (abc_list, orig_list, sf_list, ps_list)
+     PyObject* abc_list = PyTuple_GET_ITEM(arg, 0);
+     PyObject* sf_list  = PyTuple_GET_ITEM(arg, 2);
+     PyObject* ps_list  = PyTuple_GET_ITEM(arg, 3);
   
-         cudaGraphNode_t kernel_node;
-         cudaKernelNodeParams kp = {0};
-         void* args[] = { &batch_d_params[i], &batch_d_lookup[i], &batch_total_tiles[i] };
-         kp.func = entry.use_persistent ? persistent_fn : simple_fn;
-         kp.gridDim = dim3(entry.grid_size);
-         kp.blockDim = dim3(tb_size);
-         kp.sharedMemBytes = smem_size;
-         kp.kernelParams = args;
-         kp.extra = nullptr;
+     Py_ssize_t ng = PyList_GET_SIZE(ps_list);
+     int64_t* buf = g_buf;
   
-         cudaGraphAddKernelNode(&kernel_node, batch_graph, nullptr, 0, &kp);
+     PyObject* c_list = PyList_New(ng);
+  
+     for (Py_ssize_t i = 0; i < ng; i++) {
+         PyObject* abc_t = PyList_GET_ITEM(abc_list, i);
+         PyObject* sf_t  = PyList_GET_ITEM(sf_list, i);
+         PyObject* pi    = PyList_GET_ITEM(ps_list, i);
+  
+         const at::Tensor& A   = THPVariable_Unpack(PyTuple_GET_ITEM(abc_t, 0));
+         const at::Tensor& B   = THPVariable_Unpack(PyTuple_GET_ITEM(abc_t, 1));
+         PyObject* C_obj = PyTuple_GET_ITEM(abc_t, 2);
+         const at::Tensor& C   = THPVariable_Unpack(C_obj);
+         const at::Tensor& SFA = THPVariable_Unpack(PyTuple_GET_ITEM(sf_t, 0));
+         const at::Tensor& SFB = THPVariable_Unpack(PyTuple_GET_ITEM(sf_t, 1));
+  
+         buf[i]          = (int64_t)A.data_ptr();
+         buf[ng + i]     = (int64_t)B.data_ptr();
+         buf[2*ng + i]   = (int64_t)C.data_ptr();
+         buf[3*ng + i]   = (int64_t)SFA.data_ptr();
+         buf[4*ng + i]   = (int64_t)SFB.data_ptr();
+         buf[5*ng + i]   = PyLong_AsLongLong(PyTuple_GET_ITEM(pi, 0));
+         buf[6*ng + i]   = PyLong_AsLongLong(PyTuple_GET_ITEM(pi, 1));
+         buf[7*ng + i]   = PyLong_AsLongLong(PyTuple_GET_ITEM(pi, 2));
+  
+         Py_INCREF(C_obj);
+         PyList_SET_ITEM(c_list, i, C_obj);
      }
   
-     cudaGraphExec_t batch_exec;
-     cudaGraphInstantiate(&batch_exec, batch_graph, nullptr, nullptr, 0);
-     cudaGraphDestroy(batch_graph);
-     return (int64_t)(uintptr_t)batch_exec;
+     launch_group_gemm(buf, (int)ng);
+     return c_list;
  }
   
- void init_slots(int64_t num_slots, int64_t max_groups) {
-     ensure_slots((int)num_slots, (int)max_groups);
- }
+ static PyMethodDef extra_methods[] = {
+     {"fast_launch", (PyCFunction)fast_launch_impl, METH_O, "Fast launch group GEMM"},
+     {NULL, NULL, 0, NULL}
+ };
   
- #include <torch/extension.h>
+ int64_t get_methods_ptr() {
+     return (int64_t)(void*)extra_methods;
+ }
  """
   
- cpp_source = """
- #include <torch/extension.h>
- torch::Tensor group_gemm_packed(torch::Tensor packed_args, int64_t num_groups, int64_t slot_idx);
- int64_t create_batch_graph(torch::Tensor cache_indices);
- void init_slots(int64_t num_slots, int64_t max_groups);
- """
-  
  _module = load_inline(
-     name='group_gemm_fp4_v24',
+     name='group_gemm_fp4_opt31_packed_tile16',
      cpp_sources=cpp_source,
      cuda_sources=cuda_source,
-     functions=['group_gemm_packed', 'create_batch_graph', 'init_slots'],
+     functions=['_dummy_init', 'get_methods_ptr'],
      verbose=False,
      extra_cuda_cflags=[
          "-O3",
  ⋯ 7 unchanged lines
      extra_ldflags=["-lcuda"],
  )
   
- _group_gemm_pybind = _module.group_gemm_packed
- _create_batch_graph = _module.create_batch_graph
- _init_slots = _module.init_slots
-  
- _init_slots(16, 10)
-  
- _ptr_cache = {}
- _next_slot = 0
- _graph_launch = _libcudart.cudaGraphLaunch
-  
- _call_sequence = []
- _known_seq = None
- _batch_handle = None
- _batch_call_idx = 0
- _batch_cache = {}
-  
- def custom_kernel(data: input_t) -> output_t:
-     global _call_sequence, _known_seq, _batch_handle, _batch_call_idx, _next_slot
-  
-     abc_tensors = data[0]
-     a0 = abc_tensors[0][0]
-     c0 = abc_tensors[0][2]
-     ng = len(data[3])
-     k = (a0.data_ptr(), c0.data_ptr(), ng)
-  
-     if _known_seq is not None and _batch_handle is not None:
-         entry = _ptr_cache.get(k)
-         if entry is not None:
-             _, c_tensors, _, _ = entry
-             if _batch_call_idx == 0:
-                 if k == _known_seq[0]:
-                     _graph_launch(_batch_handle, _ZERO)
-                 else:
-                     _known_seq = None
-                     _batch_handle = None
-                     _batch_call_idx = 0
-                     _call_sequence = [k]
-                     _graph_launch(entry[0], _ZERO)
-                     return c_tensors
-             _batch_call_idx = (_batch_call_idx + 1) % len(_known_seq)
-             return c_tensors
-         else:
-             _known_seq = None
-             _batch_handle = None
-             _batch_call_idx = 0
-             _call_sequence = []
-  
-     entry = _ptr_cache.get(k)
-     if entry is not None:
-         handle, c_tensors, cache_idx, slot_idx = entry
-         _graph_launch(handle, _ZERO)
-  
-         _call_sequence.append(k)
-         if len(_call_sequence) >= 30:
-             half = len(_call_sequence) // 2
-             first = tuple(_call_sequence[:half])
-             second = tuple(_call_sequence[half:2*half])
-             if first == second:
-                 cached_batch = _batch_cache.get(first)
-                 if cached_batch is not None:
-                     _known_seq = first
-                     _batch_handle = cached_batch
-                     _batch_call_idx = 0
-                     _call_sequence = []
-                 else:
-                     indices = []
-                     for seq_k in first:
-                         _, _, ci, _ = _ptr_cache[seq_k]
-                         indices.append(ci)
-                     idx_tensor = torch.tensor(indices, dtype=torch.int64, device='cpu')
-                     batch_exec_handle = _create_batch_graph(idx_tensor)
-                     bh = ctypes.c_void_p(batch_exec_handle)
-                     _batch_cache[first] = bh
-                     _known_seq = first
-                     _batch_handle = bh
-                     _batch_call_idx = 0
-                     _call_sequence = []
-             elif len(_call_sequence) > 60:
-                 _call_sequence = _call_sequence[-30:]
-         return c_tensors
-  
-     slot_idx = _next_slot % 16
-     _next_slot += 1
-  
-     sfasfb_reordered_tensors = data[2]
-     problem_sizes = data[3]
-  
-     packed = []
-     for i in range(ng):
-         packed.append(abc_tensors[i][0].data_ptr())
-     for i in range(ng):
-         packed.append(abc_tensors[i][1].data_ptr())
-     for i in range(ng):
-         packed.append(abc_tensors[i][2].data_ptr())
-     for i in range(ng):
-         packed.append(sfasfb_reordered_tensors[i][0].data_ptr())
-     for i in range(ng):
-         packed.append(sfasfb_reordered_tensors[i][1].data_ptr())
-     for i in range(ng):
-         packed.append(problem_sizes[i][0])
-     for i in range(ng):
-         packed.append(problem_sizes[i][1])
-     for i in range(ng):
-         packed.append(problem_sizes[i][2])
-  
-     packed_tensor = torch.tensor(packed, dtype=torch.int64, device='cpu')
-     c_tensors = [abc_tensors[i][2] for i in range(ng)]
-  
-     result = _group_gemm_pybind(packed_tensor, ng, slot_idx)
-     graph_exec_handle = result[0].item()
-     cache_idx = result[1].item()
-     handle = ctypes.c_void_p(graph_exec_handle)
-     _ptr_cache[k] = (handle, c_tensors, cache_idx, slot_idx)
-  
-     _call_sequence.append(k)
-  
-     return c_tensors
+ # Register fast_launch as a raw CPython method on the module
+ import ctypes as _ct
+ _methods_ptr = _module.get_methods_ptr()
+ _py_dll = _ct.pythonapi
+ _py_dll.PyModule_AddFunctions.restype = _ct.c_int
+ _py_dll.PyModule_AddFunctions.argtypes = [_ct.py_object, _ct.c_void_p]
+ _rc = _py_dll.PyModule_AddFunctions(_module, _ct.c_void_p(_methods_ptr))
+ assert _rc == 0, f"PyModule_AddFunctions failed: {_rc}"
+ custom_kernel = _module.fast_launch