submission 75571

import torch
from task import input_t, output_t
from utils import make_match_reference


sf_vec_size = 16  # each scale covers 16 FP4 values along K

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


def to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
    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 custom_kernel(data: input_t) -> output_t:
    a, b, sfa_ref, sfb_ref, sfa_perm, sfb_perm, c = data

    _, _, L = c.shape

    for l_idx in range(L):
        sfa_slice = sfa_ref[:, :, l_idx]
        sfb_slice = sfb_ref[:, :, l_idx]

        scale_a = to_blocked(sfa_slice)  
        scale_b = to_blocked(sfb_slice)

        res = torch._scaled_mm(
            a[:, :, l_idx],
            b[:, :, l_idx].transpose(0, 1),
            scale_a,        
            scale_b,       
            bias=None,
            out_dtype=torch.float16,
        )

        c[:, 0, l_idx] = res[:, 0]

    return c