submission 697682

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X
"""
High-Performance MXFP4 GEMM kernel for AMD MI355X
Corrected version ensuring tensor compatibility
"""
import torch

try:
    from aiter import dtypes
    from aiter.ops.triton.quant import dynamic_mxfp4_quant
    from aiter.utility.fp4_utils import e8m0_shuffle
    from aiter import gemm_a4w4
except ImportError:
    pass

def custom_kernel(data):
    A, B, B_q, B_shuffle, B_scale_sh = data

    # 确保输入张量是连续的
    A_c = A.contiguous()

    # 执行量化操作
    a_fp4, a_scale = dynamic_mxfp4_quant(A_c)

    # 执行GEMM操作
    result = gemm_a4w4(
        a_fp4.view(dtypes.fp4x2),
        B_shuffle,
        e8m0_shuffle(a_scale).view(dtypes.fp8_e8m0),
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True
    )

    return result