submission 494881

#!POPCORN leaderboard nvfp4_group_gemm
#!POPCORN gpu B200

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

# ----------------------------
# CUDA / C++ extension
# ----------------------------

CUDA_SRC_COMMON = r"""
#include <cuda.h>
#include <cudaTypedefs.h>
#include <cuda_fp16.h>
#include <stdint.h>

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

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

// cache hint (from CUTLASS copy_sm90_desc)
constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;
constexpr uint64_t EVICT_FIRST  = 0x12F0000000000000;
constexpr uint64_t EVICT_LAST   = 0x14F0000000000000;

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

// elect one thread in warp
__device__ uint32_t elect_sync() {
  uint32_t pred = 0;
  asm volatile(
    "{\n\t"
    ".reg .pred %%px;\n\t"
    "elect.sync _|%%px, %1;\n\t"
    "@%%px mov.s32 %0, 1;\n\t"
    "}"
    : "+r"(pred)
    : "r"(0xFFFFFFFF)
  );
  return pred;
}

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

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

// TMA bulk copy (1D)
__device__ inline
void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
  asm volatile(
    "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint "
    "[%0], [%1], %2, [%3], %4;"
    :: "r"(dst), "l"(src), "r"(size), "r"(mbar_addr), "l"(cache_policy)
  );
}

// TMA tensor 3D copy
__device__ inline
void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
  asm volatile(
    "cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
    "[%0], [%1, {%2, %3, %4}], [%5], %6;"
    :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy)
    : "memory"
  );
}

// tcgen05 cp scale (nvfp4)
__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));
}

// tcgen05 mma (nvfp4 blockscale)
__device__ inline
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)
  );
}

// Minimal tcgen05 ld: need 32 regs for 16x256b.x8
struct SHAPE {
  static constexpr char _16x256b[] = ".16x256b";
};

struct NUM {
  static constexpr char x8[] = ".x8";
};

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

__device__ inline
void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
  tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
}

void check_cu(CUresult err) {
  if (err == CUDA_SUCCESS) return;
  const char *error_msg_ptr;
  if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS)
    error_msg_ptr = "unable to get error string";
  TORCH_CHECK(false, "CUDA driver error: ", error_msg_ptr);
}

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
) {
  // NOTE: matches champion gemm's fp4 tensor map (16U4)
  constexpr uint32_t rank = 3;
  uint64_t globalDim[rank]       = {256, global_height, global_width / 256};
  uint64_t globalStrides[rank-1] = {global_width / 2, 128};  // bytes
  uint32_t boxDim[rank]          = {256, shared_height, shared_width / 256};
  uint32_t elementStrides[rank]  = {1, 1, 1};

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

CUDA_SRC_GEMM = r"""
template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
__global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
void kernel_gemm(
  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
) {
  const int tid = threadIdx.x;
  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 % grid_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; // BLOCK_M=128 => 6 warps

  // smem
  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;  // fp4 packed
  constexpr int B_size   = BLOCK_N * BLOCK_K / 2;
  constexpr int SFA_size = 128 * BLOCK_K / 16;     // fp8 bytes, always 128 rows
  constexpr int SFB_size = 128 * BLOCK_K / 16;
  constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;

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

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

  constexpr int num_iters = K / BLOCK_K;

  // TMA warp
  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    constexpr uint64_t cache_A = EVICT_LAST;
    constexpr uint64_t 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;

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

      // scale layout assumed: [mn/128, rest_k, 32,4,4] contiguous in memory (512 bytes each block)
      const int rest_k = K / 16 / 4; // = K/64
      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++)
      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);
    }
  }
  // MMA warp
  else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    // fp4 MMA uses MMA_M=128 always
    constexpr uint32_t i_desc = (1U << 7U)   // atype=E2M1
                             | (1U << 10U)  // btype=E2M1
                             | ((uint32_t)BLOCK_N >> 3U << 17U)  // MMA_N
                             | ((uint32_t)128     >> 7U << 27U); // MMA_M

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

      // smem->tmem for scales
      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);
      }

      // mma over BLOCK_K
      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);

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

    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                 :: "r"(mainloop_mbar_addr) : "memory");
  }
  // epilogue warps (write row-major C: [M,N])
  else if (tid < BLOCK_M) {
    mbarrier_wait(mainloop_mbar_addr, 0);
    asm volatile("tcgen05.fence::after_thread_sync;");

    // C is row-major (N-major in their naming)
    for (int m = 0; m < 32 / 16; m++) {
      float tmp[BLOCK_N / 2]; // BLOCK_N=64 => 32 floats
      tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 0);
      asm volatile("tcgen05.wait::ld.sync.aligned;");

      for (int i = 0; i < BLOCK_N / 8; i++) {
        const int row = off_m + warp_id * 32 + m * 16 + lane_id / 4;
        const int col = off_n + i * 8 + (lane_id % 4) * 2;

        reinterpret_cast<half2 *>(C_ptr + (row + 0) * N + col)[0] =
          __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
        reinterpret_cast<half2 *>(C_ptr + (row + 8) * N + col)[0] =
          __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
      }
    }

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

