submission 645644

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

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

The stock aiter wrapper falls back to a default kernel for several AMD qualifier
shapes. This submission keeps the reference quantization path but overrides the
GEMM kernel for the shapes where a tuned asm kernel is known to work well.
"""

from task import input_t, output_t


_ASM_32X128 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E"
_ASM_64X128 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E"

# Curated from ROCm/aiter a4w4_blockscale_tuned_gemm.csv. We keep the override
# table small and explicit so every non-reference path is easy to reason about.
_EXACT_KERNEL_OVERRIDES = {
    (8, 2112, 7168): (_ASM_32X128, 0),
    (16, 3072, 1536): (_ASM_32X128, 0),
    (64, 3072, 1536): (_ASM_32X128, 0),
    (64, 7168, 2048): (_ASM_32X128, 0),
    (256, 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

    # The tuned table has no exact entries for the qualifier's low-latency
    # K=512 shapes, so we use a dedicated small-M path instead of the generic
    # wrapper fallback.
    if k == 512 and m <= 32 and n in (2880, 4096):
        return (_ASM_64X128, 0)

    # The tuned table contains M=1/8 entries for this family and they all pick
    # the same 32x128 asm kernel, so reuse it for the benchmark/test shape.
    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
    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]