submission 518127

"""
MXFP4 GEMM submission v1 - AITER a4w4 baseline.
Flow: bf16 A -> MXFP4 quant A -> gemm_a4w4(A_q, B_shuffled) -> bf16 C
"""
from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    import aiter
    from aiter import QuantType, dtypes

    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B.shape[0]

    # Quantize A to MXFP4 with shuffle for CK kernel
    quant_func = aiter.get_triton_quant(QuantType.per_1x32)
    A_q, A_scale_sh = quant_func(A, shuffle=True)

    # GEMM via AITER CK kernel (pre-shuffled weights)
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out_gemm