submission 745270

"""
amd-mxfp4-mm — GPU MODE Hackathon submission
Competition : amd-mxfp4-mm
"""

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

BLOCK_M      = 16
BLOCK_N      = 32
BLOCK_K      = 64
NUM_STAGES   = 3
NUM_WARPS    = 8
SPLIT_K      = 2
FUSE_QUANT   = True
PRESHUFFLE_A = True

SCALE_GROUP_SIZE = 32


def _quant_mxfp4(x: torch.Tensor, shuffle: bool = True):
    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 = 42):
    """Generate MXFP4-MM input tensors on CUDA."""
    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)


def custom_kernel(data: input_t) -> output_t:
    """
    MXFP4 per-1x32 quantize A, then gemm_a4w4 -> bf16 C.
    """
    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()

    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=PRESHUFFLE_A)

    out = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out


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