submission 742099

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

"""
FP4 quant + FP4 GEMM reference: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Quant logic follows aiter op_tests/test_gemm_a4w4.py (get_triton_quant(QuantType.per_1x32)).
"""

import torch
import triton
import triton.language as tl

from task import input_t, output_t


@triton.jit
def _dynamic_mxfp4_quant_kernel_shuffled(
    x_ptr,
    x_fp4_ptr,
    bs_ptr,
    stride_x_m,
    stride_x_n,
    stride_x_fp4_m,
    stride_x_fp4_n,
    M: tl.constexpr,
    N: tl.constexpr,
    scaleN_valid: tl.constexpr,
    scaleM_pad: tl.constexpr,
    scaleN_pad: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    stride_x_m = tl.cast(stride_x_m, tl.int64)
    stride_x_n = tl.cast(stride_x_n, tl.int64)
    stride_x_fp4_m = tl.cast(stride_x_fp4_m, tl.int64)
    stride_x_fp4_n = tl.cast(stride_x_fp4_n, tl.int64)

    x_offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    x_offs_n = pid_n * MXFP4_QUANT_BLOCK_SIZE + tl.arange(0, MXFP4_QUANT_BLOCK_SIZE)
    x_offs = x_offs_m[:, None] * stride_x_m + x_offs_n[None, :] * stride_x_n
    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)

    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)
    quant_scale = tl.exp2(-scale_e8m0_unbiased)
    qx = x * quant_scale
    bs_e8m0 = scale_e8m0_unbiased.to(tl.uint8) + 127

    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

    qx = qx.to(tl.uint32, bitcast=True)
    sign = qx & 0x80000000
    qx = qx ^ sign

    qx_fp32 = qx.to(tl.float32, bitcast=True)
    saturate_mask = qx_fp32 >= MAX_NORMAL
    denormal_mask = (~saturate_mask) & (qx_fp32 < MIN_NORMAL)
    normal_mask = ~(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([BLOCK_SIZE, MXFP4_QUANT_BLOCK_SIZE], 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 = sign >> (MBITS_F32 + EBITS_F32 - MBITS_FP4 - EBITS_FP4)
    sign_lp = sign_lp.to(tl.uint8)
    e2m1_value = e2m1_value | sign_lp

    packed = tl.reshape(e2m1_value, [BLOCK_SIZE, MXFP4_QUANT_BLOCK_SIZE // 2, 2])
    evens, odds = tl.split(packed)
    out_tensor = evens | (odds << 4)

    out_offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    out_offs_n = pid_n * MXFP4_QUANT_BLOCK_SIZE // 2 + tl.arange(
        0, MXFP4_QUANT_BLOCK_SIZE // 2
    )
    out_offs = (
        out_offs_m[:, None] * stride_x_fp4_m + out_offs_n[None, :] * stride_x_fp4_n
    )
    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_offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    bs_offs_n = pid_n

    bs_offs_0 = bs_offs_m[:, None] // 32
    bs_offs_1 = bs_offs_m[:, None] % 32
    bs_offs_2 = bs_offs_1 % 16
    bs_offs_1 = bs_offs_1 // 16
    bs_offs_3 = bs_offs_n[None, :] // 8
    bs_offs_4 = bs_offs_n[None, :] % 8
    bs_offs_5 = bs_offs_4 % 4
    bs_offs_4 = bs_offs_4 // 4
    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 * scaleN_valid
    )
    bs_mask_valid = (bs_offs_m < M)[:, None] & (bs_offs_n < scaleN_valid)[None, :]
    bs_mask_padded = (bs_offs_m < scaleM_pad)[:, None] & (bs_offs_n < scaleN_pad)[
        None, :
    ]
    bs_e8m0 = tl.where(bs_mask_valid, bs_e8m0, 127)
    tl.store(bs_ptr + bs_offs, bs_e8m0, mask=bs_mask_padded)


def _quant_mxfp4_shuffled_inline(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    m, n = x.shape
    mxfp4_quant_block_size = 32
    x_fp4 = torch.empty((m, n // 2), dtype=torch.uint8, device=x.device)
    scale_n_valid = triton.cdiv(n, mxfp4_quant_block_size)
    scale_n_pad = triton.cdiv(scale_n_valid, 8) * 8
    blockscale_e8m0 = torch.empty(
        (triton.cdiv(m, 256) * 256, scale_n_pad), dtype=torch.uint8, device=x.device
    )

    block_size = 128
    grid = (triton.cdiv(m, block_size), scale_n_pad)
    _dynamic_mxfp4_quant_kernel_shuffled[grid](
        x,
        x_fp4,
        blockscale_e8m0,
        *x.stride(),
        *x_fp4.stride(),
        M=m,
        N=n,
        scaleN_valid=scale_n_valid,
        scaleM_pad=triton.cdiv(m, 32) * 32,
        scaleN_pad=scale_n_pad,
        BLOCK_SIZE=block_size,
        MXFP4_QUANT_BLOCK_SIZE=mxfp4_quant_block_size,
    )
    return x_fp4, blockscale_e8m0


def custom_kernel(data: input_t) -> output_t:
    """
    Reference: MXFP4 per-1x32 quant on A; B_shuffle, B_scale_sh from generate_input.
    gemm_a4w4 with bpreshuffle=True.
    """
    import aiter
    from aiter import dtypes

    A, _, _, B_shuffle, B_scale_sh = data
    m, k = A.shape
    n = B_shuffle.shape[0]

    A_q_raw, A_scale_sh_raw = _quant_mxfp4_shuffled_inline(A)
    A_q = A_q_raw.view(dtypes.fp4x2)
    A_scale_sh = A_scale_sh_raw.view(dtypes.fp8_e8m0)
    if m < 32 or (m == 32 and n in (2880, 4096) and k <= 1024):
        kernel_name = (
            "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E"
            if m == 32 and n == 2880 and k <= 1024
            else "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E"
        )
        out = torch.empty(((m + 31) // 32 * 32, n), dtype=dtypes.bf16, device=A.device)
        out_gemm = aiter.gemm_a4w4_asm(
            A_q,
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            out,
            kernel_name,
            bpreshuffle=True,
            log2_k_split=0,
        )
        return out_gemm[:m]

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