submission 531852

#!POPCORN leaderboard amd-mxfp4-mm
"""
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 torch

from task import input_t, output_t

SPLIT_ROW_SHAPE = (32, 2880, 512)
SPLIT_CHUNK_M = 16


def _run_quant_gemm(aiter, quant_func, dtypes, A: torch.Tensor, B_shuffle: torch.Tensor, B_scale_sh: torch.Tensor) -> torch.Tensor:
    A_q, A_scale_sh = quant_func(A, shuffle=True)
    return aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )


def _run_split_row_two_pass(
    aiter,
    quant_func,
    dtypes,
    A: torch.Tensor,
    B_shuffle: torch.Tensor,
    B_scale_sh: torch.Tensor,
) -> torch.Tensor:
    top = _run_quant_gemm(aiter, quant_func, dtypes, A[:SPLIT_CHUNK_M], B_shuffle, B_scale_sh)
    bottom = _run_quant_gemm(aiter, quant_func, dtypes, A[SPLIT_CHUNK_M:], B_shuffle, B_scale_sh)
    return torch.cat((top, bottom), dim=0)


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.
    """
    import aiter
    from aiter import QuantType, dtypes

    A, _B, _B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B_shuffle.shape[0]

    quant_func = aiter.get_triton_quant(QuantType.per_1x32)
    # Keep the split-row survivor narrowly targeted to the one benchmark shape most likely to benefit.
    if (m, n, k) == SPLIT_ROW_SHAPE:
        return _run_split_row_two_pass(aiter, quant_func, dtypes, A, B_shuffle, B_scale_sh)

    return _run_quant_gemm(aiter, quant_func, dtypes, A, B_shuffle, B_scale_sh)