submission 711628

import torch
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):
    """Reference quantization logic from aiter."""
    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):
    """
    The entrypoint MUST be called custom_kernel.
    The evaluation script passes 'data' which is the tuple:
    (A, B, B_q, B_shuffle, B_scale_sh)
    """
    A, B, B_q, B_shuffle, B_scale_sh = data
    
    # 1) Quantize A to MXFP4 (Required by benchmark flow)
    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=True)
    
    # 2) GEMM using aiter's optimized MI355X kernel
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    
    return out_gemm