submission 695351

"""
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)).
"""
import aiter
import torch
from aiter import 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

FALLBACK_SHAPES = {
    (4, 2880, 512),
    (16, 2112, 7168),
    (32, 4096, 512),
    (32, 2880, 512),
}
ASM_KERNEL_32X128 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E"


def _quant_a(A):
    A_q, A_scale_sh = dynamic_mxfp4_quant(A)
    return A_q.view(dtypes.fp4x2), e8m0_shuffle(A_scale_sh).view(dtypes.fp8_e8m0)


def _run_default(A_q, B_shuffle, A_scale_sh, B_scale_sh):
    return aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )


def _run_forced_asm(A_q, B_shuffle, A_scale_sh, B_scale_sh):
    m = A_q.numel() // A_q.shape[-1]
    n = B_shuffle.shape[0]
    out = torch.empty(((m + 31) // 32 * 32, n), dtype=torch.bfloat16, device=A_q.device)
    return aiter.gemm_a4w4_asm(
        A_q.view(m, -1),
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        out,
        ASM_KERNEL_32X128,
        bpreshuffle=True,
        log2_k_split=0,
    )[:m]


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.
    """
    A, _, _, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B_shuffle.shape[0]
    A_q, A_scale_sh = _quant_a(A)

    if (m, n, k) in FALLBACK_SHAPES:
        return _run_forced_asm(A_q, B_shuffle, A_scale_sh, B_scale_sh)

    return _run_default(A_q, B_shuffle, A_scale_sh, B_scale_sh)