submission 526982

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

"""
MXFP4 matrix multiplication on AMD MI355X.

Pipeline:
  bf16 A  ->  per-1x32 MXFP4 quant (with hardware shuffle)  ->  a4w4 GEMM  ->  bf16 C

B is provided pre-quantized and pre-shuffled (B_shuffle, B_scale_sh), so only A
needs to be quantized at runtime.  The quantization function is compiled once at
module import time to avoid Triton JIT overhead on the first timed call.
"""

import aiter
from aiter import QuantType, dtypes
from task import input_t, output_t

# Compile the Triton quantization kernel once at import time.
# Keeps it out of the timed hot path and ensures the JIT cache is warm.
_quant_func = aiter.get_triton_quant(QuantType.per_1x32)


def custom_kernel(data: input_t) -> output_t:
    A, B, B_q, B_shuffle, B_scale_sh = data

    # Quantize A: bf16 -> MXFP4 (e2m1) with per-1x32 scales.
    # shuffle=True produces the layout expected by gemm_a4w4 with bpreshuffle=True.
    A_q, A_scale_sh = _quant_func(A.contiguous(), shuffle=True)

    # Hardware-accelerated 4-bit GEMM using pre-shuffled B weights.
    return aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )