submission 689569

"""
V2 MXFP4 GEMM: fused A16WFP4 preshuffle (eliminates quant overhead) + ASM fallback.
"""
import torch
import aiter
from aiter import dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle
from aiter.ops.gemm_op_a4w4 import get_GEMM_config
from aiter.jit.utils.chip_info import get_cu_num
from task import input_t, output_t

try:
    from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
    _HAS_FUSED = True
except Exception:
    gemm_a16wfp4_preshuffle = None
    _HAS_FUSED = False

_K32 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E"
_FUSED_FAILED = set()  # (n, k) that failed fused path


def _patch_configs():
    get_GEMM_config(1, 512, 4096)
    d = get_GEMM_config.gemm_dict
    cu = get_cu_num()
    for key in [(cu, 4, 2880, 512), (cu, 16, 2112, 7168),
                (cu, 32, 4096, 512), (cu, 32, 2880, 512)]:
        if key not in d:
            d[key] = {'kernelId': 21, 'splitK': 0, 'us': 0,
                      'kernelName': _K32, 'tflops': 0, 'bw': 0, 'errRatio': 0.0}

_patch_configs()


def _fused_kernel(A, B_shuffle, B_scale_sh):
    """Fused path: A quant happens inside Triton kernel, no separate quant step."""
    n, k_half = B_shuffle.shape
    scale_m, scale_n = B_scale_sh.shape
    # Reshape to preshuffle layout: [N//16, K_half*16]
    w = B_shuffle.view(torch.uint8).view(n // 16, k_half * 16)
    w_scales = B_scale_sh.view(torch.uint8).view(scale_m // 32, scale_n * 32)
    # config=None lets AITER auto-select from tuned configs
    return gemm_a16wfp4_preshuffle(A, w, w_scales, prequant=True, dtype=torch.bfloat16)


def _fallback_kernel(A, B_shuffle, B_scale_sh):
    """Original path: separate quant + ASM GEMM."""
    A_fp4, A_scale = dynamic_mxfp4_quant(A)
    A_scale_sh = e8m0_shuffle(A_scale)
    A_q = A_fp4.view(dtypes.fp4x2)
    A_scale_sh = A_scale_sh.view(dtypes.fp8_e8m0)
    return aiter.gemm_a4w4(A_q, B_shuffle, A_scale_sh, B_scale_sh,
                           dtype=dtypes.bf16, bpreshuffle=True)


def custom_kernel(data: input_t) -> output_t:
    A, B, B_q, B_shuffle, B_scale_sh = data
    m = A.shape[0]
    nk = (B_shuffle.shape[0], A.shape[1])
    # Fused path wins big for small M (quant overhead dominates).
    # For M>=64 the GEMM itself dominates and ASM kernel is faster.
    if _HAS_FUSED and m <= 32 and nk not in _FUSED_FAILED:
        try:
            return _fused_kernel(A, B_shuffle, B_scale_sh)
        except Exception:
            _FUSED_FAILED.add(nk)
    return _fallback_kernel(A, B_shuffle, B_scale_sh)