submission 526261

#!POPCORN leaderboard amd-mxfp4-mm
"""
FP4 quant + FP4 GEMM reference: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Quant logic follows aiter op_tests/test_gemm_a4w4.py (get_triton_quant(QuantType.per_1x32)).
"""
import os

import torch

from task import input_t, output_t

VARIANT = os.environ.get("MXFP4_MM_VARIANT", "baseline")
SMALL_M_PAD_THRESHOLD = 32
PADDED_M = 64


def _maybe_pad_a(A: torch.Tensor) -> tuple[torch.Tensor, int]:
    orig_m = A.shape[0]
    if orig_m > SMALL_M_PAD_THRESHOLD:
        return A, orig_m

    padded = torch.zeros((PADDED_M, A.shape[1]), dtype=A.dtype, device=A.device)
    padded[:orig_m].copy_(A)
    return padded, orig_m


def custom_kernel(data: input_t) -> output_t:
    """
    Reference: MXFP4 per-1x32 quant on A; B_shuffle, B_scale_sh from generate_input.
    gemm_a4w4 with bpreshuffle=True.
    """
    import aiter
    from aiter import QuantType, dtypes

    A, _B, _B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()

    if VARIANT == "small_m_pad":
        A_work, orig_m = _maybe_pad_a(A)
    else:
        A_work, orig_m = A, A.shape[0]

    quant_func = aiter.get_triton_quant(QuantType.per_1x32)
    A_q, A_scale_sh = quant_func(A_work, shuffle=True)
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out_gemm[:orig_m]