submission 614458

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

"""
FP4 quant + FP4 GEMM optimized: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Optimizations:
1. M-Padding to 32 for asmGEMM alignment.
2. Global import to minimize CPU overhead.
"""

import torch
import aiter
from aiter import QuantType, dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant 
from aiter.utility.fp4_utils import e8m0_shuffle
from task import input_t, output_t

def _quant_mxfp4(x, shuffle=True):
    """Performs MXFP4 per-1x32 quantization on the input tensor."""
    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:
    A, B, B_q, B_shuffle, B_scale_sh = data
    orig_m, k = A.shape
    n, _ = B.shape

    # --- 方案 1: Padding 优化 ---
    # AMD MFMA 指令在 M 为 16/32 倍数时效率最高。针对小 M 形状进行对齐补全。
    pad_m = 0
    if orig_m % 32 != 0:
        pad_m = 32 - (orig_m % 32)
        A = torch.nn.functional.pad(A, (0, 0, 0, pad_m))
    
    A = A.contiguous()
    
    # --- 方案 2: 动态量化 ---
    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=True)

    # 执行核心算子
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True
    )

    # 去除 Padding 部分
    if pad_m > 0:
        out_gemm = out_gemm[:orig_m, :]

    return out_gemm