submission 101118

import torch
from task import input_t, output_t

import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_ptr
import cutlass.utils.blockscaled_layout as blockscaled_utils

from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass._mlir.dialects import nvvm

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

# Tile sizes for GEMV: M, N, K
#   - 128 rows of M per block
#   - 1 column of N (GEMV)
#   - 64 elements of K per tile
mma_tiler_mnk = (128, 1, 64)

ab_dtype = cutlass.Float4E2M1FN   # FP4 for A and B
sf_dtype = cutlass.Float8E4M3FN   # FP8 for scale factors
c_accum_dtype = cutlass.Float32   # FP32 accumulation buffer
c_out_dtype = cutlass.Float16     # Final FP16 output
sf_vec_size = 16                  # 16 FP4 elements share one FP8 scale
threads_per_cta = 128             # 128 threads per block (1 thread per M row)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

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


@dsl_user_op
def atomic_add_fp32(a: float | c_accum_dtype, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> None:
    """
    Atomic add on a float32 value in *global* memory.
    Wrapper around NVVM atomic rmw (FADD).
    """
    nvvm.atomicrmw(
        res=T.f32(),
        op=nvvm.AtomicOpKind.FADD,
        ptr=gmem_ptr.llvm_ptr,
        a=c_accum_dtype(a).ir_value(),
    )

# ---------------------------------------------------------------------------
# Kernel: extra blocks along K with atomic add into FP32 C buffer
# ---------------------------------------------------------------------------

@cute.kernel
def kernel(
    mA_mkl: cute.Tensor,   # [M, K, L] in FP4
    mB_nkl: cute.Tensor,   # [N_padded=128, K, L] in FP4
    mSFA_mkl: cute.Tensor, # scale A: special SF layout
    mSFB_nkl: cute.Tensor, # scale B: special SF layout
    mC_mnl: cute.Tensor,   # [M, N, L] in FP32 (accum buffer)
):
    # Block and thread indices
    bidx, bidy, bidz = cute.arch.block_idx()     # x: M tiles, y: K tiles, z: batch L
    tidx, _, _ = cute.arch.thread_idx()          # one thread per M row in the tile

    # Tiled views
    # A: uses M and K (no N)
    gA_mkl = cute.local_tile(
        mA_mkl,
        cute.slice_(mma_tiler_mnk, (None, 0, None)),
        (None, None, None),
    )
    gSFA_mkl = cute.local_tile(
        mSFA_mkl,
        cute.slice_(mma_tiler_mnk, (None, 0, None)),
        (None, None, None),
    )

    # B: uses N and K (no M)
    gB_nkl = cute.local_tile(
        mB_nkl,
        cute.slice_(mma_tiler_mnk, (0, None, None)),
        (None, None, None),
    )
    gSFB_nkl = cute.local_tile(
        mSFB_nkl,
        cute.slice_(mma_tiler_mnk, (0, None, None)),
        (None, None, None),
    )

    # C: uses M and N (no K)
    gC_mnl = cute.local_tile(
        mC_mnl,
        cute.slice_(mma_tiler_mnk, (None, None, 0)),
        (None, None, None),
    )

    # One thread = one output row within this M-tile (for N=1)
    # Note: we fix N-tile index to 0 here because GEMV has N=1.
    tCgC = gC_mnl[tidx, None, bidx, 0, bidz]
    tCgC = cute.make_tensor(tCgC.iterator, 1)  # scalar tensor view
    res = cute.zeros_like(tCgC, c_accum_dtype) # FP32 accumulator in registers

    # Each block now handles exactly ONE K-tile: bidy
    # So no outer loop over k_tile_cnt; bidy *is* the k_tile index.
    tAgA   = gA_mkl[tidx, None, bidx, bidy, bidz]
    tBgB   = gB_nkl[0,    None, 0,    bidy, bidz]
    tAgSFA = gSFA_mkl[tidx, None, bidx, bidy, bidz]
    tBgSFB = gSFB_nkl[0,    None, 0,    bidy, bidz]

    # Register-memory tensors (local fragments) for this tile
    tArA   = cute.make_rmem_tensor_like(tAgA,   c_out_dtype)    # A in FP16
    tBrB   = cute.make_rmem_tensor_like(tBgB,   c_out_dtype)    # B in FP16
    tArSFA = cute.make_rmem_tensor_like(tAgSFA, c_accum_dtype)  # SF A in FP32
    tBrSFB = cute.make_rmem_tensor_like(tBgSFB, c_accum_dtype)  # SF B in FP32

    tABrAB = cute.make_rmem_tensor_like(tAgA,   c_out_dtype)    # A*B (FP16)
    tSFrSF = cute.make_rmem_tensor_like(tAgSFA, c_accum_dtype)  # SF_A*SF_B (FP32)

    # Load NVFP4 / FP8 from GMEM
    a_val_nvfp4 = tAgA.load()
    b_val_nvfp4 = tBgB.load()
    sfa_val_fp8 = tAgSFA.load()
    sfb_val_fp8 = tBgSFB.load()

    # Convert to usable compute types
    a_val  = a_val_nvfp4.to(c_out_dtype)      # A → FP16
    b_val  = b_val_nvfp4.to(c_out_dtype)      # B → FP16
    sfa_val = sfa_val_fp8.to(c_accum_dtype)   # SF A → FP32
    sfb_val = sfb_val_fp8.to(c_accum_dtype)   # SF B → FP32

    # Store into local RMEM tensors
    tArA.store(a_val)
    tBrB.store(b_val)
    tArSFA.store(sfa_val)
    tBrSFB.store(sfb_val)

    # Pre-multiply A*B and SF_A*SF_B so inner loop is lighter
    tABrAB.store(tArA.load()   * tBrB.load())    # FP16
    tSFrSF.store(tArSFA.load() * tBrSFB.load())  # FP32

    # Inner loop over elements within this K-tile (64 elements)
    for i in cutlass.range_constexpr(mma_tiler_mnk[2]):  # 0..63
        res += tABrAB[i] * tSFrSF[i]   # FP16 * FP32 → FP32 accumulate

    # This block computed a partial sum over its K-chunk.
    # Accumulate it into the global FP32 buffer using atomic add.
    atomic_add_fp32(res[0], tCgC.iterator)
    return

# ---------------------------------------------------------------------------
# JIT wrapper: set up layouts, grid, and launch
# ---------------------------------------------------------------------------

@cute.jit
def my_kernel(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    sfa_ptr: cute.Pointer,
    sfb_ptr: cute.Pointer,
    c_ptr: cute.Pointer,
    problem_size: tuple,
):
    """
    Host-side JIT wrapper: builds CuTe tensors/layouts and launches the kernel.
    problem_size = (m, n, k, l)
    """
    m, n, k, l = problem_size

    # A: [M, K, L]
    a_tensor = cute.make_tensor(
        a_ptr,
        cute.make_layout(
            (m, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
        ),
    )

    # B: [N_padded=128, K, L]
    # N is padded to 128 in generate_input() for torch._scaled_mm compatibility.
    n_padded_128 = 128
    b_tensor = cute.make_tensor(
        b_ptr,
        cute.make_layout(
            (n_padded_128, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(n_padded_128 * k, 32)),
        ),
    )

    # C accumulation buffer: [M, N=1, L] in FP32
    c_tensor = cute.make_tensor(
        c_ptr,
        cute.make_layout((cute.assume(m, 32), n, l), stride=(1, 1, m)),
    )

    # Scale factor tensors in MMA layout:
    # (((32,4), REST_M), ((SF_K,4), REST_K), (1, REST_L))
    sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, sf_vec_size)
    sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)

    sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, sf_vec_size)
    sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)

    # ---- Grid: parallelize over M tiles *and* K tiles ----
    #
    #   grid.x = number of M-tiles  (ceil(M / 128))
    #   grid.y = number of K-tiles  (ceil(K / 64))
    #   grid.z = batch size L
    #
    m_blocks = cute.ceil_div(c_tensor.shape[0], mma_tiler_mnk[0])   # M_tiles
    k_blocks = cute.ceil_div(a_tensor.shape[1], mma_tiler_mnk[2])   # K_tiles
    grid = (m_blocks, k_blocks, c_tensor.shape[2])

    # Launch kernel
    kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
        grid=grid,
        block=[threads_per_cta, 1, 1],
        cluster=(1, 1, 1),
    )
    return

