submission 75219

from dataclasses import dataclass
from typing import Callable, Dict

import cutlass
import cutlass.cute as cute
import cutlass.utils.blockscaled_layout as blockscaled_utils
import torch
from cutlass.cute.runtime import make_ptr
from task import input_t, output_t


@dataclass(frozen=True)
class KernelConfig:
    """
    Describes one GEMV kernel variant tuned for a specific workload shape.
    """

    tile_m: int
    tile_n: int = 1
    tile_k: int = 64
    threads_per_cta: int = 128
    pipeline_stages: int = 1
    name: str = "default"

    @property
    def mma_tiler_mnk(self):
        return (self.tile_m, self.tile_n, self.tile_k)


# FP formats
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
c_dtype = cutlass.Float16
sf_vec_size = 16

# GB200-oriented presets
STREAMING_CONFIG = KernelConfig(
    tile_m=64,
    tile_k=128,
    threads_per_cta=64,
    pipeline_stages=2,
    name="streaming_k128",
)
BATCH_CONFIG = KernelConfig(
    tile_m=128,
    tile_k=64,
    threads_per_cta=128,
    pipeline_stages=1,
    name="batch_m128",
)
LATENCY_CONFIG = KernelConfig(
    tile_m=64,
    tile_k=64,
    threads_per_cta=64,
    pipeline_stages=1,
    name="latency_m64",
)
FALLBACK_CONFIG = KernelConfig(
    tile_m=128,
    tile_k=64,
    threads_per_cta=128,
    pipeline_stages=2,
    name="fallback_double_buffer",
)


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


def _build_kernel_for_config(config: KernelConfig):
    """
    Build and JIT-compile a CuTe kernel for the given configuration.
    """

    mma_tiler_mnk = config.mma_tiler_mnk
    threads_per_cta = config.threads_per_cta
    pipeline_stages = config.pipeline_stages

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

        stage_a = [
            cute.make_rmem_tensor_like(
                gA_mkl[tidx, None, bidx, 0, bidz], cutlass.Float32
            )
            for _ in range(pipeline_stages)
        ]
        stage_b = [
            cute.make_rmem_tensor_like(gB_nkl[0, None, bidy, 0, bidz], cutlass.Float32)
            for _ in range(pipeline_stages)
        ]
        stage_sfa = [
            cute.make_rmem_tensor_like(
                gSFA_mkl[tidx, None, bidx, 0, bidz], cutlass.Float32
            )
            for _ in range(pipeline_stages)
        ]
        stage_sfb = [
            cute.make_rmem_tensor_like(
                gSFB_nkl[0, None, bidy, 0, bidz], cutlass.Float32
            )
            for _ in range(pipeline_stages)
        ]

        k_tile_cnt = gA_mkl.layout[3].shape
        k_tile = 0
        while k_tile < k_tile_cnt:
            stages = pipeline_stages
            remaining = k_tile_cnt - k_tile
            if remaining < stages:
                stages = remaining

            for stage in cutlass.range_constexpr(pipeline_stages):
                if stage < stages:
                    kt = k_tile + stage
                    tAgA = gA_mkl[tidx, None, bidx, kt, bidz]
                    tBgB = gB_nkl[0, None, bidy, kt, bidz]
                    tAgSFA = gSFA_mkl[tidx, None, bidx, kt, bidz]
                    tBgSFB = gSFB_nkl[0, None, bidy, kt, bidz]

                    stage_a[stage].store(tAgA.load().to(cutlass.Float32))
                    stage_b[stage].store(tBgB.load().to(cutlass.Float32))
                    stage_sfa[stage].store(tAgSFA.load().to(cutlass.Float32))
                    stage_sfb[stage].store(tBgSFB.load().to(cutlass.Float32))

            for stage in cutlass.range_constexpr(pipeline_stages):
                if stage < stages:
                    tArA = stage_a[stage]
                    tBrB = stage_b[stage]
                    tArSFA = stage_sfa[stage]
                    tBrSFB = stage_sfb[stage]

                    for i in cutlass.range_constexpr(mma_tiler_mnk[2]):
                        res += tArA[i] * tArSFA[i] * tBrB[i] * tBrSFB[i]

            k_tile += stages

        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], mma_tiler_mnk[0]),
            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

    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)

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


_compiled_kernel_cache: Dict[KernelConfig, Callable] = {}


def compile_kernel(config: KernelConfig) -> Callable:
    """
    Compile (or fetch) the CUDA kernel for the provided configuration.
    """

    cached = _compiled_kernel_cache.get(config)
    if cached is not None:
        return cached
    compiled = _build_kernel_for_config(config)
    _compiled_kernel_cache[config] = compiled
    return compiled


def _select_kernel_config(m: int, k: int, l: int) -> KernelConfig:
    """
    Heuristic tuned for the three benchmark points plus nearby workloads on GB200.
    """

    # Mapping:
    # (7168, 16384, 1)  -> STREAMING_CONFIG (tile_k=128, 2-stage pipeline)
    # (4096, 7168, 8)   -> BATCH_CONFIG (larger CTA tile on M for reuse across L)
    # (7168, 2048, 4)   -> LATENCY_CONFIG (more CTAs, shorter K loop)
    if l == 1 or (l <= 2 and k >= 8192):
        return STREAMING_CONFIG
    if l >= 6:
        return BATCH_CONFIG
    if k <= 3072:
        return LATENCY_CONFIG
    return FALLBACK_CONFIG


def custom_kernel(data: input_t) -> output_t:
    """
    Execute the block-scaled GEMV kernel optimized for GB200 workloads.
    """

    a, b, _, _, sfa_permuted, sfb_permuted, c = data

    m, k, l = a.shape
    k = k * 2  # torch e2m1_x2 packing halves the logical K.
    n = 1

    config = _select_kernel_config(m, k, l)
    compiled_func = compile_kernel(config)

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

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

    return c