submission 69828

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):
    rows, cols = input_matrix.shape
    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 _block_scales_across_L(scale_2d_per_L: torch.Tensor) -> torch.Tensor:
    """
    Vectorizes your to_blocked across L.
    Expects input of shape (R, C, L), returns (L, flat), where
    each row l is to_blocked( scale_2d_per_L[:,:,l] ).
    """
    R, C, L = scale_2d_per_L.shape
    n_row_blocks = ceil_div(R, 128)
    n_col_blocks = ceil_div(C, 4)
    # (R, C, L) -> (n_row_blocks,128, n_col_blocks,4, L)
    t = scale_2d_per_L.view(n_row_blocks, 128, n_col_blocks, 4, L)
    # -> (L, n_row_blocks, n_col_blocks, 128, 4)
    t = t.permute(4, 0, 2, 1, 3).contiguous()
    # -> (L, n_row_blocks*n_col_blocks, 4, 32, 4)
    t = t.view(L, n_row_blocks * n_col_blocks, 4, 32, 4)
    # -> (L, n_row_blocks*n_col_blocks, 32, 4, 4)
    t = t.transpose(2, 3).contiguous()
    # -> (L, n_row_blocks*n_col_blocks, 32, 16)
    t = t.view(L, n_row_blocks * n_col_blocks, 32, 16)
    # -> (L, flat)
    return t.view(L, -1)

@torch.inference_mode()
def custom_kernel(data: input_t) -> output_t:
    # Matches the evaluator’s tuple: (a, b, sfa, sfb, _, _, c)
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, *_ , c_ref = data

    M, K, L = a_ref.shape
    dev = a_ref.device

    # 1) Pre-block scales across L and move to device ONCE
    # Assumptions (consistent with your working code):
    #   sfa_ref_cpu: (M, K//16, L)
    #   sfb_ref_cpu: (K//16, 1, L)   <-- note: many harnesses use this order for “K-major”
    # If your sfb is actually (1, K//16, L), transpose it below.
    sfa_blocked = _block_scales_across_L(sfa_ref_cpu)            # (L, flat) on CPU
    # Ensure sfb has rows = K//16, cols = 1 for the blocking math
    if sfb_ref_cpu.shape[0] == 1 and sfb_ref_cpu.shape[1] == (K // 16):
        sfb_2dL = sfb_ref_cpu.permute(1, 0, 2)                   # (K//16, 1, L)
    else:
        sfb_2dL = sfb_ref_cpu                                    # assume already (K//16, 1, L)
    sfb_blocked = _block_scales_across_L(sfb_2dL)                 # (L, flat)

    sfa_blocked = sfa_blocked.to(dev, non_blocking=True)
    sfb_blocked = sfb_blocked.to(dev, non_blocking=True)

    # 2) Per-L call into scaled_mm (no CPU<->GPU copies inside the loop)
    for l_idx in range(L):
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],                     # (M, K) nvfp4
            b_ref[:, :, l_idx].transpose(0, 1),     # (K, 1) nvfp4
            sfa_blocked[l_idx],                     # 1D blocked scales
            sfb_blocked[l_idx],
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, 0, l_idx].copy_(res[:, 0])

    return c_ref