submission 754758

#!POPCORN gpu MI355X
"""
Kernel C: MXFP4 GEMM — Block-scaled MXFP4 matrix multiplication
y = Q(A) × Q(B)^T where Q() is MXFP4 per-1×32 block quantization

Target: Beat AITER gemm_a4w4 reference on AMD MI355X (CDNA 4)

Optimizations over reference:
  1. Removed unnecessary .contiguous() call (A is already contiguous from generate_input)
  2. Removed unused imports to reduce module load time
  3. Direct dtype imports to avoid attribute lookup overhead
"""
import subprocess, sys
for _pkg in ["aiter"]:
    try:
        __import__(_pkg)
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", _pkg])

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

# Pre-resolve dtype constants to avoid repeated attribute lookups
_FP4X2 = dtypes.fp4x2
_FP8_E8M0 = dtypes.fp8_e8m0
_BF16 = dtypes.bf16


def custom_kernel(data):
    """
    MXFP4 GEMM: bf16 A → quant to MXFP4 → gemm_a4w4 with pre-quantized B → bf16 C.

    Input:
        data = (A, B, B_q, B_shuffle, B_scale_sh)
        A:          [m, k]    bf16
        B:          [n, k]    bf16 (unused — we use pre-quantized B)
        B_q:        [n, k//2] fp4x2 (raw, unused)
        B_shuffle:  [n, k//2] fp4x2 (pre-shuffled to (16,16) tile layout)
        B_scale_sh: [*, k//32] e8m0  (pre-shuffled scales, * = padded)

    Output:
        C: [m, n] bf16
    """
    A, B, B_q, B_shuffle, B_scale_sh = data

    # Quantize A to MXFP4 with shuffled scales
    A_fp4, A_scale_e8m0 = dynamic_mxfp4_quant(A)
    A_q = A_fp4.view(_FP4X2)
    A_scale_sh = e8m0_shuffle(A_scale_e8m0).view(_FP8_E8M0)

    # GEMM using CK backend with pre-shuffled weights
    return aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=_BF16,
        bpreshuffle=True,
    )