template <int K>
static inline void launch_one(
  const at::Tensor& A,
  const at::Tensor& B,
  const at::Tensor& SFA_blk,
  const at::Tensor& SFB_blk,
        at::Tensor& C
) {
  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_N = 64;
  constexpr int BLOCK_K = 256;
  constexpr int NUM_STAGES = 6;

  const int M = A.size(0);
  const int N = B.size(0);

  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_blk.data_ptr());
  auto SFB_ptr = reinterpret_cast<const char*>(SFB_blk.data_ptr());
  auto C_ptr   = reinterpret_cast<half*>(C.data_ptr());

  CUtensorMap A_tmap, B_tmap;
  init_AB_tmap(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
  init_AB_tmap(&B_tmap, B_ptr, N, K, BLOCK_N, BLOCK_K);

  const int grid = (M / BLOCK_M) * (N / BLOCK_N);
  const int tb_size = 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;

  auto ker = kernel_gemm<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
  if (smem_size > 48'000)
    cudaFuncSetAttribute(ker, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);

  ker<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, M, N);
}

at::Tensor gemm(
  const at::Tensor& A,
  const at::Tensor& B,
  const at::Tensor& SFA_blk,
  const at::Tensor& SFB_blk,
        at::Tensor& C
) {
  const int K = A.size(1) * 2;

  if (false) {}
  else if (K == 7168)  launch_one<7168>(A, B, SFA_blk, SFB_blk, C);
  else if (K == 4096)  launch_one<4096>(A, B, SFA_blk, SFB_blk, C);
  else if (K == 2048)  launch_one<2048>(A, B, SFA_blk, SFB_blk, C);
  else if (K == 1536)  launch_one<1536>(A, B, SFA_blk, SFB_blk, C);
  else if (K == 2304)  launch_one<2304>(A, B, SFA_blk, SFB_blk, C);
  else if (K == 512)   launch_one<512>(A, B, SFA_blk, SFB_blk, C);
  else if (K == 256)   launch_one<256>(A, B, SFA_blk, SFB_blk, C);
  else {
    // leave to python fallback for correctness
  }

  return C;
}

