submission 674628

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

"""
AMD E2E SpeedRun - MXFP4 GEMM Kernel (1000 pts)
=================================================
Block-scale MXFP4 matrix multiplication for MI355X.

Flow: bf16 A -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C

Input: (A, B, B_q, B_shuffle, B_scale_sh)
  A:           [m, k] bf16 (K-major)
  B:           [n, k] bf16 (K-major)
  B_q:         [n, k//2] MXFP4
  B_shuffle:   [n, k//2] MXFP4 (16x16 tile coalesced)
  B_scale_sh:  [*, k//32] E8M0 (padded)

Output: [m, n] bf16

TODO: Optimize using custom Triton/HIP kernels for MI355X
"""
from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    """
    Baseline MXFP4 GEMM using aiter's gemm_a4w4.
    """
    import aiter
    from aiter import QuantType, dtypes
    from aiter.ops.triton.quant import dynamic_mxfp4_quant 
    from aiter.utility.fp4_utils import e8m0_shuffle

    def _quant_mxfp4(x, shuffle=True):
        x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
        if shuffle:
            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()
    B = B.contiguous()
    m, k = A.shape
    n, _ = B.shape

    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=True)
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out_gemm