submission 409474

  from __future__ import annotations
   
  import os
- from typing import List
+ from typing import Dict, List, Tuple
   
  import torch
  from torch.utils.cpp_extension import load_inline
   
- _EXT_READY = False
- _SCRATCH_A: dict = {}
+ _MOD = None
+ _PAD_CACHE: Dict[Tuple[int, int, int, int], torch.Tensor] = {}
   
   
- def _load_ext() -> None:
-     global _EXT_READY
-     if _EXT_READY:
-         return
+ def _get_mod():
+     global _MOD
+     if _MOD is not None:
+         return _MOD
   
      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;
+ constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;
+ constexpr uint64_t EVICT_FIRST  = 0x12F0000000000000;
+ constexpr uint64_t EVICT_LAST   = 0x14F0000000000000;
   
- __device__ __forceinline__ int64_t globaltimer() {
-   int64_t t;
-   asm volatile("mov.u64 %0, %globaltimer;" : "=l"(t) :: "memory");
-   return t;
- }
+ __device__ inline constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; }
   
- __device__ __forceinline__ constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3FFFFULL) >> 4ULL; }
-  
- __device__ __forceinline__ uint32_t elect_sync() {
+ __device__ inline uint32_t elect_sync() {
    uint32_t pred = 0;
    asm volatile(
      "{\n\t"
  ⋯ 7 unchanged lines
    return pred;
  }
   
- __device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) {
+ __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__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) {
+ __device__ void mbarrier_wait(int mbar_addr, int phase) {
    uint32_t ticks = 0x989680;
    asm volatile(
      "{\n\t"
  ⋯ 8 unchanged lines
    );
  }
   
- __device__ __forceinline__ void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
+ __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;"
  ⋯ 1 unchanged lines
    );
  }
   
- __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) {
+ __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;"
  ⋯ 2 unchanged lines
    );
  }
   
- __device__ __forceinline__ void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
+ __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__ __forceinline__ void tcgen05_mma_nvfp4(
+ __device__ inline void tcgen05_mma_nvfp4(
    uint64_t a_desc,
    uint64_t b_desc,
    uint32_t i_desc,
  ⋯ 14 unchanged lines
  }
   
  struct SHAPE {
-   static constexpr char _32x32b[] = ".32x32b";
+   static constexpr char _32x32b[]  = ".32x32b";
  };
  struct NUM {
-   static constexpr char x64[] = ".x64";
+   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) {
+ template <const char *SHAPE_, const char *NUM_>
+ __device__ inline 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, "
  ⋯ 12 unchanged lines
        "=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));
+     : "r"((row << 16) | col), "C"(SHAPE_), "C"(NUM_)
+   );
  }
   
- __device__ __forceinline__ void tcgen05_ld_32x32bx64(float *tmp, int row, int col) {
-   tcgen05_ld_64regs<SHAPE::_32x32b, NUM::x64>(tmp, row, col);
- }
+ __device__ inline 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) {
+ static void check_cu(CUresult err) {
    if (err == CUDA_SUCCESS) return;
-   const char *msg = "unknown";
-   cuGetErrorString(err, &msg);
-   TORCH_CHECK(false, msg);
+   const char *error_msg_ptr;
+   if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS)
+     error_msg_ptr = "unable to get error string";
+   TORCH_CHECK(false, "cuTensorMapEncodeTiled error: ", error_msg_ptr);
  }
   
- static __forceinline__ void init_AB_tmap(
+ static 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
+   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};
  ⋯ 18 unchanged lines
    check_cu(err);
  }
   
- template <int BLOCK_N, int NUM_STAGES>
- __global__ __launch_bounds__(128 + 2 * WARP_SIZE)
+ template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
+ __global__ __launch_bounds__(BLOCK_M + 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
+   int K,
+   int M, int N,
+   int N_valid
  ) {
-   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 bid = blockIdx.x;
   
    const int lane_id = tid % WARP_SIZE;
    const int warp_id = tid / WARP_SIZE;
   
+   const int grid_m = M / BLOCK_M;
+   const int grid_n = N / BLOCK_N;
+   const int bid_m = bid / grid_n;
+   const int bid_n = bid - bid_m * grid_n;
+  
    const int off_m = bid_m * BLOCK_M;
    const int off_n = bid_n * BLOCK_N;
   
  ⋯ 1 unchanged lines
   
    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 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;
  ⋯ 8 unchanged lines
    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);
