submission 678940

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X

"""
Low-bit layout-normalized retry for mxfp4-mm.

Keep the baseline a4w4 semantics unchanged, but normalize A-side packed buffers
and B-side preshuffled inputs to contiguous layouts before gemm_a4w4 so the
trial stays comparable while probing a smaller layout-handling overhead surface.
"""

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 _as_contiguous(x):
    if x.is_contiguous():
        return x
    return x.contiguous()


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


def custom_kernel(data: input_t) -> output_t:
    A, _, _, B_shuffle, B_scale_sh = data
    A = _as_contiguous(A)
    B_shuffle = _as_contiguous(B_shuffle)
    B_scale_sh = _as_contiguous(B_scale_sh)

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