submission 581885

"""
Optimized MXFP4 GEMM: bf16 A + preshuffled MXFP4 B -> fused fp4 quant A + gemm -> bf16 C.

Key optimization: Use gemm_a16wfp4_preshuffle which takes bf16 A and quantizes it
on-the-fly inside the GEMM kernel, eliminating the separate dynamic_mxfp4_quant step.

B_shuffle and B_scale_sh are reshaped (zero-copy views) to the format expected by
the Triton preshuffle kernel. Both e8m0_shuffle (used to create B_scale_sh) and
shuffle_scales (expected by gemm_a16wfp4_preshuffle) apply the same permutation;
they differ only in the final view, so B_scale_sh.view(N_pad//32, K)[:N//32]
gives exactly the shuffle_scales-format tensor.
"""
from task import input_t, output_t
import torch


def custom_kernel(data: input_t) -> output_t:
    from aiter import dtypes
    from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle

    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B_shuffle.shape[0]  # N (B_shuffle has shape (N, K//2) in fp4x2)

    # Reshape B_shuffle: (N, K//2) fp4x2 -> (N//16, K//2 * 16) uint8
    # This is a zero-copy view matching the preshuffled weight layout expected by the kernel.
    k_fp4x2 = k // 2
    B_w = B_shuffle.view(torch.uint8).view(n // 16, k_fp4x2 * 16)

    # Reshape B_scale_sh: (N_pad, K//32) e8m0 -> (N//32, K) uint8
    # e8m0_shuffle and shuffle_scales apply the same data permutation.
    # e8m0_shuffle views the result as (N_pad, K//32); shuffle_scales views it as (N//32, K).
    # Since K//32 is divisible by 8 for all benchmark shapes, K_pad = K//32, K_pad*32 = K.
    # N_pad is divisible by 32 (it's a multiple of 256), so view(N_pad//32, K) is exact.
    N_pad = B_scale_sh.shape[0]
    B_scale_w = B_scale_sh.view(torch.uint8).view(N_pad // 32, k)[:n // 32]

    # Fused: quantize A bf16->MXFP4 on-the-fly + GEMM with preshuffled B.
    # Uses tuned Triton configs (specialized for e.g. N=2112, K=7168 with NUM_KSPLIT=14).
    return gemm_a16wfp4_preshuffle(A, B_w, B_scale_w, dtype=dtypes.bf16)