submission 567805

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 for CUDA graphs per shape
_graph_cache = {}
_static_inputs = {}
_static_outputs = {}

def _quant_mxfp4_shuffle(x):
    x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
    bs_e8m0_sh = e8m0_shuffle(bs_e8m0)
    return x_fp4.view(dtypes.fp4x2), bs_e8m0_sh.view(dtypes.fp8_e8m0)

def _run_kernel(A, B_shuffle, B_scale_sh):
    A = A.contiguous()
    A_q, A_scale_sh = _quant_mxfp4_shuffle(A)
    return aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )

def custom_kernel(data):
    A, B, B_q, B_shuffle, B_scale_sh = data
    return _run_kernel(A, B_shuffle, B_scale_sh)