submission 108615

import torch
from task import input_t, output_t

sf_vec_size = 16


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


def to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
    """
    Convert scale factor tensor of shape (rows, cols) into the blocked layout
    expected by torch._scaled_mm for NVFP4 block-scaled matmul.

    Assumes:
      - rows is a multiple of 128
      - cols is a multiple of 4
    """
    rows, cols = input_matrix.shape

    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)

    # (rows, cols) -> (n_row_blocks, 128, n_col_blocks, 4)
    blocks = (
        input_matrix
        .view(n_row_blocks, 128, n_col_blocks, 4)
        .permute(0, 2, 1, 3)        # (n_row_blocks, n_col_blocks, 128, 4)
    )

    # -> (num_blocks, 32, 16)
    rearranged = (
        blocks
        .reshape(-1, 4, 32, 4)      # (num_blocks, 4, 32, 4)
        .transpose(1, 2)            # (num_blocks, 32, 4, 4)
        .reshape(-1, 32, 16)        # (num_blocks * 2, 32, 16)
    )

    # Flatten to 1D vector, which is what torch._scaled_mm expects
    return rearranged.flatten()


def custom_kernel(data: input_t) -> output_t:
    """
    Block-scale NVFP4 GEMV using torch._scaled_mm on NVIDIA B200.

    Args:
        data: Tuple that expands to:
            a:   torch.Tensor[float4e2m1fn]   of shape [m, k, l],
            b:   torch.Tensor[float4e2m1fn]   of shape [n_padded_128, k, l],
            sfa: torch.Tensor[float8_e4m3fnuz] of shape [m, k // 16, l],
            sfb: torch.Tensor[float8_e4m3fnuz] of shape [n_padded_128, k // 16, l],
            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:
        c: torch.Tensor[float16] of shape [m, 1, l]
    """
    # Unpack
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data

    # Dimensions from [m, k, l]
    m, k, l = a.shape

    # Loop over batch dimension L
    for l_idx in range(l):
        # Slices for this batch
        a_l = a[:, :, l_idx]          # [m, k]
        b_l = b[:, :, l_idx]          # [n_padded_128, k]
        sfa_l = sfa[:, :, l_idx]      # [m, k//16]
        sfb_l = sfb[:, :, l_idx]      # [n_padded_128, k//16]

        # Convert scale factors into blocked layout on GPU
        scale_a = to_blocked(sfa_l)
        scale_b = to_blocked(sfb_l)

        # NVFP4 block-scaled matmul via cuBLAS through torch._scaled_mm:
        #   (m, k) @ (n, k).T -> (m, n)
        res = torch._scaled_mm(
            a_l,                         # [m, k], float4_e2m1fn_x2
            b_l.transpose(0, 1),         # [k, n_padded_128], float4_e2m1fn_x2
            scale_a,                     # 1D blocked scales for A (float8_e4m3fnuz)
            scale_b,                     # 1D blocked scales for B (float8_e4m3fnuz)
            bias=None,
            out_dtype=torch.float16,     # accumulate and store as fp16
        )

        # Only the first column (logical N=1) is used as GEMV result
        c[:, 0, l_idx] = res[:, 0]

    return c