submission 700777

"""
MXFP4 GEMM — v2 optimized: fused quant+shuffle Triton kernel + CK gemm_a4w4.
Fuses the original dynamic_mxfp4_quant and e8m0_shuffle into a single Triton
kernel launch, reducing 3 kernel launches to 2 while preserving exact numerical
behavior of the original quantization.
"""
import torch
import triton
import triton.language as tl
try:
    from task import input_t, output_t
except ImportError:
    from typing import Any, Tuple
    input_t = Tuple[Any, ...]
    output_t = Any
from aiter import dtypes
import aiter
from aiter.ops.triton.quant import _mxfp4_quant_op


@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_mxfp4_quant_shuffle_kernel(
    x_ptr,
    x_fp4_ptr,
    bs_ptr,
    stride_x_m_in,
    stride_x_n_in,
    stride_x_fp4_m_in,
    stride_x_fp4_n_in,
    M,
    N,
    scaleN: tl.int64,
    scaleN_pad: tl.int64,
    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
        )

        # Store FP4 output (same as original)
        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)

        # Store blockscale with CK shuffle layout fused in
        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)

        # CK shuffle: (m, n) -> d0*32*sn + d1*256 + d2*64 + d3*4 + d4*2 + d5
        m_idx = bs_offs_m[:, None]
        n_idx = bs_offs_n[None, :]
        d0 = m_idx // 32
        d5 = (m_idx % 32) // 16
        d3 = m_idx % 16
        d1 = n_idx // 8
        d4 = (n_idx % 8) // 4
        d2 = n_idx % 4
        shuffle_offs = (
            d0 * 32 * scaleN_pad + d1 * 256 + d2 * 64 + d3 * 4 + d4 * 2 + d5
        )

        if EVEN_M_N:
            tl.store(bs_ptr + shuffle_offs, bs_e8m0)
        else:
            bs_mask = (bs_offs_m < M)[:, None] & (
                bs_offs_n < scaleN
            )[None, :]
            tl.store(
                bs_ptr + shuffle_offs,
                bs_e8m0,
                mask=bs_mask,
            )


def _quant_mxfp4_fused(x):
    M, N = x.shape
    assert (N // 2) % 2 == 0
    MXFP4_QUANT_BLOCK_SIZE = 32

    x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)

    scaleN = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE
    scaleN_pad = (scaleN + 7) // 8 * 8
    sm = (M + 255) // 256 * 256
    bs_e8m0 = torch.empty(sm * scaleN_pad, dtype=torch.uint8, device=x.device)

    if M <= 32:
        NUM_ITER = 1
        BLOCK_SIZE_M = triton.next_power_of_2(M)
        BLOCK_SIZE_N = 32
        NUM_WARPS = 1
        NUM_STAGES = 1
    else:
        NUM_ITER = 4
        BLOCK_SIZE_M = 64
        BLOCK_SIZE_N = 64
        NUM_WARPS = 4
        NUM_STAGES = 2
        if N <= 16384:
            BLOCK_SIZE_M = 32
            BLOCK_SIZE_N = 128

    if N <= 1024:
        NUM_ITER = 1
        NUM_STAGES = 1
        NUM_WARPS = 4
        BLOCK_SIZE_N = min(256, triton.next_power_of_2(N))
        BLOCK_SIZE_N = max(32, BLOCK_SIZE_N)
        BLOCK_SIZE_M = min(8, triton.next_power_of_2(M))

    grid = (
        triton.cdiv(M, BLOCK_SIZE_M),
        triton.cdiv(N, BLOCK_SIZE_N * NUM_ITER),
    )

    _fused_mxfp4_quant_shuffle_kernel[grid](
        x,
        x_fp4,
        bs_e8m0,
        *x.stride(),
        *x_fp4.stride(),
        M=M,
        N=N,
        scaleN=scaleN,
        scaleN_pad=scaleN_pad,
        MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE,
        SCALING_MODE=0,
        NUM_ITER=NUM_ITER,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        NUM_STAGES=NUM_STAGES,
        num_warps=NUM_WARPS,
        waves_per_eu=0,
        num_stages=1,
    )

    return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(sm, scaleN_pad).view(dtypes.fp8_e8m0)


def custom_kernel(data: input_t) -> output_t:
    A, B, B_q, B_shuffle, B_scale_sh = data
    A_q, A_scale_sh = _quant_mxfp4_fused(A)
    return aiter.gemm_a4w4(
        A_q, B_shuffle, A_scale_sh, B_scale_sh,
        dtype=dtypes.bf16, bpreshuffle=True,
    )