submission 639201

"""
Optimized amd-mxfp4-mm kernel wrapper.

Flow: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
"""
from __future__ import annotations

from typing import Tuple

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
# from utils import make_match_reference

def _quant_mxfp4_per_1x32_shuffled(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Quantize bf16 -> MXFP4 per-1x32 and shuffle scales for GEMM.

    Returns:
        (x_fp4x2, x_scale_sh) as aiter dtypes.
    """
    x_fp4, x_scale = dynamic_mxfp4_quant(x)
    x_scale_sh = e8m0_shuffle(x_scale)
    return x_fp4.view(dtypes.fp4x2), x_scale_sh.view(dtypes.fp8_e8m0)


def custom_kernel(data: input_t) -> output_t:
    """
    Optimized path: quantize A once, use pre-shuffled B and scales.
    """
    A, _B, _B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()

    A_q, A_scale_sh = _quant_mxfp4_per_1x32_shuffled(A)

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