submission 726928

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

"""
MXFP4-MM v42 - Hybrid: fused kernel for small K, ASM for large K
K<=512: gemm_a16wfp4 (fused inline quant + Triton GEMM)
K>512: quant kernel + gemm_a4w4_asm (v37 style)
"""

import torch
import triton
import triton.language as tl
import triton._utils

# Patch Triton dtype
for k, v in {'float4_e2m1fn_x2': 'u8', 'float8_e8m0fnu': 'u8'}.items():
    if k not in triton._utils.type_canonicalisation_dict:
        triton._utils.type_canonicalisation_dict[k] = v

from aiter.utility import dtypes
from aiter.ops.gemm_op_a4w4 import gemm_a4w4_asm
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4


@triton.jit
def _fused_mxfp4_quant_shuffle_kernel(
    x_ptr, x_fp4_ptr, bs_ptr,
    stride_x_m, stride_x_n,
    stride_x_fp4_m, stride_x_fp4_n,
    stride_bs_m, stride_bs_n,
    M: tl.constexpr, N: tl.constexpr,
    scaleN: 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
    MBITS_F32: tl.constexpr = 23
    MBITS_FP4: tl.constexpr = 1
    EBITS_F32: tl.constexpr = 8
    EBITS_FP4: tl.constexpr = 2
    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 >= 6
    denormal_mask = (not saturate_mask) & (qx_fp32 < 1)
    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, MXFP4_QUANT_BLOCK_SIZE // 2, 2])
    evens, odds = tl.split(e2m1_value)
    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
    )
    bs_mask1 = (bs_offs_m < M)[:, None] & (bs_offs_n < scaleN)[None, :]
    bs_mask2 = (bs_offs_m < scaleM_pad)[:, None] & (bs_offs_n < scaleN_pad)[None, :]
    bs_e8m0 = tl.where(bs_mask1, bs_e8m0, 127)
    tl.store(bs_ptr + bs_offs, bs_e8m0, mask=bs_mask2)


K32 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E"
K64 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_64x128E"

_bufs = {}


def _run_asm_path(A, B_shuffle, B_scale_sh, m, n, k):
    """v37-style: separate quant + ASM GEMM"""
    shape_key = (m, k, n)
    if shape_key not in _bufs:
        m_pad = (m + 31) // 32 * 32
        scaleN_valid = triton.cdiv(k, 32)
        scaleN = triton.cdiv(scaleN_valid, 8) * 8
        x_fp4 = torch.empty((m, k // 2), dtype=torch.uint8, device=A.device)
        blockscale = torch.empty(
            (triton.cdiv(m, 256) * 256, scaleN), dtype=torch.uint8, device=A.device,
        )
        out = torch.empty(m_pad, n, dtype=torch.bfloat16, device="cuda")
        if m <= 16:
            kernel_name = K32 if k >= 4096 else K64
        elif m <= 32:
            kernel_name = K32 if k > 1024 else K64
        else:
            kernel_name = K32
        grid = (triton.cdiv(m, 32), scaleN_valid)
        _bufs[shape_key] = (x_fp4, blockscale, out, scaleN_valid, scaleN, m_pad, kernel_name, grid)

    x_fp4, blockscale, out, scaleN_valid, scaleN, scaleM_pad, kernel_name, grid = _bufs[shape_key]

    _fused_mxfp4_quant_shuffle_kernel[grid](
        A, x_fp4, blockscale,
        A.stride(0), A.stride(1), x_fp4.stride(0), x_fp4.stride(1),
        blockscale.stride(0), blockscale.stride(1),
        M=m, N=k, scaleN=scaleN_valid,
        scaleM_pad=scaleM_pad, scaleN_pad=scaleN,
        BLOCK_SIZE=32, MXFP4_QUANT_BLOCK_SIZE=32,
        num_warps=1,
    )

    gemm_a4w4_asm(
        x_fp4.view(dtypes.fp4x2), B_shuffle,
        blockscale.view(dtypes.fp8_e8m0), B_scale_sh, out,
        kernelName=kernel_name, bpreshuffle=True, log2_k_split=0,
    )
    return out[:m, :]


def _run_fused_path(A, B, B_q, m, n, k):
    """Fused: inline quant A + Triton GEMM via gemm_a16wfp4"""
    _, B_bs = dynamic_mxfp4_quant(B)
    B_scale = B_bs.view(dtypes.fp8_e8m0)
    return gemm_a16wfp4(A, B_q, B_scale, atomic_add=False, dtype=torch.bfloat16)[:m, :]


def custom_kernel(data):
    A, B, B_q, B_shuffle, B_scale_sh = data
    m, k = A.shape
    n = B.shape[0]

    if k <= 512:
        return _run_fused_path(A, B, B_q, m, n, k)
    else:
        return _run_asm_path(A, B_shuffle, B_scale_sh, m, n, k)