submission 642975

"""
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)).
"""
from task import input_t, output_t

_AIT = None
_DTYPES = None
_QUANTTYPE = None
_E8M0_SHUFFLE = None
_QUANT_PER_1X32 = None
_DYNAMIC_QUANT = None
_GEMM_A4W4 = None


def _init_aiter():
    """
    One-time init for imports/callables used in the hot path.
    Keeps import-time dependencies minimal for packaging.
    """
    global _AIT, _DTYPES, _QUANTTYPE, _E8M0_SHUFFLE, _DYNAMIC_QUANT, _GEMM_A4W4
    if _AIT is not None:
        return

    import aiter as _ait
    from aiter import dtypes as _dtypes
    from aiter import QuantType as _QuantType
    from aiter.utility.fp4_utils import e8m0_shuffle as _e8m0_shuffle

    _AIT = _ait
    _DTYPES = _dtypes
    _QUANTTYPE = _QuantType
    _E8M0_SHUFFLE = _e8m0_shuffle
    _GEMM_A4W4 = _ait.gemm_a4w4

    # Prefer the explicitly patched quant kernel (ROCm/aiter#975), as in the reference.
    try:
        from aiter.ops.triton.quant import dynamic_mxfp4_quant as _dyn
    except Exception:
        _dyn = None
    _DYNAMIC_QUANT = _dyn


def _get_quant_per_1x32():
    _init_aiter()
    global _QUANT_PER_1X32
    if _QUANT_PER_1X32 is None:
        _QUANT_PER_1X32 = _AIT.get_triton_quant(_QUANTTYPE.per_1x32)
    return _QUANT_PER_1X32


def _quant_mxfp4_per_1x32_shuffled(a_bf16):
    """
    Returns:
      a_fp4x2: [m, k//2] packed FP4 (dtypes.fp4x2)
      a_scale_sh: [>=m, k//32] E8M0, shuffled (dtypes.fp8_e8m0)
    """
    _init_aiter()

    dyn = _DYNAMIC_QUANT
    if dyn is not None:
        a_fp4, a_scale = dyn(a_bf16)
    else:
        a_fp4, a_scale = _get_quant_per_1x32()(a_bf16)

    a_scale_sh = _E8M0_SHUFFLE(a_scale)
    return a_fp4.view(_DTYPES.fp4x2), a_scale_sh.view(_DTYPES.fp8_e8m0)


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.
    """
    A, B, B_q, B_shuffle, B_scale_sh = data
    # Only A is consumed in bf16; B/B_q are provided for convenience but unused.
    A = A.contiguous()

    _init_aiter()
    A_q, A_scale_sh = _quant_mxfp4_per_1x32_shuffled(A)
    out_gemm = _GEMM_A4W4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=_DTYPES.bf16,
        bpreshuffle=True,
    )
    return out_gemm