submission 596705

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X

"""
AMD MXFP4 Matrix Multiplication
Target: AMD MI355X GPU (CDNA4/gfx950)

Reference implementation using aiter.gemm_a4w4 (AMD optimized kernels).
Achieves up to ~104 TFLOP/s on large M shapes.

Note: Many test shapes use default configs in aiter (not pre-tuned).
To potentially improve: Custom Triton kernel for shapes without tuned configs.
"""

import torch
from task import input_t, output_t
from utils import make_match_reference
from aiter import QuantType, dtypes
import aiter
from aiter.ops.shuffle import shuffle_weight
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle

# ---------------------------------------------------------------------------
# Quantization helpers
# ---------------------------------------------------------------------------

def _quant_mxfp4(x, shuffle=True):
    """Quantize tensor to MXFP4 format."""
    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)


def generate_input(m: int, n: int, k: int, seed: int):
    """Generate random bf16 inputs and quantized MXFP4 B."""
    assert k % 64 == 0, "k must be divisible by 64"
    gen = torch.Generator(device="cuda")
    gen.manual_seed(seed)
    A = torch.randn((m, k), dtype=torch.bfloat16, device="cuda", generator=gen)
    B = torch.randn((n, k), dtype=torch.bfloat16, device="cuda", generator=gen)
    B_q, B_scale_sh = _quant_mxfp4(B, shuffle=True)
    B_shuffle = shuffle_weight(B_q, layout=(16, 16))
    return (A, B, B_q, B_shuffle, B_scale_sh)


# ---------------------------------------------------------------------------
# Reference kernel using optimized aiter gemm
# ---------------------------------------------------------------------------

def ref_kernel(data: input_t) -> output_t:
    """
    MXFP4 GEMM using aiter.gemm_a4w4.
    
    Steps:
    1. Quantize A to MXFP4 (per-1x32)
    2. Call gemm_a4w4 with bpreshuffle=True for B
    """
    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    B = B.contiguous()  # Match reference exactly
    m, k = A.shape
    n = B.shape[0]

    # Quantize A to MXFP4
    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=True)
    
    # GEMM: A_q @ B_shuffle.T with scales
    out = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    
    return out


# Alias for compatibility
custom_kernel = ref_kernel

# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------

check_implementation = make_match_reference(ref_kernel, rtol=1e-02, atol=1e-02)


if __name__ == "__main__":
    print("Testing MXFP4 GEMM kernel...")
    for m, n, k in [(256, 256, 512), (64, 7168, 2048), (32, 4096, 512)]:
        data = generate_input(m, n, k, seed=42)
        result = ref_kernel(data)
        print(f"  [{m}x{k}] @ [{n}x{k}] -> {result.shape} dtype={result.dtype}")
    print("Done!")