submission 676519

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X

"""
MXFP4 量化矩阵乘法 - MI355X 优化版本
优先使用 aiter 的优化内核，fallback 到 gemm_a4w4
"""
from task import input_t, output_t
import torch


def custom_kernel(data: input_t) -> output_t:
    """
    MXFP4 GEMM: A @ B^T = C
    
    输入 (A, B, B_q, B_shuffle, B_scale_sh):
        A:          [M, K]      bf16
        B_shuffle:  [N, K/2]    fp4 (preshuffle)
        B_scale_sh: [N, K/32]   e8m0
    """
    A, B, B_q, B_shuffle, B_scale_sh = data
    
    A = A.contiguous()
    M, K = A.shape
    N = B.shape[0]
    
    # MXFP4 量化 A
    from aiter.ops.triton.quant import dynamic_mxfp4_quant
    from aiter.utility.fp4_utils import e8m0_shuffle
    from aiter import dtypes
    
    A_fp4, A_scale = dynamic_mxfp4_quant(A)
    A_scale_sh = e8m0_shuffle(A_scale)
    
    # 使用 aiter gemm_a4w4（MI355X 优化）
    import aiter
    
    out = aiter.gemm_a4w4(
        A_fp4.view(dtypes.fp4x2),
        B_shuffle,
        A_scale_sh.view(dtypes.fp8_e8m0),
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    
    return out