TORCH_LIBRARY(group_gemm_mod, m) {
  m.def("gemm(Tensor A, Tensor B, Tensor SFA_blk, Tensor SFB_blk, Tensor(a!) C) -> Tensor");
  m.impl("gemm", &gemm);
}
"""

load_inline(
    name="group_gemm_ext",
    cpp_sources="",
    cuda_sources=CUDA_SRC_COMMON + CUDA_SRC_GEMM,
    functions=None,
    extra_cuda_cflags=[
        "-O3",
        "-gencode=arch=compute_100a,code=sm_100a",
        "--use_fast_math",
        "--expt-relaxed-constexpr",
        "--relocatable-device-code=false",
        "-lineinfo",
        "-Xptxas=-v",
    ],
    extra_ldflags=["-lcuda"],
    with_cuda=True,
    verbose=False,
    is_python_module=False,
    no_implicit_headers=True,
)

gemm = torch.ops.group_gemm_mod.gemm

# ----------------------------
# Python helpers
# ----------------------------

def ceil_div(a: int, b: int) -> int:
    return (a + b - 1) // b

@torch.no_grad()
def reorder_sf_blocked_fast(sf_mn_k16: torch.Tensor, mn: int, k: int, mn_pad: int) -> torch.Tensor:
    """
    FAST reshape/permute version.
    input:  [mn, k//16]
    output: [mn_pad//128, (k//16)//4, 32, 4, 4]
    Layout mapping (matches your original scatter):
      mm   = i//128
      mm32 = i%32
      mm4  = (i%128)//32
      kk   = j//4
      kk4  = j%4
    => out[mm, kk, mm32, mm4, kk4] = in[i, j]
    """
    assert sf_mn_k16.dim() == 2
    assert sf_mn_k16.shape[0] == mn
    sf_k = k // 16
    assert sf_mn_k16.shape[1] == sf_k
    assert (mn_pad % 128) == 0
    assert (sf_k % 4) == 0

    device = sf_mn_k16.device
    dtype = sf_mn_k16.dtype

    # pad rows to mn_pad (cheap, contiguous)
    if mn_pad == mn:
        sf_pad = sf_mn_k16
    else:
        sf_pad = torch.zeros((mn_pad, sf_k), device=device, dtype=dtype)
        sf_pad[:mn, :] = sf_mn_k16

    rest_m = mn_pad // 128
    rest_k = sf_k // 4

    # reshape: [rest_m, 128, rest_k, 4]
    # then split 128 -> [4, 32] with order (mm4, mm32)
    # then permute to [rest_m, rest_k, 32, 4, 4]
    out = (
        sf_pad
        .view(rest_m, 128, rest_k, 4)
        .view(rest_m, 4, 32, rest_k, 4)
        .permute(0, 3, 2, 1, 4)
        .contiguous()
    )
    return out

@torch.no_grad()
def pad_fp4_rows(a_fp4: torch.Tensor, m: int, mn_pad: int) -> torch.Tensor:
    """
    Pad fp4 packed tensor [M, K//2, L] on rows to mn_pad with zeros (bitwise),
    returning float4_e2m1fn_x2 tensor.
    """
    a_u8 = a_fp4.contiguous().view(torch.uint8)
    out_u8 = torch.zeros((mn_pad, a_u8.shape[1], a_u8.shape[2]), device=a_u8.device, dtype=torch.uint8)
    out_u8[:m, :, :] = a_u8
    return out_u8.view(torch.float4_e2m1fn_x2)

@torch.no_grad()
def torch_scaled_mm_fallback(a_fp4, b_fp4, sfa, sfb, c_out, m, n, k, l):
    sf_vec_size = 16

    def to_block(input_matrix):
        rows, cols = input_matrix.shape
        n_row_blocks = ceil_div(rows, 128)
        n_col_blocks = ceil_div(cols, 4)
        padded_rows = n_row_blocks * 128
        padded_cols = n_col_blocks * 4
        if padded_rows != rows or padded_cols != cols:
            padded = torch.nn.functional.pad(
                input_matrix, (0, padded_cols - cols, 0, padded_rows - rows),
                mode="constant", value=0
            )
        else:
            padded = input_matrix
        blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
        rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
        return rearranged.flatten()

    for l_idx in range(l):
        scale_a = to_block(sfa[:, :, l_idx]).cuda()
        scale_b = to_block(sfb[:, :, l_idx]).cuda()
        res = torch._scaled_mm(
            a_fp4[:, :, l_idx].view(torch.float4_e2m1fn_x2),
            b_fp4[:, :, l_idx].transpose(0, 1).view(torch.float4_e2m1fn_x2),
            scale_a, scale_b,
            bias=None,
            out_dtype=torch.float16,
        )
        c_out[:, :, l_idx].copy_(res)

def _is_case1(problem_sizes):
    # case1: g=8, all (n=4096,k=7168,l=1), m varies
    if len(problem_sizes) != 8:
        return False
    for (m, n, k, l) in problem_sizes:
        if not (n == 4096 and k == 7168 and l == 1):
            return False
    return True

def custom_kernel(data: input_t) -> output_t:
    """
    Optimized for case1 preparation cost:
    - reorder_sf_blocked_fast: view/permute instead of meshgrid/scatter
    - keep rest logic same
    """
    if len(data) == 4:
        abc_tensors, sfasfb_tensors, _maybe_reordered, problem_sizes = data
    else:
        abc_tensors, sfasfb_tensors, problem_sizes = data
        _maybe_reordered = None

    result_tensors = []

    # (Optional) if future you finds _maybe_reordered contains ready-to-use blocked SF,
    # you can plug it here. For now, we just ignore it safely.

    supported_k = {7168, 4096, 2048, 1536, 2304, 512, 256}

    for (a, b, c), (sfa, sfb), (m, n, k, l) in zip(abc_tensors, sfasfb_tensors, problem_sizes):
        assert l == 1, "This kernel assumes L==1 (matches provided benchmark shapes)."

        if sfa.device.type != "cuda":
            sfa = sfa.cuda(non_blocking=True)
        if sfb.device.type != "cuda":
            sfb = sfb.cuda(non_blocking=True)

        mn_pad = ceil_div(m, 128) * 128

        a_pad = pad_fp4_rows(a, m, mn_pad)
        b_use = b.contiguous()

        sfa_m = sfa[:, :, 0].contiguous()
        sfb_n = sfb[:, :, 0].contiguous()

        # >>> the key change:
        sfa_blk = reorder_sf_blocked_fast(sfa_m, m, k, mn_pad)
        sfb_blk = reorder_sf_blocked_fast(sfb_n, n, k, n)  # N already multiple of 128 for your target shapes

        if mn_pad == m:
            c_out = c
        else:
            c_out = torch.empty((mn_pad, n, l), device="cuda", dtype=torch.float16)

        if k in supported_k:
            gemm(a_pad, b_use, sfa_blk, sfb_blk, c_out)
            if mn_pad != m:
                c.copy_(c_out[:m, :, :])
        else:
            torch_scaled_mm_fallback(a, b, sfa, sfb, c, m, n, k, l)

        result_tensors.append(c)

    return result_tensors