submission 625855

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

"""
FP4 quant + FP4 GEMM: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Optimized with per-shape CUDA graph caching to eliminate kernel launch overhead.
"""
from task import input_t, output_t

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

# Cache: shape_key -> (graph, static_A, static_B_shuffle, static_B_scale_sh, static_out)
_cache = {}


def custom_kernel(data: input_t) -> output_t:
    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B_shuffle.shape[0]
    key = (m, n, k)

    if key not in _cache:
        static_A = A.clone()
        static_B_shuffle = B_shuffle.clone()
        static_B_scale_sh = B_scale_sh.clone()

        # Warmup (Triton JIT + allocator)
        for _ in range(3):
            A_q, A_scale = dynamic_mxfp4_quant(static_A)
            A_q = A_q.view(dtypes.fp4x2)
            A_scale = e8m0_shuffle(A_scale).view(dtypes.fp8_e8m0)
            _ = aiter.gemm_a4w4(
                A_q, static_B_shuffle, A_scale, static_B_scale_sh,
                dtype=dtypes.bf16, bpreshuffle=True,
            )

        # Capture graph
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            A_q, A_scale = dynamic_mxfp4_quant(static_A)
            A_q = A_q.view(dtypes.fp4x2)
            A_scale = e8m0_shuffle(A_scale).view(dtypes.fp8_e8m0)
            static_out = aiter.gemm_a4w4(
                A_q, static_B_shuffle, A_scale, static_B_scale_sh,
                dtype=dtypes.bf16, bpreshuffle=True,
            )

        _cache[key] = (graph, static_A, static_B_shuffle, static_B_scale_sh, static_out)

    graph, static_A, static_B_shuffle, static_B_scale_sh, static_out = _cache[key]
    static_A.copy_(A)
    static_B_shuffle.copy_(B_shuffle)
    static_B_scale_sh.copy_(B_scale_sh)
    graph.replay()
    return static_out