submission 747449

import torch
import triton
import triton.language as tl
import aiter
from aiter import dtypes
from aiter.utility.fp4_utils import e8m0_shuffle

# ---------------------------------------------------------------------------
# ULTIMATE CACHE-PINNED QUANTIZER (BIT-EXACT AMD LOGIC)
# FIXED TRITON PACKING SYNTAX
# ---------------------------------------------------------------------------

@triton.jit
def _mxfp4_quant_op_inline(
    x,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
):
    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 = tl.reshape(x, (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

    # We use triton pointer arithmetic to load even and odd values properly
    # without relying on potentially unsupported slicing.

    m_idx = tl.arange(0, BLOCK_SIZE_M)
    q_idx = tl.arange(0, NUM_QUANT_BLOCKS)
    k_idx = tl.arange(0, MXFP4_QUANT_BLOCK_SIZE // 2)

    # Compute flat indices for evens (k*2) and odds (k*2 + 1)
    flat_evens_idx = (
        m_idx[:, None, None] * (NUM_QUANT_BLOCKS * MXFP4_QUANT_BLOCK_SIZE) +
        q_idx[None, :, None] * MXFP4_QUANT_BLOCK_SIZE +
        k_idx[None, None, :] * 2
    )
    flat_odds_idx = flat_evens_idx + 1

    # In Triton, we can't directly use array indexing to fetch from a tensor like a pointer.
    # Instead, we reshape the tensor so the 2-element pairs are their own dimension,
    # and then multiply by an indicator mask to extract the evens and odds safely!

    e2m1_pairs = tl.reshape(
        e2m1_value, (BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE // 2, 2)
    )

    evens_mask = tl.full((2,), 1, dtype=tl.int32)
    evens_mask = evens_mask * (tl.arange(0, 2) == 0) # [1, 0]

    odds_mask = tl.full((2,), 1, dtype=tl.int32)
    odds_mask = odds_mask * (tl.arange(0, 2) == 1) # [0, 1]

    # Sum along the last dimension (size 2) to collapse it
    evens = tl.sum(e2m1_pairs * evens_mask[None, None, None, :], axis=3)
    odds = tl.sum(e2m1_pairs * odds_mask[None, None, None, :], axis=3)

    x_fp4 = evens | (odds << 4)

    # Convert x_fp4 back to 2D tensor [BLOCK_SIZE_M, BLOCK_SIZE_N // 2]
    # Currently it is [BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE // 2]
    x_fp4 = tl.reshape(x_fp4, (BLOCK_SIZE_M, BLOCK_SIZE_N // 2))

    return x_fp4, bs_e8m0

@triton.jit
def fast_l2_mxfp4_quant_kernel(
    x_ptr,
    x_fp4_ptr,
    bs_ptr,
    M, N,
    stride_x_m, stride_x_n,
    stride_x_fp4_m, stride_x_fp4_n,
    stride_bs_m, stride_bs_n,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(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 + 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, other=0.0).to(tl.float32)

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

    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_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, qx_fp4, mask=out_mask, eviction_policy="evict_last")

    # Fix the scale storage shape
    NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_N // MXFP4_QUANT_BLOCK_SIZE
    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)
    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 // MXFP4_QUANT_BLOCK_SIZE))[None, :]

    # We need to reshape bs_e8m0 from [BLOCK_SIZE_M, NUM_QUANT_BLOCKS] so it stores properly
    bs_e8m0_reshaped = tl.reshape(bs_e8m0, [BLOCK_SIZE_M, NUM_QUANT_BLOCKS])

    tl.store(bs_ptr + bs_offs, bs_e8m0_reshaped, mask=bs_mask, eviction_policy="evict_last")

def custom_kernel(data):
    A, B, B_q, B_shuffle, B_scale_sh = data
    M, K = A.shape

    MXFP4_QUANT_BLOCK_SIZE = 32
    A_q = torch.empty((M, K // 2), dtype=torch.uint8, device=A.device)
    A_scale = torch.empty(
        ((K + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE, M),
        dtype=torch.uint8,
        device=A.device,
    ).T

    # Dynamic Block Tuning to maximize occupancy
    # We must ensure BLOCK_SIZE_M and BLOCK_SIZE_K are small enough to utilize all SMs
    BLOCK_SIZE_M = min(triton.next_power_of_2(M), 64)
    BLOCK_SIZE_K = min(triton.next_power_of_2(K), 128)
    if BLOCK_SIZE_K < 32:
        BLOCK_SIZE_K = 32
    if BLOCK_SIZE_M < 16:
        BLOCK_SIZE_M = 16

    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(K, BLOCK_SIZE_K))

    fast_l2_mxfp4_quant_kernel[grid](
        A, A_q, A_scale,
        M, K,
        A.stride(0), A.stride(1),
        A_q.stride(0), A_q.stride(1),
        A_scale.stride(0), A_scale.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_K,
        MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE,
        num_warps=4, num_stages=2
    )

    A_scale_sh_out = e8m0_shuffle(A_scale)

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

    return out