submission 572505

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

"""
Optimized FP4 quant + FP4 GEMM wrapper.

This uses the public aiter GEMM path and a reference-compatible MXFP4
quantization path, with compatibility fallbacks across aiter builds.

The optimization here is intentionally conservative: one-time runtime setup and
minimal per-call Python work.
"""
from task import input_t, output_t


_RUNTIME_READY = False
_GEMM_A4W4 = None
_BF16 = None
_FP4X2 = None
_FP8_E8M0 = None
_DYNAMIC_MXFP4_QUANT = None
_E8M0_SHUFFLE = None
_QUANT_MXFP4 = None


def _ensure_runtime() -> None:
    global _RUNTIME_READY
    global _GEMM_A4W4
    global _BF16
    global _FP4X2
    global _FP8_E8M0
    global _DYNAMIC_MXFP4_QUANT
    global _E8M0_SHUFFLE
    global _QUANT_MXFP4

    if _RUNTIME_READY:
        return

    import aiter
    from aiter import dtypes

    _GEMM_A4W4 = aiter.gemm_a4w4
    _BF16 = dtypes.bf16
    _FP4X2 = dtypes.fp4x2
    _FP8_E8M0 = dtypes.fp8_e8m0

    try:
        from aiter.ops.triton.quant import dynamic_mxfp4_quant
        from aiter.utility.fp4_utils import e8m0_shuffle

        _DYNAMIC_MXFP4_QUANT = dynamic_mxfp4_quant
        _E8M0_SHUFFLE = e8m0_shuffle
    except Exception:
        _DYNAMIC_MXFP4_QUANT = None
        _E8M0_SHUFFLE = None

    try:
        from aiter import QuantType

        _QUANT_MXFP4 = aiter.get_triton_quant(QuantType.per_1x32)
    except Exception:
        _QUANT_MXFP4 = None
    _RUNTIME_READY = True


def _quantize_a(A):
    if _DYNAMIC_MXFP4_QUANT is not None and _E8M0_SHUFFLE is not None:
        A_q_u8, A_scale_u8 = _DYNAMIC_MXFP4_QUANT(A)
        A_scale_u8 = _E8M0_SHUFFLE(A_scale_u8)
        return A_q_u8.view(_FP4X2), A_scale_u8.view(_FP8_E8M0)

    if _QUANT_MXFP4 is not None:
        try:
            return _QUANT_MXFP4(A, shuffle=True)
        except TypeError:
            return _QUANT_MXFP4(A)

    raise RuntimeError("No compatible MXFP4 quantization path is available.")


def custom_kernel(data: input_t) -> output_t:
    _ensure_runtime()

    A, _B, _B_q, B_shuffle, B_scale_sh = data

    if A.dim() != 2 or A.stride(-1) != 1:
        A = A.contiguous()
    if B_shuffle.dim() != 2 or B_shuffle.stride(-1) != 1:
        B_shuffle = B_shuffle.contiguous()
    if B_scale_sh.stride(-1) != 1:
        B_scale_sh = B_scale_sh.contiguous()

    gemm_a4w4 = _GEMM_A4W4

    A_q, A_scale_sh = _quantize_a(A)
    return gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=_BF16,
        bpreshuffle=True,
    )