submission 621190

"""
FP4 quant + FP4 GEMM reference: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Quant logic follows aiter op_tests/test_gemm_a4w4.py (get_triton_quant(QuantType.per_1x32)).
"""
from task import input_t, output_t
import torch
import triton
from typing import Optional
import aiter
from aiter.ops.triton.gemm.basic.gemm_afp4wfp4 import (
    gemm_afp4wfp4 as triton_gemm_afp4wfp4,
    gemm_afp4wfp4_preshuffled_weight_scales,
)
from aiter.ops.triton.gluon.gemm_afp4wfp4 import gemm_afp4wfp4 as gluon_gemm_afp4wfp4
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


    ##################### my code #####################

    import aiter.ops.triton.utils._triton.arch_info as arch_info
    from aiter.ops.triton.utils.logger import AiterTritonLogger
    from aiter.ops.triton.utils.common_utils import serialize_dict, deserialize_str
    from aiter.ops.triton._triton_kernels.gemm.basic.gemm_afp4wfp4 import (
        _gemm_afp4wfp4_kernel,
        _gemm_afp4wfp4_preshuffle_kernel,
        _gemm_afp4wfp4_kernel_preshuffle_scales,
        _gemm_afp4wfp4_reduce_kernel,
        _get_config,
    )
    from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
    from aiter.jit.utils.torch_guard import torch_compile_guard

    import os
    from aiter.utility.triton.triton_metadata_redirect import AOTMetadataContext

    _LOGGER = AiterTritonLogger()

    global _USE_GEMM_SPLITK_BF16
    _USE_GEMM_SPLITK_BF16 = False

    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 = _quant_mxfp4(A, shuffle=False)
    B_q, B_scale = _quant_mxfp4(B, shuffle=False)
    # gemm_afp4wfp4_preshuffle
    # aiter.gemm_a4w4
    # out_gemm = aiter.gemm_a4w4(
    #     A_q,
    #     B_shuffle,
    #     A_scale_sh,
    #     B_scale_sh,
    #     dtype=dtypes.bf16,
    #     bpreshuffle=True,
    # )
    # 将 dtypes.fp4x2 转换为 uint8
    # A_q_uint8 = A_q.view(torch.uint8)
    # B_shuffle_uint8 = B_shuffle.view(torch.uint8)
    # A_scale_uint8 = A_scale_sh.view(torch.uint8)
    # B_scale_uint8 = B_scale_sh.view(torch.uint8)
    # # 然后调用
    A_q_uint8 = A_q.view(torch.uint8)
    B_q_uint8 = B_q.view(torch.uint8)
    A_scale_uint8 = A_scale.view(torch.uint8)
    B_scale_uint8 = B_scale.view(torch.uint8)
    
    out_gemm = triton_gemm_afp4wfp4(
        A_q_uint8,
        B_q_uint8,
        A_scale_uint8,
        B_scale_uint8,
        dtype=dtypes.bf16,
    )
    return out_gemm