submission 407518

import torch
from task import input_t, output_t

# The harness provides scale factors in a "reordered" tensor layout
_INV_PERM = (5, 2, 4, 0, 1, 3)  # (32,4,RestM,4,RestK,L) -> (L,RestM,RestK,32,4,4)


@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    """
    Optimized implementation using torch._scaled_mm for block-scaled FP4 GEMM.
    """
    abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data

    outputs = []

    # Process each group's GEMM sequentially
    for (a, b, c), (sfa_r, sfb_r), (_, _, _, l) in zip(
        abc_tensors, sfasfb_reordered_tensors, problem_sizes
    ):
        # Pre-permute scale factors (zero-copy view transformation)
        sfa_permuted = sfa_r.permute(_INV_PERM)
        sfb_permuted = sfb_r.permute(_INV_PERM)

        for l_idx in range(l):
            # Use reshape instead of view for flexibility
            scale_a = sfa_permuted[l_idx].reshape(-1)
            scale_b = sfb_permuted[l_idx].reshape(-1)

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

        outputs.append(c)

    return outputs