submission 740565

import torch

import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

from torch.utils.cpp_extension import load_inline
from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    """
    Reference: MXFP4 per-1x32 quant on A; B_shuffle, B_scale_sh from generate_input.
    gemm_a4w4 with bpreshuffle=True.
    """
    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) # A_q stored in fp4 while A_scale_sh in fp8

    

    
    config = aiter.get_GEMM_config(m, n, k)
    print(config)

    if (config): 
        out = torch.empty(((m + 31) // 32 * 32, n), dtype=torch.bfloat16, device=A.device)
        aiter.gemm_a4w4_asm(
            A_q,
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            out,
            config.get("kernelName"),
            None,
            1.0,
            0.0,
            True,
            log2_k_split=2,
        )
    

        return out[:m].view(*A.shape[:-1], n)
    else:
        return aiter.gemm_a4w4(
            A_q,
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            dtype=dtypes.bf16,
            bpreshuffle=True,
        )