submission 409402

from __future__ import annotations

import os
from typing import List

import torch
from torch.utils.cpp_extension import load_inline

_EXT_READY = False
_SCRATCH_A: dict = {}


def _load_ext() -> None:
    global _EXT_READY
    if _EXT_READY:
        return

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

#include <torch/extension.h>
#include <torch/library.h>
#include <ATen/ATen.h>
#include <ATen/core/Tensor.h>

#include <cstdint>
#include <vector>


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

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

__device__ __forceinline__ int64_t globaltimer() {
  int64_t t;
  asm volatile("mov.u64 %0, %globaltimer;" : "=l"(t) :: "memory");
  return t;
}

__device__ __forceinline__ constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3FFFFULL) >> 4ULL; }

__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 DONE;\n\t"
    "bra.uni LAB_WAIT;\n\t"
    "DONE:\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"
  );
}

__device__ __forceinline__ 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__ __forceinline__ void tcgen05_mma_nvfp4(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  const int d_tmem = 0;
  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";
};
struct NUM {
  static constexpr char x64[] = ".x64";
};

template <const char *SHAPE_V, const char *NUM_V>
__device__ __forceinline__ void tcgen05_ld_64regs(float *tmp, int row, int col) {
  asm volatile(
    "tcgen05.ld.sync.aligned%65%66.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, %33, %34, %35, %36, %37, %38, %39, "
    " %40, %41, %42, %43, %44, %45, %46, %47, "
    " %48, %49, %50, %51, %52, %53, %54, %55, "
    " %56, %57, %58, %59, %60, %61, %62, %63}, [%64];"
    : "=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]),
      "=f"(tmp[32]), "=f"(tmp[33]), "=f"(tmp[34]), "=f"(tmp[35]), "=f"(tmp[36]), "=f"(tmp[37]), "=f"(tmp[38]), "=f"(tmp[39]),
      "=f"(tmp[40]), "=f"(tmp[41]), "=f"(tmp[42]), "=f"(tmp[43]), "=f"(tmp[44]), "=f"(tmp[45]), "=f"(tmp[46]), "=f"(tmp[47]),
      "=f"(tmp[48]), "=f"(tmp[49]), "=f"(tmp[50]), "=f"(tmp[51]), "=f"(tmp[52]), "=f"(tmp[53]), "=f"(tmp[54]), "=f"(tmp[55]),
      "=f"(tmp[56]), "=f"(tmp[57]), "=f"(tmp[58]), "=f"(tmp[59]), "=f"(tmp[60]), "=f"(tmp[61]), "=f"(tmp[62]), "=f"(tmp[63])
    : "r"((row << 16) | col), "C"(SHAPE_V), "C"(NUM_V));
}

__device__ __forceinline__ void tcgen05_ld_32x32bx64(float *tmp, int row, int col) {
  tcgen05_ld_64regs<SHAPE::_32x32b, NUM::x64>(tmp, row, col);
}

static __forceinline__ void check_cu(CUresult err) {
  if (err == CUDA_SUCCESS) return;
  const char *msg = "unknown";
  cuGetErrorString(err, &msg);
  TORCH_CHECK(false, msg);
}

static __forceinline__ 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_NONE,
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  );
  check_cu(err);
}

