submission 623529

"""
MXFP4-MM 优化版: 使用 aiter.gemm_a4w4

核心优化:
1. 使用 aiter.gemm_a4w4 (已优化的融合实现)
2. MXFP4 量化 (per-1x32 block scaling)
3. 权重预shuffle优化访问模式
"""
from task import input_t, output_t
import torch
import aiter
from aiter import dtypes
from aiter.ops.shuffle import shuffle_weight
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle

SCALE_GROUP_SIZE = 32


def _quant_mxfp4(x, shuffle=True):
    """MXFP4 量化"""
    x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
    if shuffle:
        bs_e8m0 = e8m0_shuffle(bs_e8m0)
    # 注意：返回 dtypes.fp4x2 和 dtypes.fp8_e8m0
    # 这些是aiter的自定义类型，服务器环境支持
    return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)


def custom_kernel(data: input_t) -> output_t:
    """
    优化的 MXFP4 GEMM - 使用 aiter.gemm_a4w4.

    Args:
        data: (A, B, B_q, B_shuffle, B_scale_sh)
            - A: [M, K] bf16, 需要量化
            - B: [N, K] bf16 (原始)
            - B_q: [N, K//2] fp4x2 (已量化)
            - B_shuffle: [N, K//2] fp4x2 (已shuffle)
            - B_scale_sh: [N, K//32] fp8_e8m0 (已shuffle)

    Returns:
        C: [M, N] bf16
    """
    A, B, B_q, B_shuffle, B_scale_sh = data

    # 确保A是连续的
    A = A.contiguous()
    m, k = A.shape
    n, _ = B.shape

    # 量化 A 到 MXFP4
    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=True)

    # 使用 aiter.gemm_a4w4 进行优化的矩阵乘法
    # 这是一个融合kernel，已经针对AMD GPU优化
    C = aiter.gemm_a4w4(
        A_q,           # [M, K//2] fp4x2
        B_shuffle,     # [N, K//2] fp4x2 (pre-shuffled)
        A_scale_sh,    # [M, K//32] fp8_e8m0 (shuffled)
        B_scale_sh,    # [N, K//32] fp8_e8m0 (shuffled)
        dtype=dtypes.bf16,
        bpreshuffle=True,  # B已经预shuffle
    )

    return C