submission 69271

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 custom_kernel(data: input_t) -> output_t:
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
    
    m, _, l = c_ref.shape
    
    # Batch convert all scale factors upfront
    sfa_list = [to_blocked(sfa_ref_cpu[:, :, i]).cuda() for i in range(l)]
    sfb_list = [to_blocked(sfb_ref_cpu[:, :, i]).cuda() for i in range(l)]
    
    # Process all batches with pre-converted scales
    for l_idx in range(l):
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[:, :, l_idx].transpose(0, 1),
            sfa_list[l_idx],
            sfb_list[l_idx],
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, 0, l_idx] = res[:, 0]
    
    return c_ref