submission 479896

import torch
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

#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 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;
  asm volatile(
    "{\n\t"
    ".reg .pred %%px;\n\t"
    "elect.sync _|%%px, %1;\n\t"
    "@%%px mov.s32 %0, 1;\n\t"
    "}"
    : "+r"(pred)
    : "r"(0xFFFFFFFF)
  );
  return pred;
}

__device__ inline
void mbarrier_init(int mbar_addr, int count) {
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
}

__device__
void mbarrier_wait(int mbar_addr, int phase) {
  asm volatile(
    "{\n\t"
    ".reg .pred P1;\n\t"
    "LAB_WAIT:\n\t"
    "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1;\n\t"
    "@P1 bra.uni DONE;\n\t"
    "bra.uni LAB_WAIT;\n\t"
    "DONE:\n\t"
    "}"
    :: "r"(mbar_addr), "r"(phase)
  );
}

// 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 "
              "[%0], [%1, {%2, %3, %4}], [%5], %6;"
              :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy)
              : "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
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
  );
}

// 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";
  static constexpr char x16[] = ".x16";
  static constexpr char x32[] = ".x32";
  static constexpr char x64[] = ".x64";
  static constexpr char x128[] = ".x128";
};

template <const char *SHAPE, const char *NUM>
__device__ inline
void tcgen05_ld_32regs(float *tmp, int row, int col) {
  asm volatile("tcgen05.ld.sync.aligned%33%34.b32 "
              "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
              "  %8,  %9, %10, %11, %12, %13, %14, %15, "
              " %16, %17, %18, %19, %20, %21, %22, %23, "
              " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];"
              : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
                "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
                "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
                "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31])
              : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}

template <const char *SHAPE, const char *NUM>
__device__ inline
void tcgen05_ld_16regs(float *tmp, int row, int col) {
  asm volatile("tcgen05.ld.sync.aligned%17%18.b32 "
              "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
              "  %8,  %9, %10, %11, %12, %13, %14, %15}, [%16];"
              : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
                "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15])
              : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}

__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) {
  tcgen05_ld_32regs<SHAPE, NUM>(tmp, row, col);
  tcgen05_ld_32regs<SHAPE, NUM>(tmp + 32, row, col + 64);
}
__device__ inline void tcgen05_ld_16x256bx16(float *tmp, int row, int col) {
  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;
  if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS)
    error_msg_ptr = "unable to get error string";
  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,
  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
  uint32_t boxDim[rank]          = {256, shared_height, shared_width / 256};
  uint32_t elementStrides[rank]  = {1, 1, 1};

  auto err = cuTensorMapEncodeTiled(
    tmap,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
    rank,
    (void *)ptr,
    globalDim,
    globalStrides,
    boxDim,
    elementStrides,
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
    CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  );
  check_cu(err);
}

// Kernel parameter structure with cumulative tile counts for fast lookup
struct GroupGemmParams {
  CUtensorMap A_tmap;
  CUtensorMap B_tmap;
  const char* SFA_ptr;
  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
};

// 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;
  }
  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)
void group_gemm_persistent_kernel(
  const GroupGemmParams* __restrict__ params,
  const TileInfo* __restrict__ tile_info,
  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;
  constexpr int B_size = BLOCK_N * BLOCK_K / 2;
  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;

  // 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));
  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;

  // 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 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);
      mbarrier_init(mma_mbar_addr + i * 8, 1);
    }
    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)

  // 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

    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 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;

      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 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;
      }
    }
  }
  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

    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 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);
        };
        auto make_desc_SF = [](int addr) -> uint64_t {
          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);
          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);
        }

        // 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);
          }

        // 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;
      }

      // 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;
    }
  }
  else if (tid < BLOCK_M) {
    // Epilogue warps - persistent loop with TMEM double-buffering
    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;
      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;

          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;
    }
  }

  __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)
void group_gemm_simple_kernel(
  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;
  constexpr int B_size = BLOCK_N * BLOCK_K / 2;
  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;

  // 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);
      mbarrier_init(mma_mbar_addr + i * 8, 1);
    }
    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);

  // 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 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;

  // Warp specialization
  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    // TMA warp
    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;

    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;
    }
  }
  else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    // MMA warp - single tile, single TMEM buffer
    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

    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);
      };
      auto make_desc_SF = [](int addr) -> uint64_t {
        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);
      }

      #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);
        }

      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;
    }

    // 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;

        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);
      }
    }
  }

  __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;

