submission 672650

"""
FP4 quant + FP4 GEMM reference with a small custom A4W4 tuned table.

The default aiter A4W4 tuned CSV does not contain exact entries for several
Phase 1 benchmark shapes, so those cases can fall back to a conservative
kernel choice. We install a tiny runtime CSV that only covers the contest's
test/benchmark shapes and keeps the quantization path unchanged.
"""
import functools
import os
import tempfile
from pathlib import Path

from task import input_t, output_t


_A4W4_TUNED_CSV = """cu_num,M,N,K,kernelId,splitK,us,kernelName,tflops,bw,errRatio
256,8,2112,7168,21,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E,0,0,0
256,16,3072,1536,21,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E,0,0,0
256,64,3072,1536,21,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E,0,0,0
256,256,2880,512,21,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E,0,0,0
256,4,2880,512,29,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E,0,0,0
256,16,2112,7168,21,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E,0,0,0
256,32,4096,512,21,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E,0,0,0
256,32,2880,512,21,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E,0,0,0
256,64,7168,2048,29,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E,0,0,0
256,256,3072,1536,21,0,0.001,_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E,0,0,0
"""


@functools.lru_cache(maxsize=1)
def _install_custom_a4w4_config() -> str:
    path = Path(tempfile.gettempdir()) / "amd_mxfp4_mm_a4w4_tuned_gemm.csv"
    path.write_text(_A4W4_TUNED_CSV, encoding="ascii")
    os.environ["AITER_CONFIG_GEMM_A4W4"] = str(path)
    return str(path)


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.
    """
    _install_custom_a4w4_config()

    import aiter
    from aiter import dtypes
    from aiter.ops.triton.quant import dynamic_mxfp4_quant
    from aiter.utility.fp4_utils import e8m0_shuffle

    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()

    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=True)
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out_gemm