submission 103683

import torch
from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    """
    NVFP4 block-scaled GEMV optimized with inlined operations and unrolled small-L paths.
    """
    a, b, sfa_host, sfb_host, sfa_permuted, sfb_permuted, c_out = data
    del sfa_host, sfb_host

    m, _, l = c_out.shape

    if l == 1:
        scale_a = sfa_permuted.select(5, 0).permute(2, 4, 0, 1, 3).reshape(-1)
        scale_b = sfb_permuted.select(5, 0).permute(2, 4, 0, 1, 3).reshape(-1)
        gemv = torch._scaled_mm(
            a.select(2, 0),
            b.select(2, 0).t(),
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16,
        )
        return gemv[:, :1].view(m, 1, 1)

    sfa_blocked = sfa_permuted.permute(5, 2, 4, 0, 1, 3).reshape(l, -1)
    sfb_blocked = sfb_permuted.permute(5, 2, 4, 0, 1, 3).reshape(l, -1)
    b_t = b.permute(2, 1, 0)

    if l == 2:
        g0 = torch._scaled_mm(a.select(2, 0), b_t[0], sfa_blocked[0], sfb_blocked[0], bias=None, out_dtype=torch.float16)[:, 0]
        g1 = torch._scaled_mm(a.select(2, 1), b_t[1], sfa_blocked[1], sfb_blocked[1], bias=None, out_dtype=torch.float16)[:, 0]
        return torch.stack([g0, g1], dim=1).unsqueeze(1)

    if l == 4:
        g0 = torch._scaled_mm(a.select(2, 0), b_t[0], sfa_blocked[0], sfb_blocked[0], bias=None, out_dtype=torch.float16)[:, 0]
        g1 = torch._scaled_mm(a.select(2, 1), b_t[1], sfa_blocked[1], sfb_blocked[1], bias=None, out_dtype=torch.float16)[:, 0]
        g2 = torch._scaled_mm(a.select(2, 2), b_t[2], sfa_blocked[2], sfb_blocked[2], bias=None, out_dtype=torch.float16)[:, 0]
        g3 = torch._scaled_mm(a.select(2, 3), b_t[3], sfa_blocked[3], sfb_blocked[3], bias=None, out_dtype=torch.float16)[:, 0]
        return torch.stack([g0, g1, g2, g3], dim=1).unsqueeze(1)

    if l == 8:
        a0, a1, a2, a3 = a.select(2, 0), a.select(2, 1), a.select(2, 2), a.select(2, 3)
        a4, a5, a6, a7 = a.select(2, 4), a.select(2, 5), a.select(2, 6), a.select(2, 7)
        g0 = torch._scaled_mm(a0, b_t[0], sfa_blocked[0], sfb_blocked[0], bias=None, out_dtype=torch.float16)[:, 0]
        g1 = torch._scaled_mm(a1, b_t[1], sfa_blocked[1], sfb_blocked[1], bias=None, out_dtype=torch.float16)[:, 0]
        g2 = torch._scaled_mm(a2, b_t[2], sfa_blocked[2], sfb_blocked[2], bias=None, out_dtype=torch.float16)[:, 0]
        g3 = torch._scaled_mm(a3, b_t[3], sfa_blocked[3], sfb_blocked[3], bias=None, out_dtype=torch.float16)[:, 0]
        g4 = torch._scaled_mm(a4, b_t[4], sfa_blocked[4], sfb_blocked[4], bias=None, out_dtype=torch.float16)[:, 0]
        g5 = torch._scaled_mm(a5, b_t[5], sfa_blocked[5], sfb_blocked[5], bias=None, out_dtype=torch.float16)[:, 0]
        g6 = torch._scaled_mm(a6, b_t[6], sfa_blocked[6], sfb_blocked[6], bias=None, out_dtype=torch.float16)[:, 0]
        g7 = torch._scaled_mm(a7, b_t[7], sfa_blocked[7], sfb_blocked[7], bias=None, out_dtype=torch.float16)[:, 0]
        return torch.stack([g0, g1, g2, g3, g4, g5, g6, g7], dim=1).unsqueeze(1)

    result = torch.empty_like(c_out)

    i = 0
    while i + 3 < l:
        gemv0 = torch._scaled_mm(
            a.select(2, i),
            b_t[i],
            sfa_blocked[i],
            sfb_blocked[i],
            bias=None,
            out_dtype=torch.float16,
        )
        gemv1 = torch._scaled_mm(
            a.select(2, i + 1),
            b_t[i + 1],
            sfa_blocked[i + 1],
            sfb_blocked[i + 1],
            bias=None,
            out_dtype=torch.float16,
        )
        gemv2 = torch._scaled_mm(
            a.select(2, i + 2),
            b_t[i + 2],
            sfa_blocked[i + 2],
            sfb_blocked[i + 2],
            bias=None,
            out_dtype=torch.float16,
        )
        gemv3 = torch._scaled_mm(
            a.select(2, i + 3),
            b_t[i + 3],
            sfa_blocked[i + 3],
            sfb_blocked[i + 3],
            bias=None,
            out_dtype=torch.float16,
        )
        result[:, 0, i] = gemv0[:, 0]
        result[:, 0, i + 1] = gemv1[:, 0]
        result[:, 0, i + 2] = gemv2[:, 0]
        result[:, 0, i + 3] = gemv3[:, 0]
        i += 4

    while i < l:
        gemv = torch._scaled_mm(
            a.select(2, i),
            b_t[i],
            sfa_blocked[i],
            sfb_blocked[i],
            bias=None,
            out_dtype=torch.float16,
        )
        result[:, 0, i] = gemv[:, 0]
        i += 1

    return result