submission 514200

"""
Aggressive mxfp4-mm variant:
- Keep fast reference pipeline (per-1x32 quant + pre-shuffled B)
- Force blockscale kernel with explicit splitK=3 (with safe fallback)
"""

from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    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.shape

    quant_func = aiter.get_triton_quant(QuantType.per_1x32)
    A_q, A_scale_sh = quant_func(A, shuffle=True)

    # Force explicit split-K path when available.
    try:
        out = A.new_empty(((m + 31) // 32 * 32, n), dtype=dtypes.bf16)
        # API from aiter.ops.gemm_op_a4w4.py
        aiter.gemm_a4w4_blockscale(
            A_q.view(m, k // 2),
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            out,
            splitK=3,
        )
        return out[:m]
    except Exception:
        # Safe fallback to baseline kernel path.
        return aiter.gemm_a4w4(
            A_q,
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            dtype=dtypes.bf16,
            bpreshuffle=True,
        )