submission 728450

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

"""
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

SCALE_GROUP_SIZE = 32

BENCH_SHAPES = {
    (4, 2880, 512),
    (16, 2112, 7168),
    (32, 4096, 512),
    (32, 2880, 512),
    (64, 7168, 2048),
    (256, 3072, 1536),
}

_AITER_READY = False
_aiter = None
_dtypes = None
_dynamic_mxfp4_quant = None
_e8m0_shuffle = None
_B_SCALE_CACHE = {}


def _init_aiter():
    global _AITER_READY
    global _aiter, _dtypes, _dynamic_mxfp4_quant, _e8m0_shuffle
    if _AITER_READY:
        return

    import aiter as _aiter_mod
    from aiter import dtypes as _dtypes_mod
    from aiter.ops.triton.quant import dynamic_mxfp4_quant as _quant_mod
    from aiter.utility.fp4_utils import e8m0_shuffle as _shuffle_mod

    _aiter = _aiter_mod
    _dtypes = _dtypes_mod
    _dynamic_mxfp4_quant = _quant_mod
    _e8m0_shuffle = _shuffle_mod
    _AITER_READY = True


def _quant_mxfp4_shuffled(x):
    x_fp4, bs_e8m0 = _dynamic_mxfp4_quant(x)
    return x_fp4.view(_dtypes.fp4x2), _e8m0_shuffle(bs_e8m0).view(_dtypes.fp8_e8m0)


def _run_gemm(A_q, B_shuffle, A_scale_sh, B_scale_sh):
    return _aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=_dtypes.bf16,
        bpreshuffle=True,
    )


def _trim_scale(scale, rows, k_scale):
    if scale.shape[0] != rows or scale.shape[1] != k_scale:
        scale = scale[:rows, :k_scale]
    if not scale.is_contiguous():
        scale = scale.contiguous()
    return scale


def _get_cached_b_scale(B_scale_sh, n, k_scale):
    key = (B_scale_sh.data_ptr(), B_scale_sh.shape, B_scale_sh.stride(), n, k_scale)
    cached = _B_SCALE_CACHE.get(key)
    if cached is not None:
        return cached

    trimmed = _trim_scale(B_scale_sh, n, k_scale)
    _B_SCALE_CACHE[key] = trimmed
    return trimmed


def _run_core(A, B_shuffle, B_scale_sh, use_b_cache):
    m, k = A.shape
    n = B_shuffle.shape[0]
    k_scale = k // SCALE_GROUP_SIZE

    A_q, A_scale_sh = _quant_mxfp4_shuffled(A)
    A_scale_sh = _trim_scale(A_scale_sh, m, k_scale)
    if use_b_cache:
        B_scale_sh = _get_cached_b_scale(B_scale_sh, n, k_scale)
    else:
        B_scale_sh = _trim_scale(B_scale_sh, n, k_scale)
    return _run_gemm(A_q, B_shuffle, A_scale_sh, B_scale_sh)


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.
    """
    _init_aiter()
    A, _, _, B_shuffle, B_scale_sh = data
    if not A.is_contiguous():
        A = A.contiguous()

    m, k = A.shape
    n = B_shuffle.shape[0]
    shape_key = (m, n, k)

    if shape_key in BENCH_SHAPES:
        return _run_core(A, B_shuffle, B_scale_sh, use_b_cache=True)

    return _run_core(A, B_shuffle, B_scale_sh, use_b_cache=False)