submission 670093

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

"""
MXFP4 GEMM submission with shape-specific dispatch.

This keeps the reference quantization path for correctness and only overrides
the GEMM kernel for shapes where an asm kernel is known to be beneficial.
Large-shape overrides that did not beat the public baseline are intentionally
left to AITER's default selection logic.
"""

from task import input_t, output_t


_ASM_32X128 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E"
_ASM_64X128 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E"

_EXACT_KERNEL_OVERRIDES = {
    (8, 2112, 7168): (_ASM_32X128, 0),
    (16, 3072, 1536): (_ASM_32X128, 0),
    (64, 3072, 1536): (_ASM_32X128, 0),
}


def _select_kernel(m: int, n: int, k: int):
    exact = _EXACT_KERNEL_OVERRIDES.get((m, n, k))
    if exact is not None:
        return exact

    if k == 512 and m <= 32 and n in (2880, 4096):
        return (_ASM_64X128, 0)

    if (n, k) == (2112, 7168) and m <= 16:
        return (_ASM_32X128, 0)

    return None


def custom_kernel(data: input_t) -> output_t:
    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

    try:
        gemm_a4w4_asm = aiter.gemm_a4w4_asm
    except AttributeError:
        from aiter.ops.gemm_op_a4w4 import gemm_a4w4_asm

    def _quant_mxfp4(x):
        x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
        bs_e8m0 = e8m0_shuffle(bs_e8m0)
        return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)

    a, _b, _b_q, b_shuffle, b_scale_sh = data
    if not a.is_contiguous():
        a = a.contiguous()

    m, k = a.shape
    n = b_shuffle.shape[0]

    a_q, a_scale_sh = _quant_mxfp4(a)
    kernel_plan = _select_kernel(m, n, k)

    if kernel_plan is None:
        return aiter.gemm_a4w4(
            a_q,
            b_shuffle,
            a_scale_sh,
            b_scale_sh,
            dtype=dtypes.bf16,
            bpreshuffle=True,
        )

    kernel_name, split_k = kernel_plan
    out = torch.empty((((m + 31) // 32) * 32, n), dtype=dtypes.bf16, device=a.device)
    gemm_a4w4_asm(
        a_q.view(m, k // 2),
        b_shuffle,
        a_scale_sh,
        b_scale_sh,
        out,
        kernel_name,
        None,
        1.0,
        0.0,
        True,
        split_k,
    )
    return out[:m]