submission 668523

#!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
- keep Python hot-path overhead low during quantization and dispatch
"""
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 = {
    (8, 2112, 7168): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (16, 3072, 1536): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (64, 3072, 1536): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (4, 2880, 512): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E",
    (16, 2112, 7168): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (32, 4096, 512): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (32, 2880, 512): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (64, 7168, 2048): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
    (256, 3072, 1536): "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
}

_gemm = aiter.gemm_a4w4
_gemm_asm = aiter.gemm_a4w4_asm
_quant = dynamic_mxfp4_quant
_scale_shuffle = e8m0_shuffle
_dtype_bf16 = dtypes.bf16
_fp4x2 = dtypes.fp4x2
_fp8_e8m0 = dtypes.fp8_e8m0
_torch_empty = torch.empty
_torch_bf16 = torch.bfloat16


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()
    m, k = A.shape
    n = B_shuffle.shape[0]
    kernel_name = _KERNELS.get((m, n, k))

    A_q, A_scale_sh = _quant(A)
    A_q = A_q.view(_fp4x2)
    if A_scale_sh.shape[0] != m:
        A_scale_sh = A_scale_sh[:m]
    A_scale_sh = _scale_shuffle(A_scale_sh).view(_fp8_e8m0)

    if kernel_name is None:
        return _gemm(
            A_q,
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            dtype=_dtype_bf16,
            bpreshuffle=True,
        )

    out = _torch_empty((((m + 31) // 32) * 32, n), dtype=_torch_bf16, device=A_q.device)
    _gemm_asm(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        out,
        kernel_name,
        None,
        1.0,
        0.0,
        True,
        0,
    )
    return out[:m]