submission 750737

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

"""
Probe89: Combined best of all probes.
- M=4: BSM=4, BSN=64 (probe86: 6.46μs vs 6.73μs with BSN=128)
- M=16: Default + waves_per_eu=2, num_stages=4 (probe74: 16.5→16.3μs marginal)
- M>32: NUM_ITER=1 quant kernel (probe74: 13.9/12.6μs)
"""
from task import input_t, output_t
import torch
import triton
import triton.language as tl
import aiter
from aiter import dtypes
from aiter.ops.triton._triton_kernels.quant.quant import _mxfp4_quant_op
from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
import aiter.ops.triton.gemm.basic.gemm_a16wfp4 as _gmod

# Monkey-patch _get_config for M=4 optimization
_original_get_config = _gmod._get_config

def _custom_get_config(M, N, K, preshuffle=False):
    config, extra = _original_get_config(M, N, K, preshuffle)
    if preshuffle and M <= 4:
        config = dict(config)
        config["BLOCK_SIZE_M"] = 4
        config["BLOCK_SIZE_N"] = 64
        config["num_warps"] = 4
    elif preshuffle and M <= 16:
        config = dict(config)
        config["waves_per_eu"] = 2
        config["num_stages"] = 4
    return config, extra

_gmod._get_config = _custom_get_config

_buf = {}


