submission 95843

import torch
from task import input_t, output_t


# Helper: ceil division
def ceil_div(a: int, b: int) -> int:
    return (a + b - 1) // b


# Helper: convert 2D scale-factor tensor to blocked layout expected by torch._scaled_mm
def to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
    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  # in this problem it's already padded correctly

    # (n_row_blocks, 128, n_col_blocks, 4)
    blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    # -> reshape to (-1, 4, 32, 4), then reorder
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)

    # Flatten to 1D as expected by torch._scaled_mm
    return rearranged.flatten()


def custom_kernel(data: input_t) -> output_t:
    """
    Baseline NVFP4 GEMV kernel using torch._scaled_mm,
    basically mirroring the reference implementation so we pass correctness tests.
    """

    # Unpack the 7-tuple
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data

    # c_ref has shape [m, 1, l] (after permute), so get batch size l
    _, _, l = c_ref.shape

    # For each batch index
    for l_idx in range(l):
        # Convert scale factor tensors to blocked layout (still on CPU)
        scale_a = to_blocked(sfa_ref_cpu[:, :, l_idx])
        scale_b = to_blocked(sfb_ref_cpu[:, :, l_idx])

        # (m, k) @ (n, k).T -> (m, n); here n == 1 but b is padded to 128
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],                     # [m, k], nvfp4
            b_ref[:, :, l_idx].transpose(0, 1),     # [k, n_padded_128]
            scale_a.cuda(),                         # nvfp4 scales for A
            scale_b.cuda(),                         # nvfp4 scales for B
            bias=None,
            out_dtype=torch.float16,
        )

        # Store GEMV result into c_ref[:, 0, l_idx]
        c_ref[:, 0, l_idx] = res[:, 0]

    return c_ref