submission 71538

import torch

# Scaling factor vector size
sf_vec_size = 16

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

def to_blocked(input_matrix):
    """Convert scale factors to blocked format"""
    rows, cols = input_matrix.shape
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)
    blocks = input_matrix.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(mode="max-autotune", fullgraph=True)
def fused_gemv_unrolled_l4(a0, a1, a2, a3, b0, b1, b2, b3,
                           sa0, sa1, sa2, sa3, sb0, sb1, sb2, sb3, c):
    """Fully unrolled and fused GEMV for L=4"""
    # All operations in single compiled kernel with proper keyword arguments
    r0 = torch._scaled_mm(a0, b0, sa0, sb0, bias=None, out_dtype=torch.float16)[:, 0]
    r1 = torch._scaled_mm(a1, b1, sa1, sb1, bias=None, out_dtype=torch.float16)[:, 0]
    r2 = torch._scaled_mm(a2, b2, sa2, sb2, bias=None, out_dtype=torch.float16)[:, 0]
    r3 = torch._scaled_mm(a3, b3, sa3, sb3, bias=None, out_dtype=torch.float16)[:, 0]

    c[:, 0, 0] = r0
    c[:, 0, 1] = r1
    c[:, 0, 2] = r2
    c[:, 0, 3] = r3
    return c

@torch.compile(mode="max-autotune", fullgraph=True)
def fused_gemv_unrolled_l8(a_list, b_list, sa_list, sb_list, c):
    """Fully unrolled and fused GEMV for L=8"""
    # Unroll for better optimization with proper keyword arguments
    r0 = torch._scaled_mm(a_list[0], b_list[0], sa_list[0], sb_list[0], bias=None, out_dtype=torch.float16)[:, 0]
    r1 = torch._scaled_mm(a_list[1], b_list[1], sa_list[1], sb_list[1], bias=None, out_dtype=torch.float16)[:, 0]
    r2 = torch._scaled_mm(a_list[2], b_list[2], sa_list[2], sb_list[2], bias=None, out_dtype=torch.float16)[:, 0]
    r3 = torch._scaled_mm(a_list[3], b_list[3], sa_list[3], sb_list[3], bias=None, out_dtype=torch.float16)[:, 0]
    r4 = torch._scaled_mm(a_list[4], b_list[4], sa_list[4], sb_list[4], bias=None, out_dtype=torch.float16)[:, 0]
    r5 = torch._scaled_mm(a_list[5], b_list[5], sa_list[5], sb_list[5], bias=None, out_dtype=torch.float16)[:, 0]
    r6 = torch._scaled_mm(a_list[6], b_list[6], sa_list[6], sb_list[6], bias=None, out_dtype=torch.float16)[:, 0]
    r7 = torch._scaled_mm(a_list[7], b_list[7], sa_list[7], sb_list[7], bias=None, out_dtype=torch.float16)[:, 0]

    c[:, 0, 0] = r0
    c[:, 0, 1] = r1
    c[:, 0, 2] = r2
    c[:, 0, 3] = r3
    c[:, 0, 4] = r4
    c[:, 0, 5] = r5
    c[:, 0, 6] = r6
    c[:, 0, 7] = r7
    return c

def custom_kernel(data):
    """
    Ultra-optimized NVFP4 GEMV with fully unrolled loops.
    """
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
    _, _, l = c_ref.shape

    if l == 1:
        # Single GEMV - most critical path
        scale_a = to_blocked(sfa_ref_cpu[:, :, 0])
        scale_b = to_blocked(sfb_ref_cpu[:, :, 0])

        res = torch._scaled_mm(
            a_ref[:, :, 0],
            b_ref[:, :, 0].transpose(0, 1),
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16
        )
        c_ref[:, 0, 0] = res[:, 0]

    elif l == 4:
        # Fully unrolled for L=4
        # Pre-process all data
        a0, a1, a2, a3 = (a_ref[:, :, i] for i in range(4))
        b0, b1, b2, b3 = (b_ref[:, :, i].transpose(0, 1) for i in range(4))
        sa0 = to_blocked(sfa_ref_cpu[:, :, 0])
        sa1 = to_blocked(sfa_ref_cpu[:, :, 1])
        sa2 = to_blocked(sfa_ref_cpu[:, :, 2])
        sa3 = to_blocked(sfa_ref_cpu[:, :, 3])
        sb0 = to_blocked(sfb_ref_cpu[:, :, 0])
        sb1 = to_blocked(sfb_ref_cpu[:, :, 1])
        sb2 = to_blocked(sfb_ref_cpu[:, :, 2])
        sb3 = to_blocked(sfb_ref_cpu[:, :, 3])

        c_ref = fused_gemv_unrolled_l4(
            a0, a1, a2, a3, b0, b1, b2, b3,
            sa0, sa1, sa2, sa3, sb0, sb1, sb2, sb3, c_ref
        )

    elif l == 8:
        # Fully unrolled for L=8
        a_list = [a_ref[:, :, i] for i in range(8)]
        b_list = [b_ref[:, :, i].transpose(0, 1) for i in range(8)]
        sa_list = [to_blocked(sfa_ref_cpu[:, :, i]) for i in range(8)]
        sb_list = [to_blocked(sfb_ref_cpu[:, :, i]) for i in range(8)]

        c_ref = fused_gemv_unrolled_l8(a_list, b_list, sa_list, sb_list, c_ref)

    else:
        # General case
        for l_idx in range(l):
            scale_a = to_blocked(sfa_ref_cpu[:, :, l_idx])
            scale_b = to_blocked(sfb_ref_cpu[:, :, l_idx])

            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
            )
            c_ref[:, 0, l_idx] = res[:, 0]

    return c_ref