submission 106200

from typing import cast

import torch
from task import input_t, output_t

# Try to import CuTe / cutlass – this will fail on macOS CPU wheels
try:
    import cutlass
    import cutlass.cute as cute
    from cutlass.cute.runtime import make_ptr
    import cutlass.utils.blockscaled_layout as blockscaled_utils

    HAVE_CUTE = True
except Exception:
    HAVE_CUTE = False


# -----------------------------
# Shared configuration
# -----------------------------
mma_tiler_mnk = (128, 1, 64)  # M, N, K tile
ab_dtype = cutlass.Float4E2M1FN if HAVE_CUTE else None
sf_dtype = cutlass.Float8E4M3FN if HAVE_CUTE else None
c_dtype = cutlass.Float16 if HAVE_CUTE else None
sf_vec_size = 16
threads_per_cta = 128

_compiled_kernel_cache = None  # per-process cache
_compiled_kernel_cache_l1 = None  # L=1 specialized kernel cache


# -----------------------------
# CuTe kernel path (GPU-only)
# -----------------------------
if HAVE_CUTE:

    @cute.kernel
    def _kernel(
        mA_mkl: cute.Tensor,
        mB_nkl: cute.Tensor,
        mSFA_mkl: cute.Tensor,
        mSFB_nkl: cute.Tensor,
        mC_mnl: cute.Tensor,
    ):
        bidx, bidy, bidz = cute.arch.block_idx()
        tidx, _, _ = cute.arch.thread_idx()

        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)
        )
        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)
        )
        gC_mnl = cute.local_tile(
            mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
        )

        tCgC = gC_mnl[tidx, None, bidx, bidy, bidz]
        tCgC = cute.make_tensor(tCgC.iterator, 1)
        res = cute.zeros_like(tCgC, cutlass.Float32)

        k_tile_cnt = gA_mkl.layout[3].shape
        for k_tile in range(k_tile_cnt):
            tAgA = gA_mkl[tidx, None, bidx, k_tile, bidz]
            tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
            tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, bidz]
            tBgSFB = gSFB_nkl[0, None, bidy, k_tile, bidz]

            # Load fragments and convert directly to Float32
            a_f32 = tAgA.load().to(cutlass.Float32)
            b_f32 = tBgB.load().to(cutlass.Float32)
            sfa_f32 = tAgSFA.load().to(cutlass.Float32)
            sfb_f32 = tBgSFB.load().to(cutlass.Float32)

            # Fuse scale factors: SF = SFA * SFB
            sf_f32 = sfa_f32 * sfb_f32

            # Inner accumulation: res += a * b * sf
            for i in cutlass.range_constexpr(mma_tiler_mnk[2]):
                res += a_f32[i] * b_f32[i] * sf_f32[i]

        tCgC.store(res.to(cutlass.Float16))
        return


    @cute.kernel
    def _kernel_l1(
        mA_mkl: cute.Tensor,
        mB_nkl: cute.Tensor,
        mSFA_mkl: cute.Tensor,
        mSFB_nkl: cute.Tensor,
        mC_mnl: cute.Tensor,
    ):
        # Same tiling as _kernel, but no L dimension in the grid
        bidx, bidy, bidz = cute.arch.block_idx()
        tidx, _, _ = cute.arch.thread_idx()

        # Local tiles exactly as before
        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)
        )
        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)
        )
        gC_mnl = cute.local_tile(
            mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
        )

        # For L=1, the only valid "rest_L" index is 0.
        # (bidz from arch.block_idx().z will always be 0 because we launch grid.z = 1.)
        l_idx = 0

        tCgC = gC_mnl[tidx, None, bidx, bidy, l_idx]
        tCgC = cute.make_tensor(tCgC.iterator, 1)
        res = cute.zeros_like(tCgC, cutlass.Float32)

        k_tile_cnt = gA_mkl.layout[3].shape
        for k_tile in range(k_tile_cnt):
            tAgA = gA_mkl[tidx, None, bidx, k_tile, l_idx]
            tBgB = gB_nkl[0, None, bidy, k_tile, l_idx]
            tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, l_idx]
            tBgSFB = gSFB_nkl[0, None, bidy, k_tile, l_idx]

            # Load fragments and convert directly to Float32
            a_f32 = tAgA.load().to(cutlass.Float32)
            b_f32 = tBgB.load().to(cutlass.Float32)
            sfa_f32 = tAgSFA.load().to(cutlass.Float32)
            sfb_f32 = tBgSFB.load().to(cutlass.Float32)

            # Fuse scale factors for this fragment
            sf_f32 = sfa_f32 * sfb_f32

            # Inner accumulation: res += a * b * sf
            for i in cutlass.range_constexpr(mma_tiler_mnk[2]):
                res += a_f32[i] * b_f32[i] * sf_f32[i]

        tCgC.store(res.to(cutlass.Float16))
        return


    @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,
    ):
        m, _, k, l = problem_size

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

        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_tensor = cute.make_tensor(
            c_ptr,
            cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m)),
        )

        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 = (
            cute.ceil_div(c_tensor.shape[0], 128),
            1,
            c_tensor.shape[2],
        )

        _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


    @cute.jit
    def _my_kernel_l1(
        a_ptr: cute.Pointer,
        b_ptr: cute.Pointer,
        sfa_ptr: cute.Pointer,
        sfb_ptr: cute.Pointer,
        c_ptr: cute.Pointer,
        problem_size: tuple,
    ):
        # For L=1 we still pass (m, n, k, l) but know l == 1
        m, _, k, l = problem_size

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

        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_tensor = cute.make_tensor(
            c_ptr,
            cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m)),
        )

        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)

        # L=1 -> grid.z = 1
        grid = (
            cute.ceil_div(c_tensor.shape[0], 128),
            1,
            1,
        )

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


    def _compile_kernel():
        global _compiled_kernel_cache
        if _compiled_kernel_cache is not None:
            return _compiled_kernel_cache

        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_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


    def _compile_kernel_l1():
        global _compiled_kernel_cache_l1
        if _compiled_kernel_cache_l1 is not None:
            return _compiled_kernel_cache_l1

        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_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_l1 = cute.compile(
            _my_kernel_l1, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 1)
        )
        return _compiled_kernel_cache_l1