# ---------------------------------------------------------------------------
# Compile & cache
# ---------------------------------------------------------------------------

_compiled_kernel_cache = None

def compile_kernel():
    """
    Compile CuTe kernel once and cache it.
    """
    global _compiled_kernel_cache
    if _compiled_kernel_cache is not None:
        return _compiled_kernel_cache

    # Dummy pointers just to tell CuTe the types & address spaces
    a_ptr   = make_ptr(ab_dtype,        0, cute.AddressSpace.gmem, assumed_align=16)
    b_ptr   = make_ptr(ab_dtype,        0, cute.AddressSpace.gmem, assumed_align=16)
    c_ptr   = make_ptr(c_accum_dtype,   0, cute.AddressSpace.gmem, assumed_align=16)
    sfa_ptr = make_ptr(sf_dtype,        0, cute.AddressSpace.gmem, assumed_align=32)
    sfb_ptr = make_ptr(sf_dtype,        0, cute.AddressSpace.gmem, assumed_align=32)

    _compiled_kernel_cache = cute.compile(
        my_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0)
    )
    return _compiled_kernel_cache

# ---------------------------------------------------------------------------
# Entry point used by the competition framework
# ---------------------------------------------------------------------------

def custom_kernel(data: input_t) -> output_t:
    """
    Main entry point:
      - creates FP32 accumulation buffer,
      - launches CuTe kernel with extra K-blocks and atomic adds,
      - converts FP32 result to FP16 in-place into `c` and returns it.
    """
    a, b, _, _, sfa_permuted, sfb_permuted, c = data  # note: c is FP16

    compiled_func = compile_kernel()

    # Extract logical sizes from A: [M, K_packed, L]
    m, k_packed, l = a.shape
    # A/B are FP4 packed as e2m1_x2 → true K is 2 * k_packed
    k = k_packed * 2
    n = 1  # GEMV → N = 1

    # FP32 accumulation buffer with same [M,1,L] shape as c
    c_accum = torch.zeros_like(c, dtype=torch.float32)

    # Pointers for CuTe
    a_ptr   = make_ptr(ab_dtype,      a.data_ptr(),       cute.AddressSpace.gmem, assumed_align=16)
    b_ptr   = make_ptr(ab_dtype,      b.data_ptr(),       cute.AddressSpace.gmem, assumed_align=16)
    c_ptr   = make_ptr(c_accum_dtype, c_accum.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    sfa_ptr = make_ptr(sf_dtype,      sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
    sfb_ptr = make_ptr(sf_dtype,      sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)

    # Run kernel
    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))

    # Convert FP32 accum buffer to FP16 output in-place
    c.copy_(c_accum.to(torch.float16))

    return c