// Ensure device buffers are allocated
inline void ensure_device_buffers(int num_groups, int total_tiles) {
  if (d_params == nullptr || d_params_capacity < (size_t)num_groups) {
    if (d_params) cudaFree(d_params);
    d_params_capacity = num_groups + 8;
    cudaMalloc(&d_params, d_params_capacity * sizeof(GroupGemmParams));
  }
  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;
};

static constexpr int MAX_GRAPH_CACHE = 128;
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;
  }
  return h;
}

// Hash table for fast graph cache lookup
static uint64_t g_graph_hashes[MAX_GRAPH_CACHE];

// 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);

  for (int i = 0; i < MAX_GRAPH_CACHE; i++) {
    if (!g_graph_cache[i].valid || g_graph_hashes[i] != target_hash)
      continue;

    // Hash matched, verify exact match
    if (g_graph_cache[i].num_groups != num_groups)
      continue;

    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;
      }
    }
    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;

  // 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);

  if (cache_idx >= 0) {
    // Fast path: replay cached graph
    cudaGraphLaunch(g_graph_cache[cache_idx].graph_exec, 0);
    return;
  }

  // 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);
    }
  }

  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;

  // 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;
    kernel_params.gridDim = dim3(grid_size);
    kernel_params.blockDim = dim3(tb_size);
    kernel_params.sharedMemBytes = smem_size;
    kernel_params.kernelParams = persistent_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);

  // Instantiate and execute
  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);
}


#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
);
"""

# Load with direct function binding (bypasses torch.ops dispatcher)
_module = load_inline(
    name='group_gemm_fp4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['group_gemm_cuda'],
    verbose=True,
    extra_cuda_cflags=[
        "-O3",
        "-gencode=arch=compute_100a,code=sm_100a",
        "--use_fast_math",
        "--expt-relaxed-constexpr",
        "--relocatable-device-code=false",
        "-lineinfo",
        "-Xptxas=-v",
    ],
    extra_ldflags=["-lcuda"],
)

# Direct function reference
group_gemm = _module.group_gemm_cuda

# Python-side cache (similar to gpu_mode_solution_ref.py _pointer_tensor_cache):
# ptrs_key (full pointer tuple for correctness) -> (args, c_tensors)
_pointer_tensor_cache = {}  # Cache of (ptrs_key -> (args, c_tensors)) for all seen data sets


def custom_kernel(data: input_t) -> output_t:
    """Execute the group GEMM kernel with minimal overhead using pointer cache."""
    global _pointer_tensor_cache

    abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
    num_groups = len(problem_sizes)

    # Create cache key from pointer values (full tuple for correctness)
    # Same pattern as gpu_mode_solution_ref.py lines 2566-2568
    ptrs_abc = tuple((abc_tensors[i][0].data_ptr(), abc_tensors[i][1].data_ptr(), abc_tensors[i][2].data_ptr())
                     for i in range(num_groups))
    ptrs_sfasfb = tuple((sfasfb_reordered_tensors[i][0].data_ptr(), sfasfb_reordered_tensors[i][1].data_ptr())
                        for i in range(num_groups))
    ptrs_key = (ptrs_abc, ptrs_sfasfb)

    # Check pointer tensor cache (like gpu_mode_solution_ref.py line 2570)
    if ptrs_key in _pointer_tensor_cache:
        cached_args, cached_c_tensors = _pointer_tensor_cache[ptrs_key]
        group_gemm(*cached_args)
        return cached_c_tensors

    # Cache miss: extract pointers and sizes
    a_ptrs = [abc_tensors[i][0].data_ptr() for i in range(num_groups)]
    b_ptrs = [abc_tensors[i][1].data_ptr() for i in range(num_groups)]
    c_ptrs = [abc_tensors[i][2].data_ptr() for i in range(num_groups)]
    sfa_ptrs = [sfasfb_reordered_tensors[i][0].data_ptr() for i in range(num_groups)]
    sfb_ptrs = [sfasfb_reordered_tensors[i][1].data_ptr() for i in range(num_groups)]
    m_list = [problem_sizes[i][0] for i in range(num_groups)]
    n_list = [problem_sizes[i][1] for i in range(num_groups)]
    k_list = [problem_sizes[i][2] for i in range(num_groups)]

    args = (a_ptrs, b_ptrs, c_ptrs, sfa_ptrs, sfb_ptrs, m_list, n_list, k_list, num_groups)
    c_tensors = [abc_tensors[i][2] for i in range(num_groups)]

    # Store in pointer tensor cache (like gpu_mode_solution_ref.py line 2578)
    _pointer_tensor_cache[ptrs_key] = (args, c_tensors)

    group_gemm(*args)

    return c_tensors