submission 584545

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

_K_SMALL = '_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E'
_K_MED   = '_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E'
_K_LARGE = '_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_128x128E'

out_cache = {}

def _pick_kernel(m):
    if m <= 32: return _K_SMALL
    if m <= 64: return _K_MED
    return _K_LARGE

def quant_a(a):
    qd, sc = dynamic_mxfp4_quant(a)
    return qd.view(aiter.dtypes.fp4x2), e8m0_shuffle(sc).view(aiter.dtypes.fp8_e8m0)

def custom_kernel(data: input_t) -> output_t:
    a, _, _, bsh, bsc = data
    a = a.contiguous()
    m, k = a.shape
    n = bsh.shape[0]
    key = (m, n, k)
    if key not in out_cache:
        out_cache[key] = torch.empty((m, n), dtype=torch.bfloat16, device=a.device)
    aq, asc = quant_a(a)
    return aiter.gemm_a4w4_asm(
        aq, bsh, asc, bsc,
        out_cache[key],
        _pick_kernel(m),
        bpreshuffle=True,
    )

check_implementation = make_match_reference(custom_kernel, rtol=1e-02, atol=1e-02)