submission 634418

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

"""
FP4 quant + FP4 GEMM: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Optimized: fused quant+shuffle Triton kernel eliminates separate e8m0_shuffle step.
"""
from task import input_t, output_t
import torch
import triton
import triton.language as tl
import aiter
from aiter import dtypes

MXFP4_QUANT_BLOCK_SIZE = 32


@triton.jit
def _mxfp4_quant_op(
    x,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
):
    """Inline MXFP4 quant: fp32 x -> (fp4x2 packed, e8m0 scales)."""
    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)

    # Scale: per-block amax -> power-of-2 E8M0
    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)

    # Quantize to FP4 E2M1
    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)

    max_normal: tl.constexpr = 6
    min_normal: tl.constexpr = 1
    saturate_mask = qx_fp32 >= max_normal
    denormal_mask = (not saturate_mask) & (qx_fp32 < min_normal)
    normal_mask = not (saturate_mask | denormal_mask)

    # Denormals
    denorm_exp: tl.constexpr = (127 - 1) + (23 - 1) + 1
    denorm_mask_int: tl.constexpr = denorm_exp << 23
    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)

    # Normals: cast to int32 to allow negative arithmetic, then back to uint32
    normal_x = qx.to(tl.int32, bitcast=True)
    mant_odd = (normal_x >> 22) & 1
    # ((1 - 127) << 23) + (1 << 21) - 1 = -1054867457
    normal_x += -1054867457
    normal_x += mant_odd
    normal_x = normal_x >> 22
    normal_x = normal_x.to(tl.uint8)

    # Merge
    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 >> 28).to(tl.uint8)
    e2m1_value = e2m1_value | sign_lp

    # Pack pairs into fp4x2
    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.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, stride_x_n,
    stride_fp4_m, stride_fp4_n,
    M, N,
    SCALE_N_PAD: 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,
):
    """Fused MXFP4 quant + inline e8m0 shuffle. Writes scales directly in shuffled layout."""
    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

    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)

        if EVEN_M_N:
            x = tl.load(x_ptr + x_offs_m[:, None] * stride_x_m + x_offs_n[None, :] * stride_x_n,
                        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_m[:, None] * stride_x_m + x_offs_n[None, :] * stride_x_n,
                        mask=x_mask, cache_modifier=".cg").to(tl.float32)

        out_fp4, bs_e8m0 = _mxfp4_quant_op(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)
        if EVEN_M_N:
            tl.store(x_fp4_ptr + out_offs_m[:, None] * stride_fp4_m + out_offs_n[None, :] * stride_fp4_n,
                     out_fp4)
        else:
            out_mask = (out_offs_m < M)[:, None] & (out_offs_n < (N // 2))[None, :]
            tl.store(x_fp4_ptr + out_offs_m[:, None] * stride_fp4_m + out_offs_n[None, :] * stride_fp4_n,
                     out_fp4, mask=out_mask)

        # Store scales in shuffled layout (inline e8m0_shuffle)
        # Shuffle permutation: view(sm//32,2,16,sn//8,2,4).permute(0,3,5,2,4,1).view(sm,sn)
        bs_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        bs_offs_n = pid_n * NUM_QUANT_BLOCKS + tl.arange(0, NUM_QUANT_BLOCKS)
        num_bs_cols = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE

        # Decompose into 6D indices for shuffle
        bs_offs_0 = bs_offs_m[:, None] // 32      # a
        bs_offs_1 = bs_offs_m[:, None] % 32
        bs_offs_2 = bs_offs_1 % 16                 # c
        bs_offs_1 = bs_offs_1 // 16                # b
        bs_offs_3 = bs_offs_n[None, :] // 8        # d
        bs_offs_4 = bs_offs_n[None, :] % 8
        bs_offs_5 = bs_offs_4 % 4                  # f
        bs_offs_4 = bs_offs_4 // 4                 # e

        # Compute shuffled flat offset
        bs_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 * SCALE_N_PAD
        )

        # Set padding value to 127 for out-of-range elements
        bs_valid = (bs_offs_m[:, None] < M) & (bs_offs_n[None, :] < num_bs_cols)
        bs_e8m0 = tl.where(bs_valid, bs_e8m0, 127)

        tl.store(bs_ptr + bs_offs, bs_e8m0.to(tl.uint8))


# Cached output buffers
_fp4_buf = {}
_scale_buf = {}


def fused_quant_shuffle(x: torch.Tensor):
    """Fused MXFP4 quant + e8m0 shuffle in a single kernel launch."""
    M, N = x.shape
    m_pad = (M + 255) // 256 * 256
    n_scale = N // MXFP4_QUANT_BLOCK_SIZE
    n_scale_pad = (n_scale + 7) // 8 * 8

    key = (M, N)
    if key not in _fp4_buf:
        _fp4_buf[key] = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
        # Allocate scale buffer in padded shuffled shape
        _scale_buf[key] = torch.full((m_pad * n_scale_pad,), 127, dtype=torch.uint8, device=x.device)

    x_fp4 = _fp4_buf[key]
    bs_flat = _scale_buf[key]

    # Tuning params (matching aiter's dynamic_mxfp4_quant heuristics)
    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 = 4
        BLOCK_SIZE_M = 64
        BLOCK_SIZE_N = 64
        NUM_WARPS = 4
        NUM_STAGES = 2
        if N <= 16384:
            BLOCK_SIZE_M = 32
            BLOCK_SIZE_N = 128

    if N <= 1024:
        NUM_ITER = 1
        NUM_STAGES = 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))

    grid = (
        triton.cdiv(M, BLOCK_SIZE_M),
        triton.cdiv(N, BLOCK_SIZE_N * NUM_ITER),
    )

    _fused_quant_shuffle_kernel[grid](
        x, x_fp4, bs_flat,
        x.stride(0), x.stride(1),
        x_fp4.stride(0), x_fp4.stride(1),
        M=M, N=N,
        SCALE_N_PAD=n_scale_pad,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        NUM_ITER=NUM_ITER,
        NUM_STAGES=NUM_STAGES,
        MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE,
        num_warps=NUM_WARPS,
        waves_per_eu=0,
        num_stages=1,
    )

    return x_fp4, bs_flat.view(m_pad, n_scale_pad)


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

    # Fused quant + shuffle in single kernel
    A_q, A_scale_sh = fused_quant_shuffle(A.contiguous())

    return aiter.gemm_a4w4(
        A_q.view(dtypes.fp4x2), B_shuffle,
        A_scale_sh.view(dtypes.fp8_e8m0), B_scale_sh,
        dtype=dtypes.bf16, bpreshuffle=True,
    )