submission 563102

"""
MXFP4 GEMM for AMD MI355X.

Reference impl (ref_kernel) shows exact API:
  1. dynamic_mxfp4_quant(A) -> (A_fp4, A_scale_e8m0)
  2. e8m0_shuffle(A_scale_e8m0) -> A_scale_sh  (shuffled scales)
  3. aiter.gemm_a4w4(A_q, B_shuffle, A_scale_sh, B_scale_sh,
                     dtype=dtypes.bf16, bpreshuffle=True)
"""
import torch
import aiter
from aiter import dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle


def custom_kernel(data: tuple) -> torch.Tensor:
    A, B, B_q, B_shuffle, B_scale_sh = data

    # Step 1: Quantize A to MXFP4 with per-1x32 block scales
    # dynamic_mxfp4_quant is the #975-patched kernel (correct vs fp4_utils.py)
    A_fp4, A_scale = dynamic_mxfp4_quant(A)

    # Step 2: Shuffle A scales to match the layout gemm_a4w4 expects
    A_scale_sh = e8m0_shuffle(A_scale)

    # Step 3: Cast to the dtypes gemm_a4w4 expects
    A_q = A_fp4.view(dtypes.fp4x2)
    A_scale_sh = A_scale_sh.view(dtypes.fp8_e8m0)

    # Step 4: MXFP4 x MXFP4 GEMM
    # B_shuffle: [N//16, K*16] pre-shuffled layout, bpreshuffle=True
    # B_scale_sh: already shuffled E8M0 scales from generate_input
    out = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out