template <int BLOCK_N, int NUM_STAGES>
__global__ __launch_bounds__(128 + 2 * WARP_SIZE)
void kernel(
  const __grid_constant__ CUtensorMap A_tmap,
  const __grid_constant__ CUtensorMap B_tmap,
  const char *SFA_ptr,
  const char *SFB_ptr,
  half *C_ptr,
  int M, int N, int K
) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_K = 256;

  const int tid = threadIdx.x;
  const int bid_n = blockIdx.x;
  const int bid_m = blockIdx.y;

  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;

  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 = K / BLOCK_K;

  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    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;
    }

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

    const int init_stage = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
    for (int iter_k = 0; iter_k < init_stage; 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 int MMA_N = BLOCK_N;
    constexpr int MMA_M = 128;
    constexpr uint32_t i_desc = (1U << 7U)
                              | (1U << 10U)
                              | ((uint32_t)MMA_N >> 3U << 17U)
                              | ((uint32_t)MMA_M >> 7U << 27U);

    auto make_desc_AB = [] __device__ (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 = [] __device__ (int addr) -> uint64_t {
      const int SBO = 8 * 16;
      return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
    };

    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;

      const uint64_t SFA_desc = make_desc_SF(0) + ((uint64_t)SFA_smem >> 4ULL);
      const uint64_t SFB_desc = make_desc_SF(0) + ((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);

          const int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4;
          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(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;");

    constexpr int WIDTH = (BLOCK_N < 64) ? BLOCK_N : 64;
    #pragma unroll
    for (int n0 = 0; n0 < BLOCK_N / WIDTH; n0++) {
      float tmp[WIDTH];
      tcgen05_ld_32x32bx64(tmp, warp_id * 32, n0 * WIDTH);
      asm volatile("tcgen05.wait::ld.sync.aligned;");

      #pragma unroll
      for (int i = 0; i < WIDTH; i += 2) {
        const int n_idx = off_n + n0 * WIDTH + i;
        const int m_idx = off_m + tid;
        if (m_idx < M) {
          if ((n_idx + 1) < N) {
            reinterpret_cast<half2 *>(C_ptr + m_idx * N + n_idx)[0] = __float22half2_rn({tmp[i + 0], tmp[i + 1]});
          } else if (n_idx < N) {
            C_ptr[m_idx * N + n_idx] = __float2half_rn(tmp[i + 0]);
          }
        }
      }
    }

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

at::Tensor gemm(
  const at::Tensor& A,
  const at::Tensor& B,
  const at::Tensor& SFA,
  const at::Tensor& SFB,
        at::Tensor& C,
  int64_t M,
  int64_t N,
  int64_t K
) {
  TORCH_CHECK(A.is_cuda() && B.is_cuda() && SFA.is_cuda() && SFB.is_cuda() && C.is_cuda(), "CUDA only");
  TORCH_CHECK(A.element_size() == 1 && B.element_size() == 1, "A/B must be packed bytes");
  TORCH_CHECK(C.scalar_type() == at::kHalf, "C must be float16");
  TORCH_CHECK(A.is_contiguous() && B.is_contiguous() && C.is_contiguous(), "A/B/C must be contiguous");
  TORCH_CHECK(A.dim() == 3 && B.dim() == 3 && C.dim() == 3, "A/B/C must be 3D");
  TORCH_CHECK(A.size(2) == 1 && B.size(2) == 1 && C.size(2) == 1, "L must be 1");

  TORCH_CHECK((K % 256) == 0, "K must be multiple of 256");
  TORCH_CHECK(int64_t(A.size(1)) * 2 == K, "A K mismatch");
  TORCH_CHECK(int64_t(B.size(1)) * 2 == K, "B K mismatch");
  TORCH_CHECK(int64_t(B.size(0)) == N, "B N mismatch");
  TORCH_CHECK(int64_t(C.size(0)) == M && int64_t(C.size(1)) == N, "C shape mismatch");

  const int64_t Apad = A.size(0);
  TORCH_CHECK(Apad >= M, "A pad too small");

  const char *A_ptr = reinterpret_cast<const char *>(A.data_ptr());
  const char *B_ptr = reinterpret_cast<const char *>(B.data_ptr());
  const char *SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
  const char *SFB_ptr = reinterpret_cast<const char *>(SFB.data_ptr());
  half *C_ptr = reinterpret_cast<half *>(C.data_ptr<at::Half>());

  CUtensorMap A_tmap, B_tmap;
  init_AB_tmap(&A_tmap, A_ptr, (uint64_t)Apad, (uint64_t)K, 128, 256);
  init_AB_tmap(&B_tmap, B_ptr, (uint64_t)N,    (uint64_t)K, 64,  256);

  constexpr int BLOCK_N = 64;
  constexpr int NUM_STAGES = 4;
  const int grid_m = int((M + 127) / 128);
  const int grid_n = int((N + BLOCK_N - 1) / BLOCK_N);

  const int tb_size = 128 + 2 * WARP_SIZE;
  const int A_size = 128 * 256 / 2;
  const int B_size = BLOCK_N * 256 / 2;
  const int SF_size = 128 * 256 / 16;
  const int smem_size = (A_size + B_size + SF_size * 2) * NUM_STAGES;

  auto k = kernel<BLOCK_N, NUM_STAGES>;
  cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
  dim3 grid(grid_n, grid_m, 1);
  k<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, (int)M, (int)N, (int)K);
  auto err = cudaGetLastError();
  TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err));
  return C;
}

template <int BLOCK_N, int NUM_STAGES>
__global__ __launch_bounds__(128 + 2 * WARP_SIZE)
void kernel_grouped(
  const CUtensorMap *A_tmaps,
  const CUtensorMap *B_tmaps,
  const uint64_t *SFA_ptrs,
  const uint64_t *SFB_ptrs,
  const uint64_t *C_ptrs,
  const int *Ms,
  const int *Ns,
  const int *Ks
) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_K = 256;

  const int gid = (int)blockIdx.z;
  const int M = Ms[gid];
  const int N = Ns[gid];
  const int K = Ks[gid];

  const int grid_m = (M + 127) / 128;
  const int grid_n = (N + BLOCK_N - 1) / BLOCK_N;

  const int bid_n = (int)blockIdx.x;
  const int bid_m = (int)blockIdx.y;
  if (bid_n >= grid_n || bid_m >= grid_m) return;

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

  const int tid = threadIdx.x;
  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;

  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 = K / BLOCK_K;

  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    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;
    }

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

    const int init_stage = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
    for (int iter_k = 0; iter_k < init_stage; 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 int MMA_N = BLOCK_N;
    constexpr int MMA_M = 128;
    constexpr uint32_t i_desc = (1U << 7U)
                              | (1U << 10U)
                              | ((uint32_t)MMA_N >> 3U << 17U)
                              | ((uint32_t)MMA_M >> 7U << 27U);

    auto make_desc_AB = [] __device__ (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 = [] __device__ (int addr) -> uint64_t {
      const int SBO = 8 * 16;
      return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
    };

    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;

      const uint64_t SFA_desc = make_desc_SF(0) + ((uint64_t)SFA_smem >> 4ULL);
      const uint64_t SFB_desc = make_desc_SF(0) + ((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);

          const int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4;
          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(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;");

    constexpr int WIDTH = (BLOCK_N < 64) ? BLOCK_N : 64;
    #pragma unroll
    for (int n0 = 0; n0 < BLOCK_N / WIDTH; n0++) {
      float tmp[WIDTH];
      tcgen05_ld_32x32bx64(tmp, warp_id * 32, n0 * WIDTH);
      asm volatile("tcgen05.wait::ld.sync.aligned;");

      #pragma unroll
      for (int i = 0; i < WIDTH; i += 2) {
        const int n_idx = off_n + n0 * WIDTH + i;
        const int m_idx = off_m + tid;
        if (m_idx < M) {
          if ((n_idx + 1) < N) {
            reinterpret_cast<half2 *>(C_ptr + m_idx * N + n_idx)[0] = __float22half2_rn({tmp[i + 0], tmp[i + 1]});
          } else if (n_idx < N) {
            C_ptr[m_idx * N + n_idx] = __float2half_rn(tmp[i + 0]);
          }
        }
      }
    }

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

void gemm_grouped_ptr(
  const at::Tensor& A_ptrs_h,
  const at::Tensor& B_ptrs_h,
  const at::Tensor& SFA_ptrs_h,
  const at::Tensor& SFB_ptrs_h,
  const at::Tensor& C_ptrs_h,
  const at::Tensor& Apads_h,
  const at::Tensor& Ms_h,
  const at::Tensor& Ns_h,
  const at::Tensor& Ks_h,
  int64_t dev
) {
  TORCH_CHECK(!A_ptrs_h.is_cuda() && !B_ptrs_h.is_cuda(), "ptr tensors must be CPU");
  TORCH_CHECK(A_ptrs_h.is_contiguous() && B_ptrs_h.is_contiguous(), "ptr tensors must be contiguous");
  TORCH_CHECK(A_ptrs_h.scalar_type() == at::kLong, "ptr tensors must be int64");
  TORCH_CHECK(B_ptrs_h.scalar_type() == at::kLong, "ptr tensors must be int64");
  TORCH_CHECK(SFA_ptrs_h.scalar_type() == at::kLong, "ptr tensors must be int64");
  TORCH_CHECK(SFB_ptrs_h.scalar_type() == at::kLong, "ptr tensors must be int64");
  TORCH_CHECK(C_ptrs_h.scalar_type() == at::kLong, "ptr tensors must be int64");
  TORCH_CHECK(Apads_h.scalar_type() == at::kLong, "size tensors must be int64");
  TORCH_CHECK(Ms_h.scalar_type() == at::kLong && Ns_h.scalar_type() == at::kLong && Ks_h.scalar_type() == at::kLong, "size tensors must be int64");

  const int64_t G = A_ptrs_h.numel();
  TORCH_CHECK(G > 0, "empty group");
  TORCH_CHECK(B_ptrs_h.numel() == G && SFA_ptrs_h.numel() == G && SFB_ptrs_h.numel() == G && C_ptrs_h.numel() == G, "len mismatch");
  TORCH_CHECK(Apads_h.numel() == G && Ms_h.numel() == G && Ns_h.numel() == G && Ks_h.numel() == G, "len mismatch");

  cudaSetDevice((int)dev);

  auto A_ptrs = A_ptrs_h.data_ptr<int64_t>();
  auto B_ptrs = B_ptrs_h.data_ptr<int64_t>();
  auto SFA_ptrs = SFA_ptrs_h.data_ptr<int64_t>();
  auto SFB_ptrs = SFB_ptrs_h.data_ptr<int64_t>();
  auto C_ptrs = C_ptrs_h.data_ptr<int64_t>();
  auto Apads = Apads_h.data_ptr<int64_t>();
  auto Ms = Ms_h.data_ptr<int64_t>();
  auto Ns = Ns_h.data_ptr<int64_t>();
  auto Ks = Ks_h.data_ptr<int64_t>();

  std::vector<CUtensorMap> A_tmaps((size_t)G);
  std::vector<CUtensorMap> B_tmaps((size_t)G);

  int max_grid_m = 0;
  int max_grid_n = 0;

  constexpr int BLOCK_N = 64;
  for (int i = 0; i < (int)G; i++) {
    const int M = (int)Ms[i];
    const int N = (int)Ns[i];
    const int K = (int)Ks[i];
    TORCH_CHECK((K % 256) == 0, "K must be multiple of 256");

    const char *A_ptr = reinterpret_cast<const char *>(A_ptrs[i]);
    const char *B_ptr = reinterpret_cast<const char *>(B_ptrs[i]);
    init_AB_tmap(&A_tmaps[(size_t)i], A_ptr, (uint64_t)Apads[i], (uint64_t)K, 128, 256);
    init_AB_tmap(&B_tmaps[(size_t)i], B_ptr, (uint64_t)N,        (uint64_t)K, 64,  256);

    const int grid_m = (M + 127) / 128;
    const int grid_n = (N + BLOCK_N - 1) / BLOCK_N;
    if (grid_m > max_grid_m) max_grid_m = grid_m;
    if (grid_n > max_grid_n) max_grid_n = grid_n;
  }

  at::TensorOptions opt_u8 = at::TensorOptions().device(at::kCUDA, (int)dev).dtype(at::kByte);
  at::TensorOptions opt_i64 = at::TensorOptions().device(at::kCUDA, (int)dev).dtype(at::kLong);
  at::TensorOptions opt_i32 = at::TensorOptions().device(at::kCUDA, (int)dev).dtype(at::kInt);

  auto A_tmaps_d = at::empty({G, (int64_t)sizeof(CUtensorMap)}, opt_u8);
  auto B_tmaps_d = at::empty({G, (int64_t)sizeof(CUtensorMap)}, opt_u8);

  auto SFA_ptrs_d = at::empty({G}, opt_i64);
  auto SFB_ptrs_d = at::empty({G}, opt_i64);
  auto C_ptrs_d = at::empty({G}, opt_i64);

  auto Ms_d = at::empty({G}, opt_i32);
  auto Ns_d = at::empty({G}, opt_i32);
  auto Ks_d = at::empty({G}, opt_i32);

  TORCH_CHECK(cudaMemcpy(A_tmaps_d.data_ptr(), A_tmaps.data(), (size_t)G * sizeof(CUtensorMap), cudaMemcpyHostToDevice) == cudaSuccess, "memcpy fail");
  TORCH_CHECK(cudaMemcpy(B_tmaps_d.data_ptr(), B_tmaps.data(), (size_t)G * sizeof(CUtensorMap), cudaMemcpyHostToDevice) == cudaSuccess, "memcpy fail");

  TORCH_CHECK(cudaMemcpy(SFA_ptrs_d.data_ptr<int64_t>(), SFA_ptrs, (size_t)G * sizeof(int64_t), cudaMemcpyHostToDevice) == cudaSuccess, "memcpy fail");
  TORCH_CHECK(cudaMemcpy(SFB_ptrs_d.data_ptr<int64_t>(), SFB_ptrs, (size_t)G * sizeof(int64_t), cudaMemcpyHostToDevice) == cudaSuccess, "memcpy fail");
  TORCH_CHECK(cudaMemcpy(C_ptrs_d.data_ptr<int64_t>(), C_ptrs, (size_t)G * sizeof(int64_t), cudaMemcpyHostToDevice) == cudaSuccess, "memcpy fail");

  std::vector<int> Ms_i32((size_t)G);
  std::vector<int> Ns_i32((size_t)G);
  std::vector<int> Ks_i32((size_t)G);
  for (int i = 0; i < (int)G; i++) {
    Ms_i32[(size_t)i] = (int)Ms[i];
    Ns_i32[(size_t)i] = (int)Ns[i];
    Ks_i32[(size_t)i] = (int)Ks[i];
  }

  TORCH_CHECK(cudaMemcpy(Ms_d.data_ptr<int>(), Ms_i32.data(), (size_t)G * sizeof(int), cudaMemcpyHostToDevice) == cudaSuccess, "memcpy fail");
  TORCH_CHECK(cudaMemcpy(Ns_d.data_ptr<int>(), Ns_i32.data(), (size_t)G * sizeof(int), cudaMemcpyHostToDevice) == cudaSuccess, "memcpy fail");
  TORCH_CHECK(cudaMemcpy(Ks_d.data_ptr<int>(), Ks_i32.data(), (size_t)G * sizeof(int), cudaMemcpyHostToDevice) == cudaSuccess, "memcpy fail");

  constexpr int NUM_STAGES = 4;
  const int tb_size = 128 + 2 * WARP_SIZE;
  const int A_size = 128 * 256 / 2;
  const int B_size = BLOCK_N * 256 / 2;
  const int SF_size = 128 * 256 / 16;
  const int smem_size = (A_size + B_size + SF_size * 2) * NUM_STAGES;

  auto k = kernel_grouped<BLOCK_N, NUM_STAGES>;
  cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
  dim3 grid((unsigned)max_grid_n, (unsigned)max_grid_m, (unsigned)G);
  k<<<grid, tb_size, smem_size>>>(
    reinterpret_cast<const CUtensorMap *>(A_tmaps_d.data_ptr()),
    reinterpret_cast<const CUtensorMap *>(B_tmaps_d.data_ptr()),
    reinterpret_cast<const uint64_t *>(SFA_ptrs_d.data_ptr<int64_t>()),
    reinterpret_cast<const uint64_t *>(SFB_ptrs_d.data_ptr<int64_t>()),
    reinterpret_cast<const uint64_t *>(C_ptrs_d.data_ptr<int64_t>()),
    Ms_d.data_ptr<int>(),
    Ns_d.data_ptr<int>(),
    Ks_d.data_ptr<int>()
  );
  auto err = cudaGetLastError();
  TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err));
}

TORCH_LIBRARY(nvfp4_group_gemm_opt, m) {
  m.def("gemm(Tensor A, Tensor B, Tensor SFA, Tensor SFB, Tensor(a!) C, int M, int N, int K) -> Tensor");
  m.impl("gemm", &gemm);
  m.def("gemm_grouped_ptr(Tensor A_ptrs, Tensor B_ptrs, Tensor SFA_ptrs, Tensor SFB_ptrs, Tensor C_ptrs, Tensor Apads, Tensor Ms, Tensor Ns, Tensor Ks, int dev) -> ()");
  m.impl("gemm_grouped_ptr", &gemm_grouped_ptr);
}
"""

    build_dir = os.path.join(os.path.dirname(__file__), ".build_nvfp4_group_gemm_opt")
    os.makedirs(build_dir, exist_ok=True)

    load_inline(
        name="nvfp4_group_gemm_opt_ext",
        cpp_sources="",
        cuda_sources=cuda_src,
        functions=None,
        extra_cflags=["-O3"],
        extra_cuda_cflags=[
            "-O3",
            "-gencode=arch=compute_100a,code=sm_100a",
            "--use_fast_math",
            "--expt-extended-lambda",
            "--expt-relaxed-constexpr",
            "--relocatable-device-code=false",
            "-std=c++17",
            "-lineinfo",
        ],
        extra_ldflags=["-lcuda"],
        with_cuda=True,
        is_python_module=False,
        no_implicit_headers=True,
        build_directory=build_dir,
        verbose=False,
    )

    _EXT_READY = True


def _as_u8(x: torch.Tensor) -> torch.Tensor:
    if x.dtype == torch.uint8:
        return x
    if x.element_size() != 1:
        raise RuntimeError("packed tensor must have 1-byte elements")
    return x.view(torch.uint8)


def _reorder_scale_from_raw(scale_u8_2d: torch.Tensor, rows_pad: int) -> torch.Tensor:
    if scale_u8_2d.dim() != 2:
        raise RuntimeError("scale must be 2D")
    rows = int(scale_u8_2d.size(0))
    k16 = int(scale_u8_2d.size(1))
    if (k16 % 4) != 0:
        raise RuntimeError("K//16 must be multiple of 4")
    if (rows_pad % 128) != 0:
        raise RuntimeError("rows_pad must be multiple of 128")
    blk_m = rows_pad // 128
    blk_k = k16 // 4
    buf = torch.empty((rows_pad, k16), device=scale_u8_2d.device, dtype=torch.uint8)
    buf[:rows].copy_(scale_u8_2d)
    v = buf.view(blk_m, 32, 4, blk_k, 4).permute(0, 3, 1, 2, 4).contiguous()
    return v


def _get_scratch_a(device: torch.device, m_pad: int, k2: int) -> torch.Tensor:
    key = (int(device.index), int(m_pad), int(k2))
    buf = _SCRATCH_A.get(key)
    if buf is None or (not buf.is_cuda) or buf.numel() != (m_pad * k2):
        buf = torch.empty((m_pad, k2, 1), device=device, dtype=torch.uint8)
        _SCRATCH_A[key] = buf
    return buf


def custom_kernel(data):
    abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data
    _load_ext()
    gemm = torch.ops.nvfp4_group_gemm_opt.gemm
    gemm_grouped_ptr = torch.ops.nvfp4_group_gemm_opt.gemm_grouped_ptr

    g = len(problem_sizes)
    all_l1 = True
    for i in range(g):
        if int(problem_sizes[i][3]) != 1:
            all_l1 = False
            break
    all_l1 = False

    outs: List[torch.Tensor] = []

    if all_l1:
        dev = int(abc_tensors[0][0].device.index)

        keep_a: List[torch.Tensor] = []
        keep_b: List[torch.Tensor] = []
        keep_sfa: List[torch.Tensor] = []
        keep_sfb: List[torch.Tensor] = []
        keep_c: List[torch.Tensor] = []

        a_ptrs: List[int] = []
        b_ptrs: List[int] = []
        sfa_ptrs: List[int] = []
        sfb_ptrs: List[int] = []
        c_ptrs: List[int] = []
        apads: List[int] = []
        ms: List[int] = []
        ns: List[int] = []
        ks: List[int] = []

        c_out_list: List[torch.Tensor] = []
        c_tmp_list: List[torch.Tensor] = []

        for i in range(g):
            a, b, c = abc_tensors[i]
            sfa, sfb = sfasfb_tensors[i]
            sfa_p, sfb_p = sfasfb_reordered_tensors[i]
            m, n, k, _l = problem_sizes[i]

            m_int = int(m)
            n_int = int(n)
            k_int = int(k)

            c_out = c
            if not c_out.is_contiguous():
                c_tmp = torch.empty_like(c_out, memory_format=torch.contiguous_format)
            else:
                c_tmp = c_out

            a_u8 = _as_u8(a).contiguous()
            b_u8 = _as_u8(b).contiguous()

            m_pad = ((m_int + 127) // 128) * 128
            if a_u8.size(0) != m_pad:
                a_pad = torch.empty((m_pad, a_u8.size(1), 1), device=a_u8.device, dtype=torch.uint8)
                a_pad[: a_u8.size(0)].copy_(a_u8)
            else:
                a_pad = a_u8

            n_pad = ((n_int + 127) // 128) * 128
            ok_sfp = (
                sfa_p.is_cuda
                and sfb_p.is_cuda
                and (sfa_p.dim() == 6)
                and (sfb_p.dim() == 6)
                and (sfa_p.element_size() == 1)
                and (sfb_p.element_size() == 1)
                and (int(sfa_p.storage_offset()) == 0)
                and (int(sfb_p.storage_offset()) == 0)
                and sfa_p.permute(2, 4, 0, 1, 3, 5).is_contiguous()
                and sfb_p.permute(2, 4, 0, 1, 3, 5).is_contiguous()
            )
            if ok_sfp:
                sfa_arg = sfa_p
                sfb_arg = sfb_p
            else:
                sfa2 = _as_u8(sfa[..., 0]).contiguous()
                sfb2 = _as_u8(sfb[..., 0]).contiguous()
                sfa_arg = _reorder_scale_from_raw(sfa2, m_pad)
                sfb_arg = _reorder_scale_from_raw(sfb2, n_pad)

            keep_a.append(a_pad)
            keep_b.append(b_u8)
            keep_sfa.append(sfa_arg)
            keep_sfb.append(sfb_arg)
            keep_c.append(c_tmp)

            a_ptrs.append(int(a_pad.data_ptr()))
            b_ptrs.append(int(b_u8.data_ptr()))
            sfa_ptrs.append(int(sfa_arg.data_ptr()))
            sfb_ptrs.append(int(sfb_arg.data_ptr()))
            c_ptrs.append(int(c_tmp.data_ptr()))
            apads.append(int(a_pad.size(0)))
            ms.append(m_int)
            ns.append(n_int)
            ks.append(k_int)

            c_out_list.append(c_out)
            c_tmp_list.append(c_tmp)

        gemm_grouped_ptr(
            torch.tensor(a_ptrs, dtype=torch.long),
            torch.tensor(b_ptrs, dtype=torch.long),
            torch.tensor(sfa_ptrs, dtype=torch.long),
            torch.tensor(sfb_ptrs, dtype=torch.long),
            torch.tensor(c_ptrs, dtype=torch.long),
            torch.tensor(apads, dtype=torch.long),
            torch.tensor(ms, dtype=torch.long),
            torch.tensor(ns, dtype=torch.long),
            torch.tensor(ks, dtype=torch.long),
            dev,
        )

        for i in range(g):
            c_out = c_out_list[i]
            c_tmp = c_tmp_list[i]
            if c_tmp is not c_out:
                c_out.copy_(c_tmp)
            outs.append(c_out)

        return outs

    for i in range(g):
        a, b, c = abc_tensors[i]
        sfa, sfb = sfasfb_tensors[i]
        sfa_p, sfb_p = sfasfb_reordered_tensors[i]
        m, n, k, l = problem_sizes[i]

        m_int = int(m)
        n_int = int(n)
        k_int = int(k)
        l_int = int(l)

        c_out = c
        if not c_out.is_contiguous():
            c_tmp = torch.empty_like(c_out, memory_format=torch.contiguous_format)
        else:
            c_tmp = c_out

        if l_int == 1:
            a_u8 = _as_u8(a).contiguous()
            b_u8 = _as_u8(b).contiguous()

            m_pad = ((m_int + 127) // 128) * 128
            if a_u8.size(0) != m_pad:
                a_pad = _get_scratch_a(a_u8.device, m_pad, int(a_u8.size(1)))
                a_pad[: a_u8.size(0)].copy_(a_u8)
            else:
                a_pad = a_u8

            n_pad = ((n_int + 127) // 128) * 128
            ok_sfp = (
                sfa_p.is_cuda
                and sfb_p.is_cuda
                and (sfa_p.dim() == 6)
                and (sfb_p.dim() == 6)
                and (sfa_p.element_size() == 1)
                and (sfb_p.element_size() == 1)
                and (int(sfa_p.storage_offset()) == 0)
                and (int(sfb_p.storage_offset()) == 0)
                and sfa_p.permute(2, 4, 0, 1, 3, 5).is_contiguous()
                and sfb_p.permute(2, 4, 0, 1, 3, 5).is_contiguous()
            )
            if ok_sfp:
                sfa_arg = sfa_p
                sfb_arg = sfb_p
            else:
                sfa2 = _as_u8(sfa[..., 0]).contiguous()
                sfb2 = _as_u8(sfb[..., 0]).contiguous()
                sfa_arg = _reorder_scale_from_raw(sfa2, m_pad)
                sfb_arg = _reorder_scale_from_raw(sfb2, n_pad)

            gemm(a_pad, b_u8, sfa_arg, sfb_arg, c_tmp, m_int, n_int, k_int)
        else:
            for li in range(l_int):
                a2 = _as_u8(a[..., li]).contiguous().unsqueeze(-1)
                b2 = _as_u8(b[..., li]).contiguous().unsqueeze(-1)

                m_pad = ((m_int + 127) // 128) * 128
                if a2.size(0) != m_pad:
                    a_pad = _get_scratch_a(a2.device, m_pad, int(a2.size(1)))
                    a_pad[: a2.size(0)].copy_(a2)
                else:
                    a_pad = a2

                sfa2 = _as_u8(sfa[..., li]).contiguous()
                sfb2 = _as_u8(sfb[..., li]).contiguous()
                sfa_r = _reorder_scale_from_raw(sfa2, m_pad)
                sfb_r = _reorder_scale_from_raw(sfb2, ((n_int + 127) // 128) * 128)

                c2 = torch.empty((m_int, n_int, 1), device=c_tmp.device, dtype=torch.float16)
                gemm(a_pad, b2, sfa_r, sfb_r, c2, m_int, n_int, k_int)
                c_tmp[..., li].copy_(c2[..., 0])

        if c_tmp is not c_out:
            c_out.copy_(c_tmp)
        outs.append(c_out)

    return outs


__all__ = ["custom_kernel"]