submission 630208

"""
FP4 quant + FP4 GEMM reference: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
"""
from task import input_t, output_t
import aiter
from aiter import dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle

def _quant_mxfp4(x, shuffle=True):
    """Quantize input tensor to MXFP4 format with optional scale shuffle."""
    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)

def custom_kernel(data: input_t) -> output_t:
    """
    Performs MXFP4 GEMM:
    1. Quantizes A (bf16) to MXFP4 with per-1x32 scaling.
    2. Calls aiter.gemm_a4w4 with pre-quantized B.
    """
    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()                # ensure contiguous layout

    # Quantize A to MXFP4 and shuffle its scales
    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=True)

    # Perform GEMM using aiter's optimized kernel
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,              # scales are already shuffled
    )
    return out_gemm