submission 69391

import torch
from task import input_t, output_t


# 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
@torch.compile()
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()


def _inner(a_ref, b_ref, sfa_ref, sfb_ref, l_idx):
    scale_a = to_blocked(sfa_ref[:, :, l_idx])
    scale_b = to_blocked(sfb_ref[:, :, l_idx])
    # (m, k) @ (n, k).T -> (m, n)
    # (m, k) @ (n, k).T -> (m, n)
    res = torch._scaled_mm(
        a_ref[:, :, l_idx],
        b_ref[:, :, l_idx].transpose(0, 1),
        scale_a,
        scale_b,
        bias=None,
        out_dtype=torch.float16,
    )
    return res[:, 0]
    
def custom_kernel(
    data: input_t,
) -> output_t:
    """
    PyTorch reference implementation of NVFP4 block-scaled GEMV.
    """
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data

    # Get dimensions from MxNxL layout
    _, _, l = c_ref.shape

    sfa_ref_gpu = sfa_ref_cpu.cuda()
    sfb_ref_gpu = sfb_ref_cpu.cuda()

    # Call torch._scaled_mm to compute the GEMV result
    for l_idx in range(l):
        c_ref[:, 0, l_idx] = _inner(a_ref, b_ref, sfa_ref_gpu, sfb_ref_gpu, l_idx)

    return c_ref