submission 407485

import os
from typing import List

import torch
from torch.utils.cpp_extension import load_inline

from task import input_t, output_t

os.environ.setdefault("CUTE_DSL_DISABLE_FILE_CACHING", "1")
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "10.0a")

_EXT: torch.nn.Module | None = None

CPP_SRC = r"""
#include <torch/extension.h>
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {}
"""


def _get_ext() -> torch.nn.Module:
    global _EXT
    if _EXT is not None:
        return _EXT

    cu_src = CUDA_SRC
    _EXT = load_inline(
        name="nvfp4_group_gemm_ext_mod",
        cpp_sources=CPP_SRC,
        cuda_sources=cu_src,
        functions=None,
        extra_cuda_cflags=[
            "-O3",
            "--use_fast_math",
            "--expt-relaxed-constexpr",
            "--extra-device-vectorization",
        ],
        extra_ldflags=["-lcuda"],
        with_cuda=True,
        verbose=False,
    )
    return _EXT


CUDA_SRC = r"""
#include <cuda.h>
#include <cudaTypedefs.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>

#include <torch/library.h>

constexpr int WARP_SIZE = 32;
constexpr int MMA_K = 64;

constexpr int BLOCK_M = 128;
constexpr int BLOCK_N = 64;
constexpr int BLOCK_K = 256;
constexpr int NUM_STAGES = 8;

constexpr uint64_t EVICT_NORMAL = 0x1000000000000000ULL;
constexpr uint64_t EVICT_FIRST  = 0x12F0000000000000ULL;
constexpr uint64_t EVICT_LAST   = 0x14F0000000000000ULL;

__device__ inline constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };
__device__ __forceinline__ int ceil_div_int(int a, int b) { return (a + b - 1) / b; }

__device__ __forceinline__ 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__ __forceinline__ void mbarrier_init(int mbar_addr, int count) {
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
}

__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) {
  uint32_t ticks = 0x989680;
  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, %2;\n\t"
    "@!P1 bra.uni LAB_WAIT;\n\t"
    "}"
    :: "r"(mbar_addr), "r"(phase), "r"(ticks)
  );
}

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

template <int CTA_GROUP = 1>
__device__ __forceinline__ void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
  asm volatile("tcgen05.cp.cta_group::%2.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc), "n"(CTA_GROUP));
}

template <int CTA_GROUP = 1>
__device__ __forceinline__ 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::%7.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), "n"(CTA_GROUP)
  );
}

struct SHAPE { static constexpr char _16x256b[] = ".16x256b"; };

template <int NUM_REGS, const char *SHAPE_, int NUM>
__device__ __forceinline__ void tcgen05_ld(float *tmp, int row, int col) {
  const int addr = (row << 16) | col;
  if constexpr (NUM_REGS == 32) {
    asm volatile("tcgen05.ld.sync.aligned%33.x%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"(addr), "C"(SHAPE_), "n"(NUM));
  }
}

__device__ __forceinline__ void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
  tcgen05_ld<32, SHAPE::_16x256b, 8>(tmp, row, col);
}

static void check_cu(CUresult err) {
  if (err == CUDA_SUCCESS) return;
  const char *error_msg_ptr = nullptr;
  cuGetErrorString(err, &error_msg_ptr);
  TORCH_CHECK(false, "cuTensorMapEncodeTiled error: ", (error_msg_ptr ? error_msg_ptr : "unknown"));
}

static void check_cuda(cudaError_t err) {
  if (err == cudaSuccess) return;
  TORCH_CHECK(false, cudaGetErrorString(err));
}

static void init_AB_tmap(
  CUtensorMap *tmap,
  const void *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]       = {256ULL, global_height, global_width / 256ULL};
  uint64_t globalStrides[rank-1] = {global_width / 2ULL, 128ULL};
  uint32_t boxDim[rank]          = {256U, shared_height, shared_width / 256U};
  uint32_t elementStrides[rank]  = {1U, 1U, 1U};

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

__global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
void grouped_kernel(
  const CUtensorMap *A_tmaps,
  const CUtensorMap *B_tmaps,
  const uint64_t *SFA_ptrs,
  const uint64_t *SFB_ptrs,
  uint64_t *C_ptrs,
  const int *Ms,
  const int *Ns,
  const int *Ks,
  const int *tile_offsets,
  int num_groups
) {
  const int tid = threadIdx.x;
  const int bid = blockIdx.x;
  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;

  int group = 0;
  #pragma unroll
  for (int i = 0; i < 8; i++) {
    if (i + 1 < num_groups) {
      if (bid >= tile_offsets[i + 1]) group = i + 1;
    }
  }

  const int M = Ms[group];
  const int N = Ns[group];
  const int K = Ks[group];

  const int tiles_n = ceil_div_int(N, BLOCK_N);
  const int local_bid = bid - tile_offsets[group];
  const int bid_m = local_bid / tiles_n;
  const int bid_n = local_bid - bid_m * tiles_n;
  const int off_m = bid_m * BLOCK_M;
  const int off_n = bid_n * BLOCK_N;

  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 + 1];
  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;

  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 * 2 + 1; i++) mbarrier_init(tma_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 * 2));
  }
  __syncthreads();

  const int num_iters = ceil_div_int(K, BLOCK_K);

  const CUtensorMap *A_tmap = A_tmaps + group;
  const CUtensorMap *B_tmap = B_tmaps + group;
  const char *SFA_ptr = reinterpret_cast<const char *>(SFA_ptrs[group]);
  const char *SFB_ptr = reinterpret_cast<const char *>(SFB_ptrs[group]);
  half *C_ptr = reinterpret_cast<half *>(C_ptrs[group]);

  uint64_t cache_A, cache_B;
  if (M > N) { cache_A = EVICT_FIRST; cache_B = EVICT_LAST; }
  else { cache_A = EVICT_LAST; cache_B = EVICT_FIRST; }

  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    auto issue_tma = [&](int iter_k, int stage_id) {
      const int mbar_addr = tma_mbar_addr + stage_id * 8;
      const int A_smem = smem + stage_id * 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");
    };

    for (int iter_k = 0; iter_k < NUM_STAGES && iter_k < num_iters; iter_k++) issue_tma(iter_k, iter_k);
    for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int mma_phase = (iter_k / NUM_STAGES - 1) % 2;
      mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
      issue_tma(iter_k, stage_id);
    }
  } else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U) | ((uint32_t)BLOCK_N >> 3U << 17U) | ((uint32_t)128 >> 7U << 27U);

    for (int iter_k = 0; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int tma_phase = (iter_k / NUM_STAGES) % 2;
      mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);

      const int A_smem = smem + stage_id * 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);

      for (int k = 0; k < BLOCK_K / MMA_K; k++) {
        const uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
        const 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);
      }

      for (int k1 = 0; k1 < BLOCK_K / 256; k1++)
        for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
          const uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
          const uint64_t b_desc = make_desc_AB(B_smem + k1 * BLOCK_N * 128 + k2 * 32);

          const int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4 + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
          const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
          const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
          tcgen05_mma_nvfp4(0, a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
        }

      asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                  :: "r"(mma_mbar_addr + stage_id * 8) : "memory");
    }

    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                :: "r"(mainloop_mbar_addr) : "memory");
  } else if (tid < BLOCK_M) {
    mbarrier_wait(mainloop_mbar_addr, 0);
    asm volatile("tcgen05.fence::after_thread_sync;");

    for (int m0 = 0; m0 < 32 / 16; m0++) {
      float tmp[BLOCK_N / 2];
      tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m0 * 16, 0);
      asm volatile("tcgen05.wait::ld.sync.aligned;");

      #pragma unroll
      for (int i = 0; i < BLOCK_N / 8; i++) {
        const int row = off_m + warp_id * 32 + m0 * 16 + lane_id / 4;
        const int col = off_n + i * 8 + (lane_id % 4) * 2;
        if (row >= M) continue;

        const float v00 = tmp[i * 4 + 0];
        const float v01 = tmp[i * 4 + 1];
        const float v10 = tmp[i * 4 + 2];
        const float v11 = tmp[i * 4 + 3];

        if (col + 1 < N) {
          reinterpret_cast<half2 *>(C_ptr + (row + 0) * N + col)[0] = __floats2half2_rn(v00, v01);
        } else if (col < N) {
          C_ptr[(row + 0) * N + col] = __float2half(v00);
        }
        if (row + 8 < M) {
          if (col + 1 < N) {
            reinterpret_cast<half2 *>(C_ptr + (row + 8) * N + col)[0] = __floats2half2_rn(v10, v11);
          } else if (col < N) {
            C_ptr[(row + 8) * N + col] = __float2half(v10);
          }
        }
      }
    }

    asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
    if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 2));
  }
}

// (No clustered G=2 kernel; we use the unified grouped kernel for all group counts.)

static void group_gemm(
  c10::List<at::Tensor> A_list,
  c10::List<at::Tensor> B_list,
  c10::List<at::Tensor> C_list,
  c10::List<at::Tensor> SFA_list,
  c10::List<at::Tensor> SFB_list,
  at::Tensor problem_sizes
) {
  const int64_t G = A_list.size();
  TORCH_CHECK(G == B_list.size() && G == C_list.size() && G == SFA_list.size() && G == SFB_list.size(), "group list sizes mismatch");
  TORCH_CHECK(problem_sizes.device().is_cpu(), "problem_sizes must be a CPU tensor");
  TORCH_CHECK(problem_sizes.scalar_type() == at::kInt && problem_sizes.dim() == 2 && problem_sizes.size(0) == G && problem_sizes.size(1) == 4,
              "problem_sizes must be int32 CPU tensor of shape [G,4]");
  TORCH_CHECK(G >= 1 && G <= 8, "expected 1..8 groups");

  std::vector<uint64_t> hC(G), hSFA(G), hSFB(G);
  std::vector<int> hM(G), hN(G), hK(G);
  std::vector<int> hOffsets(G + 1, 0);

  auto ps = problem_sizes.contiguous();
  const int *ps_ptr = ps.data_ptr<int>();

  for (int i = 0; i < (int)G; i++) {
    const int M = ps_ptr[i * 4 + 0];
    const int N = ps_ptr[i * 4 + 1];
    const int K = ps_ptr[i * 4 + 2];
    hM[i] = M; hN[i] = N; hK[i] = K;

    auto A = A_list.get(i);
    auto B = B_list.get(i);
    auto C = C_list.get(i);
    auto SFA = SFA_list.get(i);
    auto SFB = SFB_list.get(i);

    TORCH_CHECK(A.is_cuda() && B.is_cuda() && C.is_cuda() && SFA.is_cuda() && SFB.is_cuda(), "all tensors must be CUDA");
    hC[i] = (uint64_t)C.data_ptr();
    hSFA[i] = (uint64_t)SFA.data_ptr();
    hSFB[i] = (uint64_t)SFB.data_ptr();

    const int tiles_m = (M + BLOCK_M - 1) / BLOCK_M;
    const int tiles_n = (N + BLOCK_N - 1) / BLOCK_N;
    hOffsets[i + 1] = hOffsets[i] + tiles_m * tiles_n;
  }

  const int total_tiles = hOffsets[G];
  if (total_tiles == 0) return;

  std::vector<CUtensorMap> hAmap(G), hBmap(G);
  for (int i = 0; i < (int)G; i++) {
    auto A = A_list.get(i);
    auto B = B_list.get(i);
    init_AB_tmap(&hAmap[i], A.data_ptr(), (uint64_t)hM[i], (uint64_t)hK[i], (uint32_t)BLOCK_M, (uint32_t)BLOCK_K);
    init_AB_tmap(&hBmap[i], B.data_ptr(), (uint64_t)hN[i], (uint64_t)hK[i], (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
  }

  auto opts_i64 = at::TensorOptions().dtype(at::kLong).device(at::kCUDA);
  auto opts_i32 = at::TensorOptions().dtype(at::kInt).device(at::kCUDA);

  at::Tensor dC = at::empty({G}, opts_i64);
  at::Tensor dSFA = at::empty({G}, opts_i64);
  at::Tensor dSFB = at::empty({G}, opts_i64);
  at::Tensor dM = at::empty({G}, opts_i32);
  at::Tensor dN = at::empty({G}, opts_i32);
  at::Tensor dK = at::empty({G}, opts_i32);
  at::Tensor dOffsets = at::empty({G + 1}, opts_i32);

  at::Tensor dAmap = at::empty({G, (int64_t)(sizeof(CUtensorMap) / sizeof(int64_t))}, opts_i64);
  at::Tensor dBmap = at::empty({G, (int64_t)(sizeof(CUtensorMap) / sizeof(int64_t))}, opts_i64);

  check_cuda(cudaMemcpy(dC.data_ptr(), hC.data(), G * sizeof(uint64_t), cudaMemcpyHostToDevice));
  check_cuda(cudaMemcpy(dSFA.data_ptr(), hSFA.data(), G * sizeof(uint64_t), cudaMemcpyHostToDevice));
  check_cuda(cudaMemcpy(dSFB.data_ptr(), hSFB.data(), G * sizeof(uint64_t), cudaMemcpyHostToDevice));
  check_cuda(cudaMemcpy(dM.data_ptr(), hM.data(), G * sizeof(int), cudaMemcpyHostToDevice));
  check_cuda(cudaMemcpy(dN.data_ptr(), hN.data(), G * sizeof(int), cudaMemcpyHostToDevice));
  check_cuda(cudaMemcpy(dK.data_ptr(), hK.data(), G * sizeof(int), cudaMemcpyHostToDevice));
  check_cuda(cudaMemcpy(dOffsets.data_ptr(), hOffsets.data(), (G + 1) * sizeof(int), cudaMemcpyHostToDevice));
  check_cuda(cudaMemcpy(dAmap.data_ptr(), hAmap.data(), G * sizeof(CUtensorMap), cudaMemcpyHostToDevice));
  check_cuda(cudaMemcpy(dBmap.data_ptr(), hBmap.data(), G * sizeof(CUtensorMap), cudaMemcpyHostToDevice));

  dim3 grid(total_tiles, 1, 1);
  const int tb = BLOCK_M + 2 * WARP_SIZE;
  const int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
  const int SFAB_size = 128 * (BLOCK_K / 16) * 2;
  const int smem_size = (AB_size + SFAB_size) * NUM_STAGES;

  if (smem_size > 48'000) cudaFuncSetAttribute(grouped_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
  grouped_kernel<<<grid, tb, smem_size>>>(
    (const CUtensorMap*)dAmap.data_ptr(),
    (const CUtensorMap*)dBmap.data_ptr(),
    (const uint64_t*)dSFA.data_ptr<int64_t>(),
    (const uint64_t*)dSFB.data_ptr<int64_t>(),
    (uint64_t*)dC.data_ptr<int64_t>(),
    (const int*)dM.data_ptr<int>(),
    (const int*)dN.data_ptr<int>(),
    (const int*)dK.data_ptr<int>(),
    (const int*)dOffsets.data_ptr<int>(),
    (int)G
  );
}

TORCH_LIBRARY(nvfp4_group_gemm_ext, m) {
  m.def("group_gemm(Tensor[] A, Tensor[] B, Tensor[] C, Tensor[] SFA, Tensor[] SFB, Tensor problem_sizes) -> ()");
  m.impl("group_gemm", &group_gemm);
}
"""


def custom_kernel(data: input_t) -> output_t:
    abc_tensors, _sfasfb_cpu, sfasfb_reordered_tensors, problem_sizes = data
    g = len(problem_sizes)

    a_list: List[torch.Tensor] = [abc_tensors[i][0] for i in range(g)]
    b_list: List[torch.Tensor] = [abc_tensors[i][1] for i in range(g)]
    c_list: List[torch.Tensor] = [abc_tensors[i][2] for i in range(g)]
    sfa_list: List[torch.Tensor] = [sfasfb_reordered_tensors[i][0] for i in range(g)]
    sfb_list: List[torch.Tensor] = [sfasfb_reordered_tensors[i][1] for i in range(g)]

    _get_ext()

    ps = torch.tensor(problem_sizes, dtype=torch.int32, device="cpu")
    torch.ops.nvfp4_group_gemm_ext.group_gemm(a_list, b_list, c_list, sfa_list, sfb_list, ps)
    return c_list