submission 72157

import torch
from task import input_t, output_t

# Kernel configuration parameters
sf_vec_size = 16


# Helper function for ceiling division
def ceil_div(a, b):
    return (a + b - 1) // b


# Helper function to convert scale factor tensor to blocked format
def to_blocked(input_matrix):
    rows, cols = input_matrix.shape

    # Please ensure rows and cols are multiples of 128 and 4 respectively
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)

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

@torch.compile
def to_blocked_2(input_matrix):
    rows, cols, l = input_matrix.shape

    # Please ensure rows and cols are multiples of 128 and 4 respectively
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)

    padded = input_matrix
    blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4, l).permute(4, 0, 2, 1, 3)
    rearranged = blocks.reshape(l, -1, 4, 32, 4).transpose(2, 3).reshape(l, -1, 32, 16)

    return rearranged.reshape(l, -1)


def custom_kernel(
    data: input_t,
) -> output_t:
    """
    Reference implementation of block-scale fp8 gemv
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float4e2m1fn] of shape [m, k, l],
            b: torch.Tensor[float4e2m1fn] of shape [1, k, l],
            sfa: torch.Tensor[float8_e4m3fnuz] of shape [m, k // 16, l], used by reference implementation
            sfb: torch.Tensor[float8_e4m3fnuz] of shape [1, k // 16, l], used by reference implementation
            sfa_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_m, 4, rest_k, l],
            sfb_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_n, 4, rest_k, l],
            c: torch.Tensor[float16] of shape [m, 1, l]
    Returns:
        Tensor containing output in float16
        c: torch.Tensor[float16] of shape [m, 1, l]
    """
    """
    PyTorch reference implementation of NVFP4 block-scaled GEMV.
    """
    a_ref, b_ref, _, _, sfa_permuted, sfb_permuted, c_ref = data

    # Get dimensions from MxNxL layout
    m, k, l = a_ref.shape
    n_padded_128 = 128

    # a_ref = a_ref.view(torch.uint8).permute(2, 0, 1).contiguous().view(torch.float4_e2m1fn_x2) # [l, m, k]
    # b_ref = b_ref.view(torch.uint8).permute(2, 1, 0).contiguous().view(torch.float4_e2m1fn_x2) # [l, k, 1]
    # Convert the scale factor tensor to blocked format
    # scale_a = to_blocked_2(sfa_ref_cpu.cuda())
    # scale_b = to_blocked_2(sfb_ref_cpu.cuda())

    # Call torch._scaled_mm to compute the GEMV result
    tmp_c = torch.empty((l, 64, m), dtype=torch.float16, device="cuda")
    for l_idx in range(l):

        # (m, k) @ (n, k).T -> (m, n)
        torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[0:64, :, l_idx].transpose(0, 1),
            sfa_permuted[:,:,:,:,:,l_idx].permute(2, 4, 0, 1, 3).flatten(),
            sfb_permuted[:,:,:,:,:,l_idx].permute(2, 4, 0, 1, 3).flatten(),
            bias=None,
            out_dtype=torch.float16,
            out=tmp_c[l_idx,:,:].transpose(0,1),
        )
    return tmp_c[:, 0:1, :].permute(2, 1, 0)