submission 701042

"""
Fused quant+shuffle Triton kernel - eliminates shuffle kernel launch.
Writes blockscale directly in e8m0 shuffled layout within the quantization kernel.
"""
from task import input_t, output_t
import torch
import triton
import triton.language as tl

_states = {}
_gemm_asm = None
_fp4x2 = None
_fp8_e8m0 = None


@triton.jit
def _mxfp4_quant_op(x, BLOCK_SIZE_N, BLOCK_SIZE_M, MXFP4_QUANT_BLOCK_SIZE):
    """MXFP4 quantization - copied from aiter."""
    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 = x.reshape(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
    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_shuffled_ptr,
    stride_x_m_in, stride_x_n_in,
    stride_x_fp4_m_in, stride_x_fp4_n_in,
    M, N, sn,
    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, SCALING_MODE: tl.constexpr,
):
    pid_m = tl.program_id(0)
    start_n = tl.program_id(1) * NUM_ITER
    stride_x_m = tl.cast(stride_x_m_in, tl.int64)
    stride_x_n = tl.cast(stride_x_n_in, tl.int64)
    stride_x_fp4_m = tl.cast(stride_x_fp4_m_in, tl.int64)
    stride_x_fp4_n = tl.cast(stride_x_fp4_n_in, tl.int64)
    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)
        x_offs = x_offs_m[:, None] * stride_x_m + x_offs_n[None, :] * stride_x_n
        if EVEN_M_N:
            x = tl.load(x_ptr + x_offs, 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, mask=x_mask, cache_modifier=".cg").to(tl.float32)
        out_tensor, bs_e8m0 = _mxfp4_quant_op(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
        if EVEN_M_N:
            tl.store(x_fp4_ptr + out_offs, out_tensor)
        else:
            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_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        bs_n = pid_n * NUM_QUANT_BLOCKS + tl.arange(0, NUM_QUANT_BLOCKS)
        a = bs_m // 32
        b = (bs_m % 32) // 16
        c = bs_m % 16
        d = bs_n // 8
        e = (bs_n % 8) // 4
        f = bs_n % 4
        bs_shuffled_offs = (a[:, None] * 32 * sn +
                            d[None, :] * 256 + f[None, :] * 64 +
                            c[:, None] * 4 + e[None, :] * 2 + b[:, None])
        if EVEN_M_N:
            tl.store(bs_shuffled_ptr + bs_shuffled_offs, bs_e8m0)
        else:
            n_scale = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE
            bs_mask = (bs_m < M)[:, None] & (bs_n < n_scale)[None, :]
            tl.store(bs_shuffled_ptr + bs_shuffled_offs, bs_e8m0, mask=bs_mask)


def _init(device):
    global _gemm_asm, _fp4x2, _fp8_e8m0
    import aiter
    from aiter import dtypes
    _gemm_asm = torch.ops.aiter.gemm_a4w4_asm.default
    _fp4x2 = dtypes.fp4x2
    _fp8_e8m0 = dtypes.fp8_e8m0


def _init_shape(m, k, n, device):
    from aiter.ops.gemm_op_a4w4 import get_GEMM_config
    Q = 32
    n_scale = (k + Q - 1) // Q
    sm = (m + 255) // 256 * 256
    sn = (n_scale + 7) // 8 * 8
    m_pad = (m + 31) // 32 * 32

    x_fp4 = torch.empty((m, k // 2), dtype=torch.uint8, device=device)
    shuffle_out = torch.zeros(sm, sn, dtype=torch.uint8, device=device)
    gemm_out = torch.empty((m_pad, n), dtype=torch.bfloat16, device=device)

    if m <= 32:
        NI, BSM, BSN, NW, NS = 1, triton.next_power_of_2(m), 32, 1, 1
    else:
        NI, BSM, BSN, NW, NS = 4, 64, 64, 4, 2
        if k <= 16384: BSM, BSN = 32, 128
    if k <= 1024:
        NI, NS, NW = 1, 1, 4
        BSN = max(32, min(256, triton.next_power_of_2(k)))
        BSM = min(8, triton.next_power_of_2(m))

    qg = (triton.cdiv(m, BSM), triton.cdiv(k, BSN * NI))

    config = get_GEMM_config(m, n, k)
    kn, sk = "", 0
    if config is not None:
        sk = config.get("splitK", 0) or 0
        kn = config.get("kernelName", "")

    from aiter import dtypes as _dt
    x_fp4_view = x_fp4.view(_dt.fp4x2)
    shuffle_out_view = shuffle_out.view(_dt.fp8_e8m0)
    xf_s0, xf_s1 = x_fp4.stride()

    return (x_fp4, shuffle_out, gemm_out, qg,
            NI, BSM, BSN, NS, NW, kn, sk,
            sn, x_fp4_view, shuffle_out_view,
            xf_s0, xf_s1)


def custom_kernel(data: input_t) -> output_t:
    global _gemm_asm, _fp4x2, _fp8_e8m0

    A = data[0]
    if not A.is_contiguous():
        A = A.contiguous()
    B_sh = data[3]
    B_sc = data[4]
    m = A.shape[0]
    k = A.shape[1]
    n = B_sh.shape[0]

    if _gemm_asm is None:
        _init(A.device)

    key = (m, k, n)
    if key not in _states:
        _states[key] = _init_shape(m, k, n, A.device)
    s = _states[key]

    # Fused quant + shuffle (single kernel launch!)
    _fused_quant_shuffle_kernel[s[3]](
        A, s[0], s[1],
        k, 1, s[14], s[15],
        M=m, N=k, sn=s[11],
        BLOCK_SIZE_M=s[5], BLOCK_SIZE_N=s[6],
        NUM_ITER=s[4], NUM_STAGES=s[7],
        MXFP4_QUANT_BLOCK_SIZE=32, SCALING_MODE=0,
        num_warps=s[8], waves_per_eu=0, num_stages=1,
    )

    # GEMM
    _gemm_asm(s[12], B_sh, s[13], B_sc,
              s[2], s[9], None, 1.0, 0.0, True, s[10])
    return s[2][:m]