submission 684603

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

"""
Optimized MXFP4 GEMM with fused quant+shuffle Triton kernel.

Pipeline: bf16 A -> [fused_mxfp4_quant_shuffle] -> gemm_a4w4 -> bf16 C

Key optimization: Fuse dynamic_mxfp4_quant + e8m0_shuffle into ONE Triton kernel.
This eliminates one kernel launch (~3-5 us) and avoids the huge zero-padded
memory copy that standard e8m0_shuffle requires.
"""

import torch
import triton
import triton.language as tl
from task import input_t, output_t

import aiter
from aiter import dtypes

_gemm = aiter.gemm_a4w4
_fp4x2 = dtypes.fp4x2
_fp8_e8m0 = dtypes.fp8_e8m0
_bf16 = dtypes.bf16

# Per-shape cache: (m, k) -> (fp4_buf, scale_buf)
_buf_cache: dict = {}


@triton.jit
def _mxfp4_quant_op_inline(
    x,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
):
    """Quantize x [BLOCK_SIZE_M, BLOCK_SIZE_N] to MXFP4 with E8M0 block scales."""
    EXP_BIAS_FP32: tl.constexpr = 127
    EXP_BIAS_FP4: tl.constexpr = 1
    EBITS_F32: tl.constexpr = 8
    EBITS_FP4: tl.constexpr = 2
    MBITS_F32: tl.constexpr = 23
    MBITS_FP4: tl.constexpr = 1
    max_normal: tl.constexpr = 6
    min_normal: tl.constexpr = 1

    NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_N // MXFP4_QUANT_BLOCK_SIZE
    x = x.reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE)

    amax = tl.max(tl.abs(x), axis=-1, keep_dims=True)
    amax = amax.to(tl.int32, bitcast=True)
    amax = (amax + 0x200000).to(tl.uint32, bitcast=True) & 0xFF800000
    amax = amax.to(tl.float32, bitcast=True)
    scale_e8m0_unbiased = tl.log2(amax).floor() - 2
    scale_e8m0_unbiased = tl.clamp(scale_e8m0_unbiased, min=-127, max=127)

    bs_e8m0 = scale_e8m0_unbiased.to(tl.uint8) + 127
    quant_scale = tl.exp2(-scale_e8m0_unbiased)

    qx = x * quant_scale
    qx = qx.to(tl.uint32, bitcast=True)

    s = qx & 0x80000000
    qx = qx ^ s

    qx_fp32 = qx.to(tl.float32, bitcast=True)
    saturate_mask = qx_fp32 >= max_normal
    denormal_mask = (not saturate_mask) & (qx_fp32 < min_normal)
    normal_mask = not (saturate_mask | denormal_mask)

    denorm_exp: tl.constexpr = (EXP_BIAS_FP32 - EXP_BIAS_FP4) + (MBITS_F32 - MBITS_FP4) + 1
    denorm_mask_int: tl.constexpr = denorm_exp << MBITS_F32
    denorm_mask_float: tl.constexpr = tl.cast(denorm_mask_int, tl.float32, bitcast=True)

    denormal_x = qx_fp32 + denorm_mask_float
    denormal_x = denormal_x.to(tl.uint32, bitcast=True)
    denormal_x -= denorm_mask_int
    denormal_x = denormal_x.to(tl.uint8)

    normal_x = qx
    mant_odd = (normal_x >> (MBITS_F32 - MBITS_FP4)) & 1
    val_to_add = ((EXP_BIAS_FP4 - EXP_BIAS_FP32) << MBITS_F32) + (1 << 21) - 1
    normal_x += val_to_add
    normal_x += mant_odd
    normal_x = normal_x >> (MBITS_F32 - MBITS_FP4)
    normal_x = normal_x.to(tl.uint8)

    e2m1_value = tl.full(qx.type.get_block_shapes(), 0x7, dtype=tl.uint8)
    e2m1_value = tl.where(normal_mask, normal_x, e2m1_value)
    e2m1_value = tl.where(denormal_mask, denormal_x, e2m1_value)

    sign_lp = s >> (MBITS_F32 + EBITS_F32 - MBITS_FP4 - EBITS_FP4)
    sign_lp = sign_lp.to(tl.uint8)
    e2m1_value = e2m1_value | sign_lp

    e2m1_value = tl.reshape(e2m1_value, [BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE // 2, 2])
    evens, odds = tl.split(e2m1_value)
    x_fp4 = evens | (odds << 4)
    x_fp4 = x_fp4.reshape(BLOCK_SIZE_M, BLOCK_SIZE_N // 2)

    return x_fp4, bs_e8m0.reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS)


@triton.jit
def _fused_quant_shuffle_kernel(
    x_ptr,
    x_fp4_ptr,
    bs_shuffled_ptr,    # Output: shuffled scale, flat buffer of size M_pad * K_scale_pad
    stride_x_m, stride_x_n,
    stride_fp4_m, stride_fp4_n,
    M,
    N,
    M_pad: tl.constexpr,
    K_scale_pad: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    NUM_ITER: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
    EVEN_M_N: tl.constexpr,
):
    """Fused MXFP4 quantization + e8m0 scale shuffle in one kernel."""
    pid_m = tl.program_id(0)
    start_n = tl.program_id(1) * NUM_ITER
    NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_N // MXFP4_QUANT_BLOCK_SIZE

    stride_x_m_64 = tl.cast(stride_x_m, tl.int64)
    stride_x_n_64 = tl.cast(stride_x_n, tl.int64)
    stride_fp4_m_64 = tl.cast(stride_fp4_m, tl.int64)
    stride_fp4_n_64 = tl.cast(stride_fp4_n, tl.int64)

    for pid_n in tl.range(start_n, min(start_n + NUM_ITER, N), num_stages=1):
        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_64 + x_offs_n[None, :] * stride_x_n_64

        if EVEN_M_N:
            x = tl.load(x_ptr + x_offs).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).to(tl.float32)

        out_tensor, bs_e8m0 = _mxfp4_quant_op_inline(
            x, BLOCK_SIZE_N, BLOCK_SIZE_M, MXFP4_QUANT_BLOCK_SIZE
        )

        # Store fp4 output
        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_fp4_m_64 + out_offs_n[None, :] * stride_fp4_n_64

        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)

        # Store shuffled scale directly
        # bs_e8m0 shape: [BLOCK_SIZE_M, NUM_QUANT_BLOCKS]
        bs_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)  # original row
        bs_offs_n = pid_n * NUM_QUANT_BLOCKS + tl.arange(0, NUM_QUANT_BLOCKS)  # original col

        # Compute shuffled destination for each (row, col)
        # Decompose: g = row//32, j = (row%32)//16, i = (row%32)%16
        #            h = col//8, l = (col%8)//4, k = (col%8)%4
        # dst_offset = g*(32*K_scale_pad) + h*256 + k*64 + i*4 + l*2 + j

        g = bs_offs_m // 32
        row_in_g = bs_offs_m % 32
        j = row_in_g // 16
        i = row_in_g % 16
        h = bs_offs_n // 8
        col_in_h = bs_offs_n % 8
        l = col_in_h // 4
        k = col_in_h % 4

        dst_offset = (g * (32 * K_scale_pad))[:, None] + (h * 256 + k * 64)[None, :] + (i * 4)[:, None] + (l * 2)[None, :] + j[:, None]

        bs_mask = (bs_offs_m < M_pad)[:, None] & (bs_offs_n < K_scale_pad)[None, :]
        src_mask = (bs_offs_m < M)[:, None] & (bs_offs_n < (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE)[None, :]

        # For padded rows (row >= M), store 0
        final_val = tl.where(src_mask, bs_e8m0, tl.zeros_like(bs_e8m0))
        tl.store(bs_shuffled_ptr + dst_offset, final_val, mask=bs_mask)


def _fused_mxfp4_quant_shuffle(x, m, k):
    """Fused MXFP4 quantization + e8m0 shuffle."""
    M, N = m, k
    MXFP4_QUANT_BLOCK_SIZE = 32

    K_scale = triton.cdiv(N, MXFP4_QUANT_BLOCK_SIZE)
    M_pad = triton.cdiv(M, 256) * 256
    K_scale_pad = triton.cdiv(K_scale, 8) * 8

    buf_key = (M, N)
    bufs = _buf_cache.get(buf_key)
    if bufs is None:
        x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
        scale_buf = torch.zeros(M_pad * K_scale_pad, dtype=torch.uint8, device=x.device)
        bufs = (x_fp4, scale_buf)
        _buf_cache[buf_key] = bufs
    x_fp4, scale_buf = bufs

    # Choose block sizes (matching AITER's heuristics)
    if M <= 32:
        BLOCK_SIZE_M = triton.next_power_of_2(M)
        BLOCK_SIZE_N = 32
        NUM_WARPS = 1
        NUM_ITER = 1
    else:
        BLOCK_SIZE_M = 64
        BLOCK_SIZE_N = 64
        NUM_WARPS = 4
        NUM_ITER = 4

    if N <= 16384:
        BLOCK_SIZE_M = 32
        BLOCK_SIZE_N = 128

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

    EVEN_M_N = (M % BLOCK_SIZE_M == 0) and (N % (BLOCK_SIZE_N * NUM_ITER) == 0)

    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N * NUM_ITER))
    _fused_quant_shuffle_kernel[grid](
        x, x_fp4, scale_buf,
        x.stride(0), x.stride(1),
        x_fp4.stride(0), x_fp4.stride(1),
        M, N,
        M_pad, K_scale_pad,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        NUM_ITER=NUM_ITER,
        MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE,
        EVEN_M_N=EVEN_M_N,
        num_warps=NUM_WARPS,
        num_stages=1,
    )

    return x_fp4, scale_buf.view(M_pad, K_scale_pad)


def custom_kernel(data: input_t) -> output_t:
    A, _, _, B_shuffle, B_scale_sh = data
    m, k = A.shape

    # Fused quant + shuffle in ONE kernel launch
    A_q, A_scale_sh = _fused_mxfp4_quant_shuffle(A, m, k)

    return _gemm(
        A_q.view(_fp4x2),
        B_shuffle,
        A_scale_sh.view(_fp8_e8m0),
        B_scale_sh,
        dtype=_bf16,
        bpreshuffle=True,
    )