submission 721677

"""
Hybrid: custom Triton GEMM (shuffled-scale-aware) for M≤32 k≤1024,
preshuffle Triton GEMM for M≤32 k>1024,
fused quant + ASM GEMM for M>32.
"""
from task import input_t, output_t
import torch
import triton
import triton.language as tl
import aiter
from aiter import dtypes
from aiter.ops.triton._triton_kernels.quant.quant import _mxfp4_quant_op

# Monkey-patch Triton type system for FP4 support
try:
    from triton.runtime.jit import type_canonicalisation_dict
    if 'float4_e2m1fn_x2' not in type_canonicalisation_dict:
        type_canonicalisation_dict['float4_e2m1fn_x2'] = 'u8'
    if 'float8_e8m0fnu' not in type_canonicalisation_dict:
        type_canonicalisation_dict['float8_e8m0fnu'] = 'u8'
except (ImportError, AttributeError):
    pass

from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
from aiter.ops.triton._triton_kernels.gemm.basic.gemm_afp4wfp4 import _gemm_afp4wfp4_reduce_kernel
from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid


# ─── Custom GEMM kernel: non-preshuffle B + shuffled B_scale on-the-fly ───
# Benefits: no per-K B data deshuffle (register reshape), no separate unshuffle kernel
# Only for KSPLIT=1 (k≤1024)

@triton.heuristics({
    "EVEN_K": lambda args: args["K"] % (args["BLOCK_SIZE_K"] // 2) == 0,
})
@triton.jit
def _custom_gemm_shufscale_kernel(
    a_ptr, b_ptr, c_ptr, b_scales_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    B_SCALE_SN: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    EVEN_K: tl.constexpr,
    num_warps: tl.constexpr,
    num_stages: tl.constexpr,
    waves_per_eu: tl.constexpr,
    matrix_instr_nonkdim: tl.constexpr,
    cache_modifier: tl.constexpr,
):
    """GEMM C = quant(A) @ B with B_scale loaded from shuffled layout on-the-fly.
    A: (M, K*2) bf16, B: (K, N) uint8 fp4x2 (transposed non-shuffled),
    B_scale: flat shuffled uint8 e8m0 buffer.
    K is in fp4 units (k_original // 2).
    BLOCK_SIZE_K is in bf16 units.
    """
    SCALE_GROUP_SIZE: tl.constexpr = 32
    N_SG: tl.constexpr = B_SCALE_SN // 8

    tl.assume(stride_am > 0)
    tl.assume(stride_ak > 0)
    tl.assume(stride_bk > 0)
    tl.assume(stride_bn > 0)
    tl.assume(stride_cm > 0)
    tl.assume(stride_cn > 0)

    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    pid_m, pid_n = pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M)

    tl.assume(pid_m >= 0)
    tl.assume(pid_n >= 0)

    # A pointers (bf16)
    offs_k_bf16 = tl.arange(0, BLOCK_SIZE_K)
    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k_bf16[None, :] * stride_ak

    # B pointers (non-shuffled, transposed: (K, N) in fp4x2 uint8)
    offs_k = tl.arange(0, BLOCK_SIZE_K // 2)
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn

    # B scale: compute shuffled flat indices for the (N, K//32) original scale space
    offs_ks = tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_SIZE)
    i = offs_bn[:, None]  # (BSN, 1) - N dimension
    j = offs_ks[None, :]  # (1, BSK//32) - scale group dimension
    a_s = i // 32
    b_s = (i % 32) // 16
    c_s = i % 16
    d_s = j // 8
    e_s = (j % 8) // 4
    f_s = j % 4
    b_scale_idx = ((((a_s * N_SG + d_s) * 4 + f_s) * 16 + c_s) * 2 + e_s) * 2 + b_s

    # Load B scale from shuffled buffer
    b_scales = tl.load(b_scales_ptr + b_scale_idx)

    # Load A (bf16) and B (fp4x2 packed uint8)
    if EVEN_K:
        a_bf16 = tl.load(a_ptrs)
        b = tl.load(b_ptrs, cache_modifier=cache_modifier)
    else:
        a_bf16 = tl.load(a_ptrs, mask=offs_k_bf16[None, :] < 2 * K, other=0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K, other=0, cache_modifier=cache_modifier)

    # Quantize A to MXFP4
    a, a_scales = _mxfp4_quant_op(a_bf16, BLOCK_SIZE_K, BLOCK_SIZE_M, 32)

    # Compute scaled dot product
    accumulator = tl.dot_scaled(a, a_scales, "e2m1", b, b_scales, "e2m1")

    c = accumulator.to(c_ptr.type.element_ty)

    # Store output
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)


# ─── Fused quant+shuffle kernel for M>32 path ───

@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_ptr,
    stride_x_m_in, stride_x_n_in, stride_x_fp4_m_in, stride_x_fp4_n_in,
    M, N,
    N_scale_pad: tl.constexpr, M_scale_pad: tl.constexpr,
    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
    N_scale_groups: tl.constexpr = N_scale_pad // 8

    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_local_m = tl.arange(0, BLOCK_SIZE_M)
        bs_local_n = tl.arange(0, NUM_QUANT_BLOCKS)
        bs_global_m = pid_m * BLOCK_SIZE_M + bs_local_m
        bs_global_n = pid_n * NUM_QUANT_BLOCKS + bs_local_n
        i = bs_global_m[:, None]
        j = bs_global_n[None, :]
        a = i // 32
        b = (i % 32) // 16
        c = i % 16
        d = j // 8
        e = (j % 8) // 4
        f = j % 4
        shuffled_offset = ((((a * N_scale_groups + d) * 4 + f) * 16 + c) * 2 + e) * 2 + b
        if EVEN_M_N:
            tl.store(bs_ptr + shuffled_offset, bs_e8m0)
        else:
            N_scale_valid = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE
            bs_mask = (bs_global_m < M)[:, None] & (bs_global_n < N_scale_valid)[None, :]
            tl.store(bs_ptr + shuffled_offset, bs_e8m0, mask=bs_mask)


_buf_cache = {}

def fused_mxfp4_quant_shuffle(x):
    M, N = x.shape
    MXFP4_QUANT_BLOCK_SIZE = 32
    N_scale_valid = triton.cdiv(N, MXFP4_QUANT_BLOCK_SIZE)
    M_scale_pad = triton.cdiv(M, 256) * 256
    N_scale_pad = triton.cdiv(N_scale_valid, 8) * 8

    key = (M, N)
    if key not in _buf_cache:
        fp4_buf = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
        scale_buf = torch.full((M_scale_pad * N_scale_pad,), 127, dtype=torch.uint8, device=x.device)
        _buf_cache[key] = (fp4_buf, scale_buf)
    x_fp4, blockscale_e8m0 = _buf_cache[key]

    if M <= 64:
        BLOCK_SIZE_M = 16
    else:
        BLOCK_SIZE_M = 32
    BLOCK_SIZE_N = 64
    NUM_ITER = 1
    NUM_WARPS = 4
    NUM_STAGES = 1

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

    _fused_quant_shuffle_kernel[grid](
        x, x_fp4, blockscale_e8m0,
        *x.stride(), *x_fp4.stride(),
        M=M, N=N, N_scale_pad=N_scale_pad, M_scale_pad=M_scale_pad,
        BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_N,
        NUM_ITER=NUM_ITER, NUM_STAGES=NUM_STAGES,
        MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE,
        SCALING_MODE=0, num_warps=NUM_WARPS, waves_per_eu=0, num_stages=1,
    )
    return (
        x_fp4.view(dtypes.fp4x2),
        blockscale_e8m0.view(M_scale_pad, N_scale_pad).view(dtypes.fp8_e8m0),
    )


# ─── ASM GEMM config injection ───
from aiter.ops.gemm_op_a4w4 import get_GEMM_config

def _inject_tuned_configs():
    get_GEMM_config(1, 512, 4096)
    if hasattr(get_GEMM_config, "gemm_dict"):
        d = get_GEMM_config.gemm_dict
        _k32x256 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x256E"
        _k32x128 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E"
        for M_val in [1, 2, 4, 8, 16, 32, 64, 128, 256]:
            for N in [2112, 2880, 3072, 4096, 7168]:
                for K in [512, 1536, 2048, 7168]:
                    key = (256, M_val, N, K)
                    if key not in d:
                        if K <= 1024:
                            d[key] = {"kernelName": _k32x256, "splitK": 0}
                        elif K >= 7168:
                            d[key] = {"kernelName": _k32x128, "splitK": 1}
                        else:
                            d[key] = {"kernelName": _k32x128, "splitK": 0}

_inject_tuned_configs()


# ─── Main kernel ───
_y_cache = {}

def custom_kernel(data: input_t) -> output_t:
    A, _, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B_shuffle.shape[0]

    if m <= 32:
        if k <= 1024:
            # Custom kernel: non-preshuffle B + shuffled scale on-the-fly
            b = B_q.view(torch.uint8).T  # (K//2, N) with strides (1, K//2)
            bs_flat = B_scale_sh.view(torch.uint8)  # (sm, sn) uint8
            bs_sn = bs_flat.shape[1]

            key = (m, n, k)
            if key not in _y_cache:
                _y_cache[key] = torch.empty((m, n), dtype=torch.bfloat16, device=A.device)
            y = _y_cache[key]

            K_fp4 = k // 2
            BSK = triton.next_power_of_2(2 * K_fp4)  # bf16 block size

            if m <= 4:
                BSM, BSN, cm = 4, 32, ".cg"
            elif m <= 16:
                BSM, BSN, cm = 16, 64, ".cg"
            else:
                BSM, BSN, cm = 8, 64, ".cg"

            grid = (triton.cdiv(m, BSM) * triton.cdiv(n, BSN),)
            _custom_gemm_shufscale_kernel[grid](
                A, b, y, bs_flat,
                m, n, K_fp4,
                A.stride(0), A.stride(1),
                b.stride(0), b.stride(1),
                y.stride(0), y.stride(1),
                B_SCALE_SN=bs_sn,
                BLOCK_SIZE_M=BSM, BLOCK_SIZE_N=BSN, BLOCK_SIZE_K=BSK,
                GROUP_SIZE_M=1,
                num_warps=4, num_stages=1, waves_per_eu=1,
                matrix_instr_nonkdim=16, cache_modifier=cm,
            )
            return y
        else:
            # k=7168: preshuffle Triton kernel (needs KSPLIT for parallelism)
            w = B_shuffle.view(torch.uint8).reshape(n // 16, (k // 2) * 16)
            ws = B_scale_sh.view(torch.uint8)
            m_sp, n_sp = ws.shape
            w_scales = ws.reshape(m_sp // 32, n_sp * 32)

            key = (m, n, k)
            if key not in _y_cache:
                _y_cache[key] = torch.empty((m, n), dtype=torch.bfloat16, device=A.device)
            y = _y_cache[key]

            cfg = {"BLOCK_SIZE_M": 8, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 512,
                   "GROUP_SIZE_M": 1, "NUM_KSPLIT": 7,
                   "num_warps": 4, "num_stages": 2, "waves_per_eu": 1,
                   "matrix_instr_nonkdim": 16, "cache_modifier": ".cg"}

            y_pp = gemm_a16wfp4_preshuffle(
                A, w, w_scales,
                prequant=True,
                dtype=torch.bfloat16,
                y=y,
                config=cfg,
                skip_reduce=True,
            )
            # Custom reduce (BSM=4 BSN=32 had best bench 8.69)
            ACTUAL_KSPLIT = y_pp.shape[0]
            REDUCE_BSM, REDUCE_BSN = 4, 32
            grid_r = (triton.cdiv(m, REDUCE_BSM), triton.cdiv(n, REDUCE_BSN))
            _gemm_afp4wfp4_reduce_kernel[grid_r](
                y_pp, y, m, n,
                y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
                y.stride(0), y.stride(1),
                REDUCE_BSM, REDUCE_BSN,
                ACTUAL_KSPLIT,
                triton.next_power_of_2(cfg["NUM_KSPLIT"]),
            )
            return y
    else:
        # Separate quant + ASM GEMM
        A_q, A_scale_sh = fused_mxfp4_quant_shuffle(A)
        return aiter.gemm_a4w4(
            A_q, B_shuffle, A_scale_sh, B_scale_sh,
            dtype=dtypes.bf16, bpreshuffle=True,
        )