submission 666956

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

"""
Shape-specialized MXFP4 GEMM submission for the AMD qualifier.

Optimizations over the stock template:
- bypass the generic `gemm_a4w4()` wrapper for known shapes and call
  `gemm_a4w4_asm()` directly with a fixed kernel name
- keep the proven `dynamic_mxfp4_quant()` path for correctness
- avoid generic config lookup and wrapper overhead on the benchmark/test shapes
"""
import os

os.environ.setdefault("AITER_LOG_LEVEL", "WARNING")

import aiter
import torch
from aiter import dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle

from task import input_t, output_t

_KERNELS = {
    # test shapes
    (8, 2112, 7168): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (16, 3072, 1536): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (64, 3072, 1536): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    # benchmark shapes with either exact or nearest-shape ROCm configs
    (4, 2880, 512): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E",
    (16, 2112, 7168): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (32, 4096, 512): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E",
    (32, 2880, 512): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E",
    (64, 7168, 2048): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (256, 3072, 1536): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
}

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


def _gemm_direct(A_q, B_shuffle, A_scale_sh, B_scale_sh, m, n, k):
    kernel_name = _KERNELS.get((m, n, k))
    if not kernel_name:
        return aiter.gemm_a4w4(
            A_q,
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            dtype=dtypes.bf16,
            bpreshuffle=True,
        )

    out = torch.empty(((m + 31) // 32 * 32, n), dtype=torch.bfloat16, device=A_q.device)
    aiter.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,
        0,
    )
    return out[:m]


def _run_impl(A, B_shuffle, B_scale_sh):
    m, k = A.shape
    n = B_shuffle.shape[0]
    A_q, A_scale_sh = _quant_mxfp4_shuffled(A)
    return _gemm_direct(A_q, B_shuffle, A_scale_sh, B_scale_sh, m, n, k)


def custom_kernel(data: input_t) -> output_t:
    """
    Quantize A to MXFP4 and multiply against the pre-shuffled quantized B.
    """
    A, _B, _B_q, B_shuffle, B_scale_sh = data
    if not A.is_contiguous():
        A = A.contiguous()
    return _run_impl(A, B_shuffle, B_scale_sh)