submission 498274

import torch
import ctypes
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline

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 = 8;
constexpr int MAX_TILES = 768;

constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;
constexpr uint64_t EVICT_LAST = 0x14F0000000000000;
constexpr uint64_t EVICT_NORMAL = 0x10F0000000000000;

__device__ inline
constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; }

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

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

__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;"
              :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy)
              : "memory");
}

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

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

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

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

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

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

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

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;
  int N;
  uint64_t cache_A;
  uint64_t cache_B;
};

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

// 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(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;
  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;
  #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);
      mbarrier_init(mma_mbar_addr + i * 8, 1);
    }
    for (int i = 0; i < 2; i++) {
      mbarrier_init(mainloop_mbar_addr + i * 8, 1);
      mbarrier_init(epilogue_mbar_addr + i * 8, BLOCK_M / WARP_SIZE);
    }
    asm volatile("fence.mbarrier_init.release.cluster;");
  }
  else if (warp_id == 1) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 2));
  }
  __syncthreads();

  // Parallel TMA descriptor patching: each lane patches one group
  if (warp_id == 0) {
    int g = lane_id;
    if (g < kp.ng) {
      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 < 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;
        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 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;
      }
    }
  }
  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 < 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++) {
        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++) {
          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
          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 enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
            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;
      }
      asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                  :: "r"(mainloop_mbar_addr + mainloop_stage * 8) : "memory");
      mainloop_stage = (mainloop_stage + 1) % 2;
      if (mainloop_stage == 0) epilogue_phase ^= 1;
    }
  }
  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 < 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 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];
        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]});
          // N is always a multiple of BLOCK_N, so col is always within bounds
          if (row + 0 < M)
            store_steaming_half2(C_ptr + (row + 0) * N + col, val0);
          if (row + 8 < M)
            store_steaming_half2(C_ptr + (row + 8) * N + col, val1);
        }
      }
      if (elect_sync()) {
        asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];"
                    :: "r"(epilogue_mbar_addr + mainloop_stage * 8) : "memory");
      }
      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 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(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;
  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;
  #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);
  // Read tile info early (from params in constant memory, no sync 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 int num_iters = tile_num_k_iters(ti);
  const GroupDynArgs& gp = kp.groups[g];

  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) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N));
  }
  __syncthreads();

  // Only patch the group this block needs (1 group instead of ng)
  if (warp_id == 0 && elect_sync()) {
    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);

  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;

  // TMA warp
  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    int tma_stage = 0; int mma_phase = 1;
    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;
      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 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);
    const int scale_B_offset = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
    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);
      };
      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++) {
        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
        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 enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
          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;
    }
    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                :: "r"(tma_mbar_addr) : "memory");
  }
  __syncthreads();
  asm volatile("tcgen05.fence::after_thread_sync;");
  // Epilogue
  if (tid < BLOCK_M) {
    const int local_warp_id = tid / WARP_SIZE;
    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);
      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]});
        // N is always a multiple of BLOCK_N, so col is always within bounds
        if (row + 0 < M)
          store_steaming_half2(C_ptr + (row + 0) * N + col, val0);
        if (row + 8 < M)
          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));
}

// ========== Host-side launch infrastructure ==========

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;

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;

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 KernelParams g_kp;
static SimpleKernelParams g_skp;
static void* g_kp_args[1] = { &g_kp };
static void* g_skp_args[1] = { &g_skp };

static CUlaunchConfig g_launch_config;
static CUlaunchAttribute g_launch_attrs[1];

static CUtensorMap* h_d_A_tmaps = nullptr;
static CUtensorMap* h_d_B_tmaps = nullptr;

// N, K, cache_A, cache_B are now passed in GroupDynArgs (no device globals needed)

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

    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;

    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

    cudaDeviceSynchronize();
    g_init = true;
  }
}

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

  bool config_changed = !g_templates_init || g_cached_ng != ng;
  if (!config_changed) {
    for (int g = 0; g < ng; g++) {
      if (g_cached_n[g] != (int)packed[6*ng + g] || g_cached_k[g] != (int)packed[7*ng + g]) {
        config_changed = true; break;
      }
    }
  }

  if (config_changed) {
    CUtensorMap h_tmap;
    // Compute total B size across all groups for aggregate cache policy
    int64_t total_B_bytes = 0;
    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;
      total_B_bytes += (int64_t)n * k / 2;
    }
    // Use total B across all groups for cache policy decision
    // L2 on B200 is 96MB; if total B > 72MB (3/4 L2), use EVICT_FIRST; otherwise EVICT_LAST
    // 72MB threshold: Bench2 (56MB) keeps B in L2, Bench1 (112MB) evicts
    uint64_t global_cb = (total_B_bytes > 72LL*1024*1024) ? EVICT_FIRST : EVICT_LAST;
    for (int g = 0; g < ng; g++) {
      int k = g_cached_k[g];
      int64_t B_bytes = (int64_t)g_cached_n[g] * k / 2;
      // Per-group: evict first if group B alone is > 8MB
      uint64_t cb = (B_bytes > 8*1024*1024) ? EVICT_FIRST : global_cb;
      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: use per-group policy, with K>4096 override for very large K
      g_cached_cB[g] = (k > 4096) ? EVICT_FIRST : cb;

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

      init_AB_tmap(&h_tmap, (const char*)0x1000000, g_cached_n[g], k, BLOCK_N, BLOCK_K);
      cudaMemcpy(&h_d_B_tmaps[g], &h_tmap, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
    }

    g_cached_ng = ng;
    g_templates_init = true;
    cudaDeviceSynchronize();
  }

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

  bool use_persistent = (total_tiles > NUM_SMS);

  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_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++) {
      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_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);
  }
}

// _dummy_init is defined in cpp_source

#include <torch/extension.h>
"""

cpp_source = r"""
#include <torch/extension.h>
#include <Python.h>
#include <torch/csrc/autograd/python_variable.h>

int64_t _dummy_init() { return 0; }

// Forward-declare the CUDA launch function
extern "C" void launch_group_gemm(const int64_t* packed, int ng);

static int64_t g_buf[8 * 8];

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

    Py_ssize_t ng = PyList_GET_SIZE(ps_list);
    int64_t* buf = g_buf;

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

    launch_group_gemm(buf, (int)ng);
    return c_list;
}

static PyMethodDef extra_methods[] = {
    {"fast_launch", (PyCFunction)fast_launch_impl, METH_O, "Fast launch group GEMM"},
    {NULL, NULL, 0, NULL}
};

int64_t get_methods_ptr() {
    return (int64_t)(void*)extra_methods;
}
"""

_module = load_inline(
    name='group_gemm_fp4_opt47_b_cache_72m',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['_dummy_init', 'get_methods_ptr'],
    verbose=False,
    extra_cuda_cflags=[
        "-O3",
        "-gencode=arch=compute_100a,code=sm_100a",
        "--use_fast_math",
        "--expt-relaxed-constexpr",
        "--relocatable-device-code=false",
        "-Xptxas=--allow-expensive-optimizations=true",
        "-ftz=true",
    ],
    extra_ldflags=["-lcuda"],
)

# 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