submission 638448

"""
MXFP4 GEMM: tuned kernel selection per shape on MI355X (304 CU).
Based on profiling data from remote MI355X runs.
"""
from task import input_t, output_t
import torch

# Per-shape optimal kernel names from profiling
# Format: (M ranges) -> kernel_name
# Profiling showed these GEMM-only times (quant dominates total time):
#   M=4,N=2880,K=512:   96x128 = 6.71us  
#   M=16,N=2112,K=7168: 32x128 = 7.28us  
#   M=32,N=4096,K=512:  64x256 = 6.85us  
#   M=32,N=2880,K=512:  32x768 = 6.83us  
#   M=64,N=7168,K=2048: 32x128 = 6.79us  
#   M=256,N=3072,K=1536:32x128 = 6.74us

def _get_kernel_name(m, n, k):
    """Select best kernel based on profiling data."""
    # For large K (compute bound): 32x128 is best (minimal tile overhead)
    # For small K (launch bound): larger tiles reduce launch count
    
    tile_m, tile_n = 32, 128  # default
    
    if k >= 4096:
        tile_m, tile_n = 32, 128
    elif n >= 7168:
        tile_m, tile_n = 32, 128
    elif m <= 4:
        tile_m, tile_n = 96, 128
    elif m <= 32:
        tile_m, tile_n = 64, 256
    elif m <= 64:
        tile_m, tile_n = 32, 256
    else:
        tile_m, tile_n = 32, 128

    fname = f"f4gemm_bf16_per1x32Fp4_BpreShuffle_{tile_m}x{tile_n}"
    return f"_ZN5aiter{len(fname)}{fname}E"


def custom_kernel(data: input_t) -> output_t:
    import aiter
    from aiter import dtypes
    from aiter.ops.triton.quant import dynamic_mxfp4_quant
    from aiter.utility.fp4_utils import e8m0_shuffle
    from aiter.ops.gemm_op_a4w4 import gemm_a4w4_asm

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

    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B.shape[0]

    # Quantize A
    A_q, A_scale_sh = _quant_mxfp4(A)

    # Get tuned kernel name
    kname = _get_kernel_name(m, n, k)

    # Pre-allocate output
    padded_m = (m + 31) // 32 * 32
    out = torch.empty((padded_m, n), dtype=dtypes.bf16, device=A.device)

    gemm_a4w4_asm(
        A_q.view(m, k // 2),
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        out,
        kernelName=kname,
        bpreshuffle=True,
        log2_k_split=0,
    )

    return out[:m]