def _nvfp4_blockscaled_gemv_gpu(data: input_t) -> output_t:
    """
    Real CuTe path – only called when cutlass + CUDA are available.
    """
    assert HAVE_CUTE, "CuTe path requested on a system without cutlass/CUDA."

    a, b, _sfa, _sfb, sfa_perm, sfb_perm, c = data
    m, k_half, l = a.shape  # a: [m, k/2, l] in float4_e2m1fn_x2
    k = k_half * 2
    n = 1

    if l == 1:
        compiled = _compile_kernel_l1()
    else:
        compiled = _compile_kernel()

    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_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    sfa_ptr = make_ptr(
        sf_dtype, sfa_perm.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
    )
    sfb_ptr = make_ptr(
        sf_dtype, sfb_perm.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
    )

    compiled(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))
    return c


# -----------------------------
# CPU-only fallback for macOS dev
# -----------------------------
def _fallback_gemv_cpu(data: input_t) -> output_t:
    """
    Extremely simplified fallback so you can at least run shape/syntax checks on CPU.

    IMPORTANT:
    - Only meant for local dev.
    - Does NOT try to be nvfp4-correct; it just does a float32 GEMV
      with the same outer shapes.
    """
    a, b, _sfa, _sfb, _sfa_perm, _sfb_perm, c = data

    # Interpret a,b as regular float tensors on CPU
    # For local fake tests you'll construct them as float32/float16 anyway.
    # Shapes (for fake data): a: [m, k, l], b: [1, k, l], c: [m, 1, l]
    m, k_like, l = a.shape

    out = c
    out.zero_()

    for li in range(l):
        # [m, k] @ [k] -> [m]
        a_slice = a[:, :, li].to(torch.float32)
        b_slice = b[0, :, li].to(torch.float32)
        y = a_slice @ b_slice
        out[:, 0, li] = y.to(out.dtype)

    return cast(output_t, out)


# -----------------------------
# Public entrypoint
# -----------------------------
def custom_kernel(data: input_t) -> output_t:
    """
    Entry point used by the evaluation framework.

    On GPU box with CuTe: runs the high-performance kernel.
    On macOS/CPU: runs a dumb fallback so you can at least import and smoke-test.
    """
    if HAVE_CUTE and torch.cuda.is_available():
        return _nvfp4_blockscaled_gemv_gpu(data)
    else:
        # local dev / CPU-only environment
        return _fallback_gemv_cpu(data)