submission 487985

import torch
from task import input_t, output_t


@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data
    num_groups = len(abc_tensors)
    result_tensors = []

    for i in range(num_groups):
        a_ref, b_ref, c_ref = abc_tensors[i]
        sfa_reordered, sfb_reordered = sfasfb_reordered_tensors[i]
        m, n, k, l = problem_sizes[i]

        # Reordered shape: (32, 4, n_rb, 4, n_cb, L)
        # Target blocked order matching to_blocked() output.
        scale_a = (
            sfa_reordered.permute(2, 4, 0, 1, 3, 5)
            [:, :, :, :, :, 0]
            .contiguous()
            .reshape(-1, 32, 16)
            .flatten()
        )
        scale_b = (
            sfb_reordered.permute(2, 4, 0, 1, 3, 5)
            [:, :, :, :, :, 0]
            .contiguous()
            .reshape(-1, 32, 16)
            .flatten()
        )

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

    return result_tensors