submission 123301

import torch
from typing import Tuple

# Competition harness signature
input_t = Tuple[
    torch.Tensor,  # a
    torch.Tensor,  # b
    torch.Tensor,  # sfa
    torch.Tensor,  # sfb
    torch.Tensor,  # sfa_permuted
    torch.Tensor,  # sfb_permuted
    torch.Tensor,  # c
]
output_t = torch.Tensor

_scaled_mm = torch._scaled_mm


def custom_kernel(data: input_t) -> output_t:
    """Optimized NVFP4 block-scaled GEMM using pre-permuted GPU scales."""
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c_out = data

    M, N, L = c_out.shape

    sfa_flat = sfa_permuted.permute(5, 2, 4, 0, 1, 3).reshape(L, -1)
    sfb_flat = sfb_permuted.permute(5, 2, 4, 0, 1, 3).reshape(L, -1)

    if L == 1:
        a_slice = a.view(M, -1)
        b_slice = b.view(N, -1).t()
        out = _scaled_mm(
            a_slice,
            b_slice,
            sfa_flat[0],
            sfb_flat[0],
            bias=None,
            out_dtype=torch.float16,
        )
        return out.view(M, N, 1)

    result = torch.empty((M, N, L), dtype=torch.float16, device=a.device)

    if L == 2:
        result[:, :, 0] = _scaled_mm(a[:, :, 0], b[:, :, 0].t(), sfa_flat[0], sfb_flat[0], bias=None, out_dtype=torch.float16)
        result[:, :, 1] = _scaled_mm(a[:, :, 1], b[:, :, 1].t(), sfa_flat[1], sfb_flat[1], bias=None, out_dtype=torch.float16)
        return result

    if L == 3:
        result[:, :, 0] = _scaled_mm(a[:, :, 0], b[:, :, 0].t(), sfa_flat[0], sfb_flat[0], bias=None, out_dtype=torch.float16)
        result[:, :, 1] = _scaled_mm(a[:, :, 1], b[:, :, 1].t(), sfa_flat[1], sfb_flat[1], bias=None, out_dtype=torch.float16)
        result[:, :, 2] = _scaled_mm(a[:, :, 2], b[:, :, 2].t(), sfa_flat[2], sfb_flat[2], bias=None, out_dtype=torch.float16)
        return result

    if L == 4:
        result[:, :, 0] = _scaled_mm(a[:, :, 0], b[:, :, 0].t(), sfa_flat[0], sfb_flat[0], bias=None, out_dtype=torch.float16)
        result[:, :, 1] = _scaled_mm(a[:, :, 1], b[:, :, 1].t(), sfa_flat[1], sfb_flat[1], bias=None, out_dtype=torch.float16)
        result[:, :, 2] = _scaled_mm(a[:, :, 2], b[:, :, 2].t(), sfa_flat[2], sfb_flat[2], bias=None, out_dtype=torch.float16)
        result[:, :, 3] = _scaled_mm(a[:, :, 3], b[:, :, 3].t(), sfa_flat[3], sfb_flat[3], bias=None, out_dtype=torch.float16)
        return result

    a_slices = [a[:, :, i] for i in range(L)]
    b_slices = [b[:, :, i].t() for i in range(L)]

    for l_idx in range(L):
        result[:, :, l_idx] = _scaled_mm(
            a_slices[l_idx],
            b_slices[l_idx],
            sfa_flat[l_idx],
            sfb_flat[l_idx],
            bias=None,
            out_dtype=torch.float16,
        )

    return result