+     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));
  ⋯ 30 unchanged lines
        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");
+                   :: "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);
+     const int prefetch = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
+     for (int iter_k = 0; iter_k < prefetch; iter_k++)
+       issue_tma(iter_k, iter_k);
   
-     for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
+     for (int iter_k = prefetch; 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);
  ⋯ 7 unchanged lines
                                | ((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;
  ⋯ 4 unchanged lines
        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);
+       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);
+       };
   
-       #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);
+       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 kk = 0; kk < BLOCK_K / MMA_K; kk++) {
+         uint64_t sfa_desc = SFA_desc + (uint64_t)kk * (512ULL >> 4ULL);
+         uint64_t sfb_desc = SFB_desc + (uint64_t)kk * (512ULL >> 4ULL);
+         tcgen05_cp_nvfp4(SFA_tmem + kk * 4, sfa_desc);
+         tcgen05_cp_nvfp4(SFB_tmem + kk * 4, sfb_desc);
        }
   
-       #pragma unroll
-       for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
-         #pragma unroll
+       for (int k1 = 0; k1 < BLOCK_K / 256; k1++)
          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;
+           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(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");
+                   :: "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");
+                 :: "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++) {
+     for (int n_it = 0; n_it < BLOCK_N / WIDTH; n_it++) {
        float tmp[WIDTH];
-       tcgen05_ld_32x32bx64(tmp, warp_id * 32, n0 * WIDTH);
+       tcgen05_ld_32x32bx64(tmp, warp_id * 32, n_it * 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]);
-           }
+       for (int i = 0; i < WIDTH; i++) {
+         const int row = off_n + n_it * WIDTH + i;
+         const int col = off_m + tid;
+         if (row < N_valid) {
+           C_ptr[row * M + col] = __float2half(tmp[i]);
          }
        }
      }
   
      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));
+     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(
+ template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
+ static void gemm_launch(
    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
+   int N_valid,
+   int 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");
+   const int M = A.size(0);
+   const int N = B.size(0);
   
-   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");
+   auto A_ptr   = reinterpret_cast<const char *>(A.data_ptr());
+   auto B_ptr   = reinterpret_cast<const char *>(B.data_ptr());
+   auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
+   auto SFB_ptr = reinterpret_cast<const char *>(SFB.data_ptr());
+   auto C_ptr   = reinterpret_cast<half *>(C.data_ptr());
   
-   const int64_t Apad = A.size(0);
-   TORCH_CHECK(Apad >= M, "A pad too small");
+   int new_M = M;
+   int new_N = N;
+   std::swap(A_ptr, B_ptr);
+   std::swap(SFA_ptr, SFB_ptr);
+   std::swap(new_M, new_N);
   
-   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);
+   init_AB_tmap(&A_tmap, A_ptr, new_M, K, BLOCK_M, BLOCK_K);
+   init_AB_tmap(&B_tmap, B_ptr, new_N, K, BLOCK_N, BLOCK_K);
   
-   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);
+   dim3 grid((new_M / BLOCK_M) * (new_N / BLOCK_N));
+   int tb_size = BLOCK_M + 2 * WARP_SIZE;
+   int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
+   int SFAB_size = 128 * (BLOCK_K / 16) * 2;
+   int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
   
-   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;
+   auto this_kernel = kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
+   if (smem_size > 48'000)
+     cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+   this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, K, new_M, new_N, N_valid);
  }
   
- 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
+ static 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_valid
  ) {
-   constexpr int BLOCK_M = 128;
-   constexpr int BLOCK_K = 256;
+   TORCH_CHECK(A.is_cuda() && B.is_cuda() && SFA.is_cuda() && SFB.is_cuda() && C.is_cuda(), "cuda only");
+   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, "only L=1 per call");
+   TORCH_CHECK(A.is_contiguous() && B.is_contiguous() && C.is_contiguous(), "A/B/C must be contiguous");
+   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");
   
