submission 740516

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

"""
Hybrid tile selection: auto for M=4,M=64 (where auto wins),
hardcoded for M=16,M=32,M=256 (where hardcoded wins).
Full pre-caching of tensors and views.
"""
from task import input_t, output_t
import torch
import triton
import triton.language as tl
from aiter import dtypes
from aiter.ops.triton._triton_kernels.quant.quant import _mxfp4_quant_op
from aiter.ops.gemm_op_a4w4 import gemm_a4w4_asm

def _make_kernel_name(tile):
    inner = f"f4gemm_bf16_per1x32Fp4_BpreShuffle_{tile}"
    return f"_ZN5aiter{len(inner)}{inner}E"

# Only hardcode tiles where they beat auto-selection
TILE_MAP = {
    (16, 7168, 2112):   _make_kernel_name("32x128"),
    (32, 512, 4096):    _make_kernel_name("32x256"),
    (32, 512, 2880):    _make_kernel_name("32x256"),
    (256, 1536, 3072):  _make_kernel_name("64x128"),
}
# M=4 and M=64 use "" (auto-selection wins)

_cache = {}

@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_mxfp4_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,
    bs_sm: tl.constexpr, bs_sn: tl.constexpr,
    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_idx = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        bs_n_idx = pid_n * NUM_QUANT_BLOCKS + tl.arange(0, NUM_QUANT_BLOCKS)
        g32 = bs_m_idx // 32; r32 = bs_m_idx % 32
        half_row = r32 // 16; r16 = r32 % 16
        g8 = bs_n_idx // 8; r8 = bs_n_idx % 8
        half_col = r8 // 4; r4 = r8 % 4
        stride_d0 = (bs_sn // 8) * 256
        shuffled_offs = (g32[:, None] * stride_d0 + g8[None, :] * 256 +
                        r4[None, :] * 64 + r16[:, None] * 4 +
                        half_col[None, :] * 2 + half_row[:, None] * 1)
        if EVEN_M_N:
            tl.store(bs_ptr + shuffled_offs, bs_e8m0)
        else:
            bs_mask = (bs_m_idx < M)[:, None] & (
                bs_n_idx < (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE
            )[None, :]
            tl.store(bs_ptr + shuffled_offs, bs_e8m0, mask=bs_mask)


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.shape[0]

    key = (M, K, N)
    if key not in _cache:
        MXFP4_QUANT_BLOCK_SIZE = 32
        bs_sm = ((M + 255) // 256) * 256
        bs_sn = ((K // MXFP4_QUANT_BLOCK_SIZE + 7) // 8) * 8
        padded_M = ((M + 31) // 32) * 32

        x_fp4 = torch.empty((M, K // 2), dtype=torch.uint8, device=A.device)
        scale_buf = torch.empty(bs_sm * bs_sn, dtype=torch.uint8, device=A.device)
        out = torch.empty((padded_M, N), dtype=torch.bfloat16, device=A.device)
        fp4_view = x_fp4.view(dtypes.fp4x2)
        scale_view = scale_buf.view(bs_sm, bs_sn).view(dtypes.fp8_e8m0)

        if M <= 32:
            NI, BSM, BSN = 1, triton.next_power_of_2(M), 32
            NW, NS = 1, 1
        else:
            NI, BSM, BSN = 4, 64, 64
            NW, NS = 4, 2
            if K <= 16384: BSM, BSN = 32, 128
        if K <= 1024:
            NI, NS, NW = 1, 1, 4
            BSN = max(32, min(256, triton.next_power_of_2(K)))
            BSM = min(8, triton.next_power_of_2(M))

        grid = (triton.cdiv(M, BSM), triton.cdiv(K, BSN * NI))
        kernel_name = TILE_MAP.get(key, "")  # auto-select if not in map
        _cache[key] = (x_fp4, scale_buf, out, fp4_view, scale_view,
                       bs_sm, bs_sn, grid, BSM, BSN, NW, NI, NS, kernel_name)

    x_fp4, scale_buf, out, fp4_view, scale_view, bs_sm, bs_sn, grid, BSM, BSN, NW, NI, NS, kernel_name = _cache[key]

    _fused_mxfp4_quant_shuffle_kernel[grid](
        A, x_fp4, scale_buf, *A.stride(), *x_fp4.stride(),
        M=M, N=K, bs_sm=bs_sm, bs_sn=bs_sn,
        MXFP4_QUANT_BLOCK_SIZE=32, SCALING_MODE=0,
        NUM_ITER=NI, BLOCK_SIZE_M=BSM, BLOCK_SIZE_N=BSN,
        NUM_STAGES=NS, num_warps=NW, waves_per_eu=0, num_stages=1,
    )
    gemm_a4w4_asm(fp4_view, B_shuffle, scale_view, B_scale_sh,
                  out, kernel_name, None, 1.0, 0.0, True, 0)
    return out[:M]