submission 736725

"""
MXFP4 GEMM - Kernel Agent v16 for MI355X
=========================================
Optimizations:
1. Custom quant kernel with dynamic BLOCK_SIZE (v15)
2. Avoid B.contiguous() in Triton path - B is already contiguous from generate_input
3. Pre-allocate output tensor for ASM path
4. Minimal Python overhead in hot path

Dispatch (v7 vs v9 profiling):
  M=4,K=512:   Triton 14.8µs vs ASM 19.4µs -> Triton
  M=16,K=7168: Triton 36.2µs vs ASM 33.5µs -> ASM
  M=32,K=512:  Triton 14.2µs vs ASM 20.0µs -> Triton
  M=64,K=2048: Triton 28.5µs vs ASM 24.4µs -> ASM
  M=256,K=1536: Triton 24.8µs vs ASM 23.1µs -> ASM

Rule: M<=32 AND K<=512 -> Triton (custom quant), else -> ASM (reuse B)
"""
import torch
import triton
import triton.language as tl
from task import input_t, output_t
from utils import make_match_reference
from aiter import QuantType, dtypes
import aiter
from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
from aiter.utility.fp4_utils import e8m0_shuffle

SCALE_GROUP_SIZE = 32


# ============================================================
# Optimized MXFP4 Quant Kernel (replaces aiter's hardcoded BLOCK_SIZE=128)
# ============================================================

@triton.jit
def _optimized_mxfp4_quant_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,
    BLOCK_SIZE: tl.constexpr,
    SHUFFLE: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr = 32

    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)

    # Scale calculation
    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

    # FP4 conversion
    EXP_BIAS_FP32: tl.constexpr = 127
    EXP_BIAS_FP4: tl.constexpr = 1
    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)
    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 + 8 - MBITS_FP4 - 2)
    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)

    # Store scales
    bs_offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    bs_offs_n = pid_n

    if SHUFFLE:
        scaleM_pad: tl.constexpr = ((M + 255) // 256) * 256
        scaleN_pad: tl.constexpr = ((scaleN + 7) // 8) * 8
        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)
    else:
        bs_offs = bs_offs_m[:, None] * stride_bs_m + bs_offs_n[None, :] * stride_bs_n
        bs_mask = (bs_offs_m < M)[:, None] & (bs_offs_n < N)[None, :]
        tl.store(bs_ptr + bs_offs, bs_e8m0, mask=bs_mask)


def _get_block_size(m: int) -> tuple:
    """Get optimal BLOCK_SIZE and num_warps for given M dimension."""
    if m <= 4:
        return 4, 1
    elif m <= 8:
        return 8, 1
    elif m <= 16:
        return 16, 2
    elif m <= 32:
        return 32, 4
    else:
        return 128, 8


def _dynamic_mxfp4_quant_optimized(x: torch.Tensor, shuffle: bool = False):
    """Optimized MXFP4 quantization with dynamic BLOCK_SIZE."""
    M, N = x.shape
    BLOCK_SIZE, num_warps = _get_block_size(M)

    x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
    scaleN_valid = triton.cdiv(N, 32)
    scaleN = triton.cdiv(scaleN_valid, 8) * 8
    blockscale_e8m0 = torch.empty(
        (triton.cdiv(M, 256) * 256, scaleN),
        dtype=torch.uint8,
        device=x.device,
    )

    grid = (triton.cdiv(M, BLOCK_SIZE), scaleN)
    _optimized_mxfp4_quant_kernel[grid](
        x,
        x_fp4,
        blockscale_e8m0,
        *x.stride(),
        *x_fp4.stride(),
        *blockscale_e8m0.stride(),
        M=M,
        N=N,
        scaleN=scaleN_valid,
        BLOCK_SIZE=BLOCK_SIZE,
        SHUFFLE=shuffle,
        num_warps=num_warps,
    )

    if not shuffle:
        blockscale_e8m0 = blockscale_e8m0[:M, :scaleN_valid].contiguous()

    return x_fp4.view(dtypes.fp4x2), blockscale_e8m0.view(dtypes.fp8_e8m0)


# ============================================================
# Quantization Helpers
# ============================================================

def _quant_mxfp4_raw(x):
    """Optimized quant: no shuffle, dynamic BLOCK_SIZE. Returns raw uint8."""
    x_fp4, bs_e8m0 = _dynamic_mxfp4_quant_optimized(x, shuffle=False)
    return x_fp4.view(torch.uint8), bs_e8m0.view(torch.uint8)


def _quant_a_for_asm(A):
    """Quantize A for ASM path: shuffle + custom dtype view."""
    x_fp4, bs_e8m0 = _dynamic_mxfp4_quant_optimized(A, shuffle=True)
    return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)


# ============================================================
# Main Custom Kernel
# ============================================================

def custom_kernel(data: input_t) -> output_t:
    """Entry point called by eval.py"""
    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B_shuffle.shape[0]

    if m <= 32 and k <= 512:
        # Triton path: raw uint8, no shuffle, no custom dtype view
        B = B.contiguous()
        A_q, A_scale = _quant_mxfp4_raw(A)
        B_q_raw, B_scale = _quant_mxfp4_raw(B)
        return gemm_afp4wfp4(
            A_q,
            B_q_raw,
            A_scale,
            B_scale,
            dtype=torch.bfloat16,
        )
    else:
        # ASM path: REUSE pre-quantized B from input
        A_q, A_scale_sh = _quant_a_for_asm(A)
        out = torch.empty(m, n, dtype=torch.bfloat16, device=A.device)
        return aiter.gemm_a4w4(
            A_q,
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            out,
            bpreshuffle=True,
        )


def ref_kernel(data: input_t) -> output_t:
    """Reference: aiter.gemm_a4w4 with shuffle."""
    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    B = B.contiguous()
    m, k = A.shape
    n, _ = B.shape

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


check_implementation = make_match_reference(ref_kernel, rtol=1e-02, atol=1e-02)