-   const int gid = (int)blockIdx.z;
-   const int M = Ms[gid];
-   const int N = Ns[gid];
-   const int K = Ks[gid];
+   const int K = (int)A.size(1) * 2;
+   const int N = (int)B.size(0);
+   const int M_pad = (int)A.size(0);
+   TORCH_CHECK((int)C.size(0) == (int)m_valid, "C M mismatch");
+   TORCH_CHECK((int)C.size(1) == N, "C N mismatch");
+   TORCH_CHECK(m_valid >= 0 && m_valid <= M_pad, "m_valid out of range");
   
-   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));
+   TORCH_CHECK((K % 256) == 0, "K must be multiple of 256");
+   if (K >= 2048) {
+     gemm_launch<128, 64, 256, 8>(A, B, SFA, SFB, C, (int)m_valid, K);
+   } else {
+     gemm_launch<128, 64, 256, 6>(A, B, SFA, SFB, C, (int)m_valid, K);
    }
-   __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));
-   }
+   return C;
  }
   
- 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");
+ TORCH_LIBRARY(nvfp4_group_gemm_mod, m) {
+   m.def("gemm(Tensor A, Tensor B, Tensor SFA, Tensor SFB, Tensor(a!) C, int m_valid) -> 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")
+     build_dir = os.path.join(os.path.dirname(__file__), ".build_nvfp4_group_gemm")
      os.makedirs(build_dir, exist_ok=True)
-  
      load_inline(
-         name="nvfp4_group_gemm_opt_ext",
+         name="nvfp4_group_gemm_ext",
          cpp_sources="",
          cuda_sources=cuda_src,
          functions=None,
-         extra_cflags=["-O3"],
+         with_cuda=True,
          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_cflags=["-O3"],
          extra_ldflags=["-lcuda"],
-         with_cuda=True,
          is_python_module=False,
          no_implicit_headers=True,
          build_directory=build_dir,
          verbose=False,
      )
   
-     _EXT_READY = True
+     _MOD = torch.ops.nvfp4_group_gemm_mod
+     return _MOD
   
   
- 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 _pad_a(a: torch.Tensor, m_valid: int) -> Tuple[torch.Tensor, int]:
+     if a.dim() != 3 or a.size(2) != 1:
+         raise RuntimeError("only [M, K//2, 1] supported per call")
+     if not a.is_cuda:
+         raise RuntimeError("cuda only")
+     if a.element_size() != 1:
+         raise RuntimeError("packed fp4 must have 1-byte elements")
+     if not a.is_contiguous():
+         a = a.contiguous()
   
+     k_half = int(a.size(1))
+     m_pad = (int(m_valid) + 64 - 1) // 64 * 64
+     if m_pad == int(m_valid):
+         return a, m_pad
+     key = (a.device.index if a.device.index is not None else -1, a.dtype, m_pad, k_half)
   
- 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
+     buf = _PAD_CACHE.get(key)
+     if buf is None or buf.numel() != m_pad * k_half:
+         buf = torch.empty((m_pad, k_half, 1), device=a.device, dtype=a.dtype)
+         _PAD_CACHE[key] = buf
+     buf[:m_valid].copy_(a[:m_valid])
+     return buf, m_pad
   
   
- 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
+     abc_tensors, _sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data
+     mod = _get_mod()
+     outs: List[torch.Tensor] = []
   
      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, n, _k, l = problem_sizes[i]
   
-         m_int = int(m)
-         n_int = int(n)
-         k_int = int(k)
-         l_int = int(l)
+         m = int(m)
+         l = 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 != 1:
+             raise RuntimeError("only L=1 is supported")
   
-         if l_int == 1:
-             a_u8 = _as_u8(a).contiguous()
-             b_u8 = _as_u8(b).contiguous()
+         a_pad, _ = _pad_a(a, m)
+         mod.gemm(a_pad, b, sfa_p, sfb_p, c, m)
+         outs.append(c)
   
-             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