submission 698374

"""
Candidate B promoted:
- pin (16, 2112, 7168) to 32x128 based on nearby tuned rows
- pin (32, 4096, 512) to 64x128 based on the exact tuned 4096x512 family
- keep exact tuned 32x128 choices for the larger benchmark shapes
- leave the remaining shapes on AITER's default heuristic path
"""
import functools
import os

import torch

from task import input_t, output_t

os.environ.setdefault("GPU_ARCHS", "gfx950")
os.environ.setdefault("CU_NUM", "256")

from aiter import dtypes
from aiter.ops.gemm_op_a4w4 import gemm_a4w4_asm, gemm_a4w4_blockscale, get_GEMM_config
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle


def _quant_mxfp4(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
    bs_e8m0 = e8m0_shuffle(bs_e8m0)
    return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)


_FIXED_KERNELS: dict[tuple[int, int, int], tuple[str, int, bool]] = {
    (16, 2112, 7168): (
        "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
        0,
        False,
    ),
    (32, 4096, 512): (
        "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E",
        0,
        False,
    ),
    (64, 7168, 2048): (
        "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
        0,
        False,
    ),
    (256, 3072, 1536): (
        "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E",
        0,
        False,
    ),
}


@functools.lru_cache(maxsize=128)
def _lookup_kernel(m: int, n: int, k: int) -> tuple[str, int, bool]:
    fixed = _FIXED_KERNELS.get((m, n, k))
    if fixed is not None:
        return fixed

    config = get_GEMM_config(m, n, k)
    if config is None:
        return "", 0, False
    kernel_name = config["kernelName"]
    split_k = config.get("splitK", None)
    use_blockscale = kernel_name.find("_ZN") == -1
    return kernel_name, (0 if split_k is None else split_k), use_blockscale


def custom_kernel(data: input_t) -> output_t:
    A, _B, _B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B_shuffle.shape[0]

    A_q, A_scale_sh = _quant_mxfp4(A)
    kernel_name, split_k, use_blockscale = _lookup_kernel(m, n, k)

    out = torch.empty(((m + 31) // 32 * 32, n), dtype=dtypes.bf16, device=A.device)
    if use_blockscale:
        return gemm_a4w4_blockscale(
            A_q.view(m, k // 2),
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            out,
            splitK=split_k,
        )[:m]

    gemm_a4w4_asm(
        A_q.view(m, k // 2),
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        out,
        kernel_name,
        None,
        1.0,
        0.0,
        True,
        split_k,
    )
    return out[:m]