@triton.heuristics({
    "EVEN_M_N": lambda args: args["M"] % args["BLOCK_SIZE_M"] == 0
    and args["N"] % (args["BLOCK_SIZE_N"] * args["NUM_ITER"]) == 0,
})
@triton.jit
def _fused_quant_shuffle_kernel(
    x_ptr, x_fp4_ptr, bs_ptr,
    stride_x_m_in, stride_x_n_in, stride_x_fp4_m_in, stride_x_fp4_n_in,
    M, N, scaleN, scaleM_pad, scaleN_pad,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    NUM_ITER: tl.constexpr, NUM_STAGES: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr, EVEN_M_N: tl.constexpr,
    SCALING_MODE: tl.constexpr,
):
    pid_m = tl.program_id(0)
    start_n = tl.program_id(1) * NUM_ITER
    stride_x_m = tl.cast(stride_x_m_in, tl.int64)
    stride_x_n = tl.cast(stride_x_n_in, tl.int64)
    stride_x_fp4_m = tl.cast(stride_x_fp4_m_in, tl.int64)
    stride_x_fp4_n = tl.cast(stride_x_fp4_n_in, tl.int64)
    NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_N // MXFP4_QUANT_BLOCK_SIZE
    for pid_n in tl.range(start_n, min(start_n + NUM_ITER, N), num_stages=NUM_STAGES):
        x_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        x_offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
        x_offs = x_offs_m[:, None] * stride_x_m + x_offs_n[None, :] * stride_x_n
        if EVEN_M_N:
            x = tl.load(x_ptr + x_offs, cache_modifier=".cg").to(tl.float32)
        else:
            x_mask = (x_offs_m < M)[:, None] & (x_offs_n < N)[None, :]
            x = tl.load(x_ptr + x_offs, mask=x_mask, cache_modifier=".cg").to(tl.float32)
        out_tensor, bs_e8m0 = _mxfp4_quant_op(x, BLOCK_SIZE_N, BLOCK_SIZE_M, MXFP4_QUANT_BLOCK_SIZE)
        out_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        out_offs_n = pid_n * BLOCK_SIZE_N // 2 + tl.arange(0, BLOCK_SIZE_N // 2)
        out_offs = out_offs_m[:, None] * stride_x_fp4_m + out_offs_n[None, :] * stride_x_fp4_n
        if EVEN_M_N:
            tl.store(x_fp4_ptr + out_offs, out_tensor)
        else:
            out_mask = (out_offs_m < M)[:, None] & (out_offs_n < (N // 2))[None, :]
            tl.store(x_fp4_ptr + out_offs, out_tensor, mask=out_mask)
        bs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        bs_n_base = pid_n * NUM_QUANT_BLOCKS
        bs_n = bs_n_base + tl.arange(0, NUM_QUANT_BLOCKS)
        bs_offs_0 = (bs_m // 32)[:, None]
        bs_offs_1m = (bs_m % 32)[:, None]
        bs_offs_2 = bs_offs_1m % 16
        bs_offs_1 = bs_offs_1m // 16
        bs_offs_3 = (bs_n // 8)[None, :]
        bs_offs_4m = (bs_n % 8)[None, :]
        bs_offs_5 = bs_offs_4m % 4
        bs_offs_4 = bs_offs_4m // 4
        shuffled_offs = (
            bs_offs_1 + bs_offs_4 * 2 + bs_offs_2 * 2 * 2
            + bs_offs_5 * 2 * 2 * 16 + bs_offs_3 * 2 * 2 * 16 * 4
            + bs_offs_0 * 2 * 16 * scaleN_pad
        )
        bs_mask_valid = (bs_m[:, None] < M) & (bs_n[None, :] < scaleN)
        bs_mask_pad = (bs_m[:, None] < scaleM_pad) & (bs_n[None, :] < scaleN_pad)
        bs_val_safe = tl.where(bs_mask_valid, bs_e8m0, 127)
        tl.store(bs_ptr + shuffled_offs, bs_val_safe, mask=bs_mask_pad)


def _get_kernel_params(M, N):
    if M <= 32:
        NUM_ITER = 1
        BLOCK_SIZE_M = triton.next_power_of_2(M)
        BLOCK_SIZE_N = 32
        NUM_WARPS = 1
        NUM_STAGES = 1
    else:
        NUM_ITER = 1
        BLOCK_SIZE_M = 32
        BLOCK_SIZE_N = 128
        NUM_WARPS = 4
        NUM_STAGES = 1
    if N <= 1024:
        NUM_ITER = 1
        NUM_STAGES = 1
        NUM_WARPS = 4
        BLOCK_SIZE_N = max(32, min(256, triton.next_power_of_2(N)))
        BLOCK_SIZE_M = min(8, triton.next_power_of_2(M))
    return BLOCK_SIZE_M, BLOCK_SIZE_N, NUM_ITER, NUM_STAGES, NUM_WARPS


def custom_kernel(data: input_t) -> output_t:
    A, B, B_q, B_shuffle, B_scale_sh = data
    M, K = A.shape
    N = B.shape[0]

    if M <= 32:
        # ===== Triton fused quant+GEMM =====
        B_ps = B_shuffle.view(torch.uint8).reshape(N // 16, (K // 2) * 16)
        sn = triton.cdiv(K, 32)
        scaleN_pad = triton.cdiv(sn, 8) * 8
        bs_raw = B_scale_sh.view(torch.uint8)
        padded_N = bs_raw.shape[0]
        B_scale_ps = bs_raw.reshape(padded_N // 32, scaleN_pad * 32)[:N // 32].contiguous()
        return gemm_a16wfp4_preshuffle(A, B_ps, B_scale_ps, prequant=True, dtype=torch.bfloat16)
    else:
        # ===== ASM path =====
        MXFP4_QUANT_BLOCK_SIZE = 32
        scaleN_valid = triton.cdiv(K, MXFP4_QUANT_BLOCK_SIZE)
        scaleN_pad = triton.cdiv(scaleN_valid, 8) * 8
        scaleM_pad = triton.cdiv(M, 32) * 32
        scale_M_full = triton.cdiv(M, 256) * 256

        key = (M, N, K)
        if key not in _buf:
            _buf[key] = (
                torch.empty((M, K // 2), dtype=torch.uint8, device=A.device),
                torch.empty(scale_M_full, scaleN_pad, dtype=torch.uint8, device=A.device),
            )
        x_fp4, blockscale_shuffled = _buf[key]

        BLOCK_SIZE_M, BLOCK_SIZE_N, NUM_ITER, NUM_STAGES, NUM_WARPS = _get_kernel_params(M, K)
        grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(K, BLOCK_SIZE_N * NUM_ITER))
        _fused_quant_shuffle_kernel[grid](
            A, x_fp4, blockscale_shuffled,
            *A.stride(), *x_fp4.stride(),
            M=M, N=K, scaleN=scaleN_valid, scaleM_pad=scaleM_pad, scaleN_pad=scaleN_pad,
            MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE, SCALING_MODE=0,
            NUM_ITER=NUM_ITER, BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_N,
            NUM_STAGES=NUM_STAGES, num_warps=NUM_WARPS, waves_per_eu=0, num_stages=1,
        )
        A_q = x_fp4.view(dtypes.fp4x2)
        A_scale_sh = blockscale_shuffled.view(dtypes.fp8_e8m0)
        return aiter.gemm_a4w4(A_q, B_shuffle, A_scale_sh, B_scale_sh,
                               dtype=dtypes.bf16, bpreshuffle=True)