submission 689739

import torch
import aiter
from aiter import dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle

from task import input_t, output_t


def quant_mxfp4(x):
    x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
    bs_e8m0 = e8m0_shuffle(bs_e8m0)
    return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)


def custom_kernel(data: input_t) -> output_t:
    """
    Optimized kernel wrapper
    """

    A, _, _, B_shuffle, B_scale_sh = data

    if not A.is_contiguous():
        A = A.contiguous()

    m_orig, k_orig = A.shape

    if m_orig < 32:
        A_q, A_scale_sh = quant_mxfp4(A)
        out = aiter.gemm_a4w4(
            A_q, B_shuffle, A_scale_sh, B_scale_sh,
            dtype=dtypes.bf16, bpreshuffle=True,
        )
        return out

    pad_m = (32 - m_orig % 32) % 32
    if pad_m > 0:
        # Padding A 矩阵
        A = torch.nn.functional.pad(A, (0, 0, 0, pad_m))

    A_q, A_scale_sh = quant_mxfp4(A)

    out = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )

    if  pad_m > 0:
        out = out[:m_orig, :]

    return out