submission 93595

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

# Kernel configuration parameters
   # Tile sizes for M, N, K dimensions
ab_dtype = cutlass.Float4E2M1FN  # FP4 data type for A and B
sf_dtype = cutlass.Float8E4M3FN  # FP8 data type for scale factors
c_dtype = cutlass.Float16  # FP16 output type
accum_dtype = cutlass.Float32
sf_vec_size = 16  # Scale factor block size (16 elements share one scale)
threads_per_m = 16  # Number of threads per CUDA thread block
threads_per_k = 16
mma_tiler_mnk = (threads_per_m, 1, 128)

"""
@cute.kernel
def kernel(
    mA_mkl: cute.Tensor,
    mB_nkl: cute.Tensor,
    mSFA_mkl: cute.Tensor,
    mSFB_nkl: cute.Tensor,
    mC_mnl: cute.Tensor,
):
    # 1. CONFIGURATION TENSOR CORE (MMA)
    # Utilisation de l'instruction FP16 -> FP32 (16x8x16)
    # Layout TN : A (Col-Major), B (Row-Major) est le standard pour GEMM
    mma_op = cute.MMA_Atom(cute.SM80_16x8x16_F32F16F16F32_TN)
    
    # On étend cet atome pour couvrir notre Block de threads (32 threads)
    # 1 warp suffit pour exécuter un atome 16x8x16
    tiled_mma = cute.make_tiled_mma(mma_op)

    # Get CUDA block and thread indices
    bidx, bidy, bidz = cute.arch.block_idx()
    tidx, _, _ = cute.arch.thread_idx()

    # --- Gestion des Tuiles Globales (Similaire à avant) ---
    # On garde mma_tiler_mnk pour le découpage global
    
    # A: [BlockM, BlockK]
    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: [BlockN, BlockK] - Note: BlockN doit être au moins 8 pour le MMA
    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: [BlockM, BlockN]
    gC_mnl = cute.local_tile(mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None))

    # --- Partitionnement pour le MMA (Le changement magique) ---
    
    # Accumulateur (C) : Partitionné selon le TiledMMA
    # Cela distribue la matrice C résultante sur les registres des threads
    tCgC = gC_mnl(None, None, bidx, bidy, bidz) # Vue globale de la tuile C
    tCrC = tiled_mma.get_slice(0).partition_C(tCgC) # Vue registre partitionnée
    tCrC.clear() # Initialiser à 0

    # --- Boucle K (Main Loop) ---
    k_tile_cnt = gA_mkl.layout[3].shape
    
    # Tensor Core K-block size = 16.
    # On itère par blocs de 16 (taille de l'atome K)
    for k_tile in range(k_tile_cnt):
        
        # 1. Charger les données depuis la mémoire globale
        # On sélectionne la k-ième tuile
        tAgA = gA_mkl(None, None, bidx, k_tile, bidz)
        tBgB = gB_nkl(None, None, bidy, k_tile, bidz)
        tAgSFA = gSFA_mkl(None, None, bidx, k_tile, bidz)
        tBgSFB = gSFB_nkl(None, None, bidy, k_tile, bidz)

        # 2. Partitionner les registres d'entrée (Fragments) pour le MMA
        # A (MxK) -> Fragment A
        tZrA = tiled_mma.get_slice(0).partition_fragment_A(tAgA)
        tZrSFA = tiled_mma.get_slice(0).partition_fragment_A(tAgSFA)
        
        # B (NxK) -> Fragment B
        tZrB = tiled_mma.get_slice(0).partition_fragment_B(tBgB)
        tZrSFB = tiled_mma.get_slice(0).partition_fragment_B(tBgSFB)

        # 3. Chargement (Load)
        # Note: Sur B200, on utiliserait cp.async ici
        rA_fp4 = tZrA.load()
        rB_fp4 = tZrB.load()
        rSFA_fp8 = tZrSFA.load()
        rSFB_fp8 = tZrSFB.load()

        # 4. Conversion et Application de l'Échelle (Dequantize)
        # On convertit tout en FP16 pour le Tensor Core
        rA_f16 = rA_fp4.to(cutlass.Float16)
        rB_f16 = rB_fp4.to(cutlass.Float16)
        rSFA_f16 = rSFA_fp8.to(cutlass.Float16)
        rSFB_f16 = rSFB_fp8.to(cutlass.Float16)

        # Calcul : A_scaled = A_fp16 * ScaleA_fp16
        # Cela se fait élément par élément dans les registres
        rA_scaled = rA_f16 * rSFA_f16
        rB_scaled = rB_f16 * rSFB_f16

        # 5. L'Instruction MMA (Tensor Core)
        # Équivalent à : tCrC += rA_scaled * rB_scaled
        # Mais fait en un cycle hardware pour un bloc 16x8x16
        cute.gemm(tiled_mma, tCrC, rA_scaled, rB_scaled, tCrC)

    # --- Epilogue ---
    
    # Le résultat est dans tCrC (fragmenté dans les registres)
    # On doit le re-mapper vers la mémoire globale
    # Pour GEMV N=1, on ne veut souvent que la colonne 0 si on a calculé N=8
    
    # Dans CuTe, on fait l'inverse du partitionnement C
    # On crée une "View" registre du tenseur global de sortie partitionné
    tCgC_part = tiled_mma.get_slice(0).partition_C(tCgC)
    
    # Stockage (avec conversion FP32 -> FP16)
    # Note: Ceci écrit 8 colonnes (si N>=8). Pour N=1, il faudrait masquer.
    # Pour simplifier ici, on suppose que tCgC a assez de place ou on utilise predication.
    
    # Copie simple (TMA store serait mieux sur B200)
    tCgC_part.store(tCrC.to(cutlass.Float16))

    return

"""
@cute.kernel
def kernel(
    mA_mkl: cute.Tensor,
    mB_nkl: cute.Tensor,
    mSFA_mkl: cute.Tensor,
    mSFB_nkl: cute.Tensor,
    mC_mnl: cute.Tensor,
):
    # Get CUDA block and thread indices
    bidx, bidy, bidz = cute.arch.block_idx()
    tidx, tidy, _ = cute.arch.thread_idx()

    # Extract the local tile for input matrix A (shape: [block_M, block_K, rest_M, rest_K, rest_L])
    gA_mkl = cute.local_tile(
        mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    # Extract the local tile for scale factor tensor for A (same shape as gA_mkl)
    # Here, block_M = (32, 4); block_K = (16, 4)
    gSFA_mkl = cute.local_tile(
        mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    # Extract the local tile for input matrix B (shape: [block_N, block_K, rest_N, rest_K, rest_L])
    gB_nkl = cute.local_tile(
        mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
    )
    # Extract the local tile for scale factor tensor for B (same shape as gB_nkl)
    gSFB_nkl = cute.local_tile(
        mSFB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
    )
    # Extract the local tile for output matrix C (shape: [block_M, block_N, rest_M, rest_N, rest_L])
    gC_mnl = cute.local_tile(
        mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
    )

    # Select output element corresponding to this thread and block indices
    tCgC = gC_mnl[tidx, None, bidx, bidy, bidz]
    tCgC = cute.make_tensor(tCgC.iterator, 1)
    res = cute.zeros_like(tCgC, accum_dtype)

    # Shared Memory
    allocator = cutlass.utils.SmemAllocator()
    smem_layout = cute.make_layout((threads_per_m, threads_per_k), stride = (threads_per_k, 1))
    shared_res = allocator.allocate_tensor(element_type=cutlass.Float32, layout=smem_layout)

    # Get the number of k tiles (depth dimension) for the reduction loop
    k_tile_cnt = gA_mkl.layout[3].shape
    for k_tile in range(tidy, k_tile_cnt, threads_per_k, unroll_full=True):
        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]

        tArA = cute.make_rmem_tensor_like(tAgA, c_dtype)
        tBrB = cute.make_rmem_tensor_like(tBgB, c_dtype)
        tABrAB = cute.make_rmem_tensor_like(tAgA, c_dtype)
        tArSFA = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)
        tBrSFB = cute.make_rmem_tensor_like(tBgSFB, accum_dtype)
        tSFrSF = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)

        # Load NVFP4 or FP8 values from global memory
        a_val_nvfp4 = tAgA.load()
        b_val_nvfp4 = tBgB.load()
        sfa_val_fp8 = tAgSFA.load()
        sfb_val_fp8 = tBgSFB.load()

        # Convert loaded values to float32 for computation (FFMA)
        a_val = a_val_nvfp4.to(c_dtype)
        b_val = b_val_nvfp4.to(c_dtype)
        sfa_val = sfa_val_fp8.to(sf_dtype)
        sfb_val = sfb_val_fp8.to(sf_dtype)

        # Store the converted values to RMEM CuTe tensors
        tArA.store(a_val)
        tBrB.store(b_val)
        tArSFA.store(sfa_val)
        tBrSFB.store(sfb_val)

        tABrAB.store(tArA.load() * tBrB.load())
        tSFrSF.store(tArSFA.load() * tBrSFB.load())

        # Iterate over SF vector tiles and compute the scale&matmul accumulation
        for i in cutlass.range_constexpr(mma_tiler_mnk[2]):
            res += tABrAB[i] * tSFrSF[i]
   
    shared_res[(tidx, tidy)] = res[0]
    cute.arch.sync_threads()
    
    if tidy == 0:
        out = cute.zeros_like(tCgC, accum_dtype)
        for i in cutlass.range_constexpr(threads_per_k):
            out += shared_res[(tidx, i)]

        # Store the final float16 result back to global memory
        tCgC.store(out.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,
):
    """
    Host-side JIT function to prepare tensors and launch GPU kernel.
    """
    m, _, k, l = problem_size
    # Create CuTe Tensor via pointer and 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)),
        ),
    )
    # We use n=128 to create the torch tensor to do fp4 computation via torch._scaled_mm
    # then copy torch tensor to cute tensor for cute customize kernel computation
    # therefore we need to ensure b_tensor has the right stride with this 128 padded size on n.
    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))
    )
    # Convert scale factor tensors to MMA layout
    # The layout matches Tensor Core requirements: (((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)


    cluster_dims = (2, 1, 1)
    grid_m = cute.ceil_div(c_tensor.shape[0], threads_per_m)
    if grid_m % cluster_dims[0] != 0:
        grid_m += (cluster_dims[0] - (grid_m % cluster_dims[0]))

    grid = (grid_m, 1, c_tensor.shape[2])

    kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
        grid=grid,
        block=[threads_per_m, threads_per_k, 1],
        cluster=cluster_dims,
    )
    return


_compiled_kernel_cache = None

def compile_kernel():
    """
    Compile the kernel once and cache it.
    This should be called before any timing measurements.

    Returns:
        The compiled kernel function
    """
    global _compiled_kernel_cache

    if _compiled_kernel_cache is not None:
        return _compiled_kernel_cache

    # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
    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)

    # Compile the kernel
    _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 custom_kernel(data: input_t) -> output_t:
    """
    Execute the block-scaled GEMV kernel.

    This is the main entry point called by the evaluation framework.
    It converts PyTorch tensors to CuTe tensors, launches the kernel,
    and returns the result.

    Args:
        data: Tuple of (a, b, sfa_cpu, sfb_cpu, c) PyTorch tensors
            a: [m, k, l] - Input matrix in float4e2m1fn
            b: [1, k, l] - Input vector in float4e2m1fn
            sfa_cpu: [m, k, l] - Scale factors in float8_e4m3fn
            sfb_cpu: [1, k, l] - Scale factors in float8_e4m3fn
            sfa_permuted: [32, 4, rest_m, 4, rest_k, l] - Scale factors in float8_e4m3fn
            sfb_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors in float8_e4m3fn
            c: [m, 1, l] - Output vector in float16

    Returns:
        Output tensor c with computed GEMV results
    """
    a, b, _, _, sfa_permuted, sfb_permuted, c = data

    # Ensure kernel is compiled (will use cached version if available)
    # To avoid the compilation overhead, we compile the kernel once and cache it.
    compiled_func = compile_kernel()

    # Get dimensions from MxKxL layout
    m, k, l = a.shape
    #if l==1:
    #    return custom_kernel_torch(data)
    # Torch use e2m1_x2 data type, thus k is halved
    k = k * 2
    # GEMV N dimension is always 1
    n = 1

    # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
    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
    )

    # Execute the compiled kernel
    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))

    return c