submission 689231

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X
# submission_version: v018j

"""
v018j — v018h + selective exact fused tuning for the high-K small-M path.

Changes from v018h:
  - Only special-case `(8,2112,7168)` and `(16,2112,7168)` on the fused Triton path
  - Keep all M=4 / M=32 / M>=64 behavior identical to v018h

Changes retained from v018c:
  - M>=64: pre-compute quant kernel grid/args and flatten sh_scales in alloc
  - M>=64: inline _run_vendor_gemm to avoid function call
  - M<64: pre-compute B views, splitk config, grid in alloc (not re-computed per call)
  - Both: remove intermediate function calls in hot path

Architecture:
  M < 64  -> Fused Triton kernel: bf16 A -> inline _mxfp4_quant_op in registers
             -> tl.dot_scaled with pre-shuffled B -> bf16 C.  ONE kernel launch
             (+ optional reduce for split-K).  No A-quant global memory traffic.
  M >= 64 -> Fused quant+preshuffle Triton kernel (quant + e8m0_shuffle in ONE launch)
             -> ASM GEMM.  TWO launches total.

Key wins vs prior versions:
  - M<64: zero global-memory quant traffic (inline quant in registers)
  - M>=64: fused quant+preshuffle eliminates separate shuffle kernel/index_copy
"""

from __future__ import annotations
from typing import Any, Dict, Tuple

import torch
import triton
import triton.language as tl

from aiter import dtypes
from aiter.ops.gemm_op_a4w4 import gemm_a4w4_asm, get_GEMM_config

# ── Import inline quant op (for fused kernel M<64) ──
try:
    from aiter.ops.triton._triton_kernels.quant.quant import (
        _mxfp4_quant_op as _mxfp4_quant_op_imported,
    )
    _HAS_INLINE_QUANT = True
except Exception:
    _HAS_INLINE_QUANT = False

# ── Import low-level quant kernel (for fallback raw quant) ──
try:
    from aiter.ops.triton._triton_kernels.quant.quant import _dynamic_mxfp4_quant_kernel
    _HAS_LL_QUANT = True
except Exception:
    _HAS_LL_QUANT = False

if not _HAS_LL_QUANT:
    from aiter.ops.triton.quant import dynamic_mxfp4_quant

# ── Import low-level GEMM kernels (reduce kernel for split-K) ──
try:
    from aiter.ops.triton._triton_kernels.gemm.basic.gemm_afp4wfp4 import (
        _gemm_afp4wfp4_preshuffle_kernel,
        _gemm_afp4wfp4_reduce_kernel,
    )
    _HAS_LL_GEMM = True
except Exception:
    _HAS_LL_GEMM = False

from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, remap_xcd
from aiter.ops.triton.gemm.basic.gemm_afp4wfp4 import gemm_afp4wfp4_preshuffle

# For fallback preshuffle path (M>=64 when LL quant unavailable)
from aiter.utility.fp4_utils import e8m0_shuffle

input_t = Any
output_t = Any


# =========================================================================
# INLINE QUANT OP (from v014a — used by both fused GEMM and preshuffle kernels)
# =========================================================================

@triton.jit
def _mxfp4_quant_op(
    x,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
):
    """Same math as aiter quant.py, kept local for portability."""
    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)


# =========================================================================
# FUSED QUANT + GEMM TRITON KERNEL (M < 64) — from v015c
# =========================================================================

if _HAS_INLINE_QUANT:

    @triton.heuristics(
        {
            "EVEN_K": lambda args: (args["K"] % (args["BLOCK_SIZE_K"] // 2) == 0)
            and (args["SPLITK_BLOCK_SIZE"] % args["BLOCK_SIZE_K"] == 0)
            and (args["K"] % (args["SPLITK_BLOCK_SIZE"] // 2) == 0),
        }
    )
    @triton.jit
    def _fused_quant_gemm_kernel(
        a_ptr,           # bf16 [M, K_actual=2*K]
        b_ptr,           # uint8 pre-shuffled [N//16, K*16]
        c_ptr,           # output bf16 [M, N] or partials fp32 [KSPLIT, M, N]
        b_scales_ptr,    # uint8 pre-shuffled [N_scale//32, K_scale*32]
        M, N, K,         # K = K_packed = K_actual // 2
        stride_am, stride_ak,
        stride_bn, stride_bk,
        stride_ck, stride_cm, stride_cn,
        stride_bsn, stride_bsk,
        BLOCK_SIZE_M: tl.constexpr,
        BLOCK_SIZE_N: tl.constexpr,
        BLOCK_SIZE_K: tl.constexpr,
        GROUP_SIZE_M: tl.constexpr,
        NUM_KSPLIT: tl.constexpr,
        SPLITK_BLOCK_SIZE: 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,
    ):
        """Fused bf16->MXFP4 quantization + scaled GEMM."""
        tl.assume(stride_am > 0)
        tl.assume(stride_ak > 0)
        tl.assume(stride_bn > 0)
        tl.assume(stride_bk > 0)
        tl.assume(stride_cm > 0)
        tl.assume(stride_cn > 0)
        tl.assume(stride_bsn > 0)
        tl.assume(stride_bsk > 0)

        GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
        SCALE_GROUP_SIZE: tl.constexpr = 32

        pid_unified = tl.program_id(axis=0)
        pid_unified = remap_xcd(pid_unified, GRID_MN * NUM_KSPLIT, NUM_XCDS=8)
        pid_k = pid_unified % NUM_KSPLIT
        pid = pid_unified // NUM_KSPLIT
        num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
        num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

        if NUM_KSPLIT == 1:
            pid_m, pid_n = pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M)
        else:
            pid_m = pid // num_pid_n
            pid_n = pid % num_pid_n

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

        if (pid_k * SPLITK_BLOCK_SIZE // 2) < K:
            num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE // 2, BLOCK_SIZE_K // 2)

            offs_k_bf16 = tl.arange(0, BLOCK_SIZE_K)
            offs_k_bf16_start = pid_k * SPLITK_BLOCK_SIZE + offs_k_bf16
            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_start[None, :] * stride_ak
            )

            offs_k_sh = tl.arange(0, (BLOCK_SIZE_K // 2) * 16)
            offs_k_sh_start = pid_k * (SPLITK_BLOCK_SIZE // 2) * 16 + offs_k_sh
            offs_bn = (pid_n * (BLOCK_SIZE_N // 16) + tl.arange(0, BLOCK_SIZE_N // 16)) % N
            b_ptrs = b_ptr + (
                offs_bn[:, None] * stride_bn + offs_k_sh_start[None, :] * stride_bk
            )

            offs_bsn = (
                pid_n * (BLOCK_SIZE_N // 32) + tl.arange(0, BLOCK_SIZE_N // 32)
            ) % N
            offs_ks = (pid_k * (SPLITK_BLOCK_SIZE // SCALE_GROUP_SIZE) * 32) + tl.arange(
                0, (BLOCK_SIZE_K // SCALE_GROUP_SIZE) * 32
            )
            b_scale_ptrs = (
                b_scales_ptr
                + offs_bsn[:, None] * stride_bsn
                + offs_ks[None, :] * stride_bsk
            )

            accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

            for k in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter):

                if EVEN_K:
                    a_bf16 = tl.load(a_ptrs)
                else:
                    a_bf16 = tl.load(
                        a_ptrs,
                        mask=tl.arange(0, BLOCK_SIZE_K)[None, :] < 2 * K - k * BLOCK_SIZE_K,
                        other=0.0,
                    )

                a_fp32 = a_bf16.to(tl.float32)
                a_fp4, a_scales = _mxfp4_quant_op_imported(
                    a_fp32, BLOCK_SIZE_K, BLOCK_SIZE_M, SCALE_GROUP_SIZE
                )

                b_scales = (
                    tl.load(b_scale_ptrs, cache_modifier=cache_modifier)
                    .reshape(
                        BLOCK_SIZE_N // 32,
                        BLOCK_SIZE_K // SCALE_GROUP_SIZE // 8,
                        4, 16, 2, 2, 1,
                    )
                    .permute(0, 5, 3, 1, 4, 2, 6)
                    .reshape(BLOCK_SIZE_N, BLOCK_SIZE_K // SCALE_GROUP_SIZE)
                )

                if EVEN_K:
                    b = tl.load(b_ptrs, cache_modifier=cache_modifier)
                else:
                    b = tl.load(
                        b_ptrs,
                        mask=offs_k_sh[None, :] < (K - k * (BLOCK_SIZE_K // 2)) * 16,
                        other=0,
                        cache_modifier=cache_modifier,
                    )
                b = (
                    b.reshape(1, BLOCK_SIZE_N // 16, BLOCK_SIZE_K // 64, 2, 16, 16)
                    .permute(0, 1, 4, 2, 3, 5)
                    .reshape(BLOCK_SIZE_N, BLOCK_SIZE_K // 2)
                    .trans(1, 0)
                )

                accumulator = tl.dot_scaled(
                    a_fp4, a_scales, "e2m1", b, b_scales, "e2m1", accumulator
                )

                a_ptrs += BLOCK_SIZE_K * stride_ak
                b_ptrs += (BLOCK_SIZE_K // 2) * 16 * stride_bk
                b_scale_ptrs += BLOCK_SIZE_K * stride_bsk

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

            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, :]
                + pid_k * stride_ck
            )
            c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
            tl.store(c_ptrs, c, mask=c_mask, cache_modifier=".wt")


# =========================================================================
# FUSED QUANT + PRESHUFFLE KERNEL (M >= 64) — from v014a
# =========================================================================

@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 _dynamic_mxfp4_quant_preshuffle_kernel(
    x_ptr,
    x_fp4_ptr,
    bs_sh_ptr,
    stride_x_m_in,
    stride_x_n_in,
    stride_x_fp4_m_in,
    stride_x_fp4_n_in,
    scale_n_stride_in,
    M,
    N,
    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,
):
    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)
    scale_n_stride = tl.cast(scale_n_stride_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_rows = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        bs_cols = pid_n * NUM_QUANT_BLOCKS + tl.arange(0, NUM_QUANT_BLOCKS)

        m0 = bs_rows[:, None] // 32
        rem_m = bs_rows[:, None] % 32
        m1 = rem_m // 16
        m2 = rem_m % 16

        n0 = bs_cols[None, :] // 8
        rem_n = bs_cols[None, :] % 8
        n1 = rem_n // 4
        n2 = rem_n % 4

        sh_flat = (
            m1
            + 2 * n1
            + 4 * m2
            + 64 * n2
            + 256 * n0
            + 32 * scale_n_stride * m0
        )

        if EVEN_M_N:
            tl.store(bs_sh_ptr + sh_flat, bs_e8m0)
        else:
            bs_mask = (bs_rows < M)[:, None] & (bs_cols < (N // MXFP4_QUANT_BLOCK_SIZE))[None, :]
            tl.store(bs_sh_ptr + sh_flat, bs_e8m0, mask=bs_mask)


# =========================================================================
# HELPERS
# =========================================================================

_STATE: Dict[Tuple[int, int, int, int], Dict[str, Any]] = {}


def _cdiv(x: int, y: int) -> int:
    return (x + y - 1) // y


def _get_splitk(K, block_size_k, num_ksplit):
    splitk_block_size = _cdiv((2 * _cdiv(K, num_ksplit)), block_size_k) * block_size_k
    while num_ksplit > 1 and block_size_k > 16:
        if (
            K % (splitk_block_size // 2) == 0
            and splitk_block_size % block_size_k == 0
            and K % (block_size_k // 2) == 0
        ):
            break
        elif K % (splitk_block_size // 2) != 0 and num_ksplit > 1:
            num_ksplit //= 2
        elif splitk_block_size % block_size_k != 0:
            if num_ksplit > 1:
                num_ksplit //= 2
            elif block_size_k > 16:
                block_size_k //= 2
        elif K % (block_size_k // 2) != 0 and block_size_k > 16:
            block_size_k //= 2
        else:
            break
        splitk_block_size = _cdiv((2 * _cdiv(K, num_ksplit)), block_size_k) * block_size_k
    num_ksplit = _cdiv(K, (splitk_block_size // 2))
    return splitk_block_size, block_size_k, num_ksplit


def _pick_quant_cfg(m: int, k: int) -> tuple:
    """Returns (block_m, block_n, num_iter, num_warps)."""
    if k <= 1024:
        block_n = min(256, triton.next_power_of_2(k))
        block_n = max(32, block_n)
        block_m = min(8, triton.next_power_of_2(m))
        return block_m, block_n, 1, 4

    if m <= 32:
        return triton.next_power_of_2(m), 32, 1, 1

    # M >= 64: use smaller BLOCK_SIZE_M to increase program count and CU residency.
    block_m = 16
    block_n = 128 if k <= 16384 else 64
    return block_m, block_n, 1, 4


# =========================================================================
# PER-SHAPE CONFIGS (fused kernel, M < 64)
# =========================================================================

_EXACT_FUSED_CFG: Dict[Tuple[int, int, int], Dict[str, Any]] = {
    (8, 2112, 7168): {
        "BLOCK_SIZE_M": 8,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 256,
        "GROUP_SIZE_M": 1,
        "NUM_KSPLIT": 7,
        "num_warps": 4,
        "num_stages": 2,
        "waves_per_eu": 1,
        "matrix_instr_nonkdim": 16,
        "cache_modifier": ".cg",
    },
    (16, 2112, 7168): {
        "BLOCK_SIZE_M": 16,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 256,
        "GROUP_SIZE_M": 1,
        "NUM_KSPLIT": 7,
        "num_warps": 4,
        "num_stages": 2,
        "waves_per_eu": 1,
        "matrix_instr_nonkdim": 16,
        "cache_modifier": ".cg",
    },
}


def _fused_config(m, n, k):
    """Config for fused kernel (M < 64). BLOCK_SIZE_K must be >= 256 for B scale unshuffle."""
    exact = _EXACT_FUSED_CFG.get((m, n, k))
    if exact is not None:
        return dict(exact)

    cfg = {
        "BLOCK_SIZE_M": 8,
        "BLOCK_SIZE_N": 64,
        "BLOCK_SIZE_K": 256,
        "GROUP_SIZE_M": 1,
        "NUM_KSPLIT": 1,
        "num_warps": 2,
        "num_stages": 2,
        "waves_per_eu": 1,
        "matrix_instr_nonkdim": 16,
        "cache_modifier": None,
    }
    if m < 32 and k >= 4096:
        cfg["NUM_KSPLIT"] = 7
    elif m == 32:
        cfg["BLOCK_SIZE_M"] = 16
        cfg["BLOCK_SIZE_N"] = 128
        cfg["num_warps"] = 4
    return cfg


# Config for non-fused Triton fallback (when inline quant import fails)
_CFG_SPLITK_FALLBACK = {
    "BLOCK_SIZE_M": 8,
    "BLOCK_SIZE_N": 64,
    "BLOCK_SIZE_K": 256,
    "GROUP_SIZE_M": 1,
    "NUM_KSPLIT": 7,
    "num_warps": 2,
    "num_stages": 2,
    "waves_per_eu": 1,
    "matrix_instr_nonkdim": 16,
    "cache_modifier": None,
}

_CFG_M32_FALLBACK = {
    "BLOCK_SIZE_M": 16,
    "BLOCK_SIZE_N": 128,
    "BLOCK_SIZE_K": 256,
    "GROUP_SIZE_M": 1,
    "NUM_KSPLIT": 1,
    "num_warps": 4,
    "num_stages": 2,
    "waves_per_eu": 1,
    "matrix_instr_nonkdim": 16,
    "cache_modifier": None,
}


# =========================================================================
# BUFFER ALLOCATION
# =========================================================================

@torch.no_grad()
def _alloc_state(device, m, n, k):
    scale_n = k // 32
    st: Dict[str, Any] = {"quant_ptr": -1}

    if m >= 64:
        # v014a path: fused quant+preshuffle -> ASM GEMM
        padded_m32 = _cdiv(m, 32) * 32
        padded_m256 = _cdiv(m, 256) * 256

        st["aq_u8"] = torch.empty((m, k // 2), dtype=torch.uint8, device=device)
        st["aq_fp4"] = st["aq_u8"].view(dtypes.fp4x2)
        st["asc_sh_u8"] = torch.full((padded_m256, scale_n), 127, dtype=torch.uint8, device=device)
        st["asc_sh_e8m0"] = st["asc_sh_u8"].view(dtypes.fp8_e8m0)
        st["out"] = torch.empty((padded_m32, n), dtype=torch.bfloat16, device=device)
        st["out_slice"] = st["out"][:m]

        cfg = get_GEMM_config(m, n, k)
        kernel_name = cfg["kernelName"] if cfg is not None else ""
        split_k = cfg.get("splitK", 0) if cfg is not None else 0
        if split_k is None:
            split_k = 0

        st["gemm_cfg"] = cfg
        st["kernel_name"] = kernel_name
        st["split_k"] = split_k

        # Pre-compute quant kernel args (avoid recomputing per call)
        block_m, block_n, num_iter, num_warps = _pick_quant_cfg(m, k)
        st["quant_grid"] = (triton.cdiv(m, block_m), triton.cdiv(k, block_n * num_iter))
        st["quant_cfg"] = {
            "block_m": block_m, "block_n": block_n,
            "num_iter": num_iter, "num_warps": num_warps,
            "num_stages_triton": 1 if k <= 1024 else 2,
        }
        # Pre-flatten sh_scales
        st["asc_sh_flat"] = st["asc_sh_u8"].view(-1)
        st["asc_sh_scale_n"] = st["asc_sh_u8"].shape[1]

    elif _HAS_INLINE_QUANT:
        # v015c fused path: only need output buffer (no A quant buffers!)
        st["out"] = torch.empty((m, n), dtype=torch.bfloat16, device=device)
        cfg = _fused_config(m, n, k)

        # Pre-compute full fused config (avoid recomputing per call)
        k_packed = k // 2
        if cfg["NUM_KSPLIT"] > 1:
            sbs, bsk, nks = _get_splitk(k_packed, cfg["BLOCK_SIZE_K"], cfg["NUM_KSPLIT"])
            cfg["SPLITK_BLOCK_SIZE"] = sbs
            cfg["BLOCK_SIZE_K"] = bsk
            cfg["NUM_KSPLIT"] = nks
        else:
            cfg["SPLITK_BLOCK_SIZE"] = 2 * k_packed

        if cfg["BLOCK_SIZE_K"] >= 2 * k_packed:
            cfg["BLOCK_SIZE_K"] = triton.next_power_of_2(2 * k_packed)
            cfg["SPLITK_BLOCK_SIZE"] = 2 * k_packed
            cfg["NUM_KSPLIT"] = 1

        cfg["BLOCK_SIZE_N"] = max(cfg["BLOCK_SIZE_N"], 32)
        st["fused_cfg"] = cfg

        if cfg["NUM_KSPLIT"] > 1:
            actual_ksplit = triton.cdiv(k_packed, (cfg["SPLITK_BLOCK_SIZE"] // 2))
            st["y_pp"] = torch.empty((actual_ksplit, m, n), dtype=torch.float32, device=device)
            st["actual_ksplit"] = actual_ksplit
    else:
        # Fallback: separate quant + Triton GEMM
        st["aq"] = torch.empty((m, k // 2), dtype=torch.uint8, device=device)
        st["asc_raw"] = torch.empty((m, scale_n), dtype=torch.uint8, device=device)
        st["out"] = torch.empty((m, n), dtype=torch.bfloat16, device=device)
        if m < 32 and k >= 4096:
            k_packed = k // 2
            _, _, actual_ksplit = _get_splitk(
                k_packed, _CFG_SPLITK_FALLBACK["BLOCK_SIZE_K"], _CFG_SPLITK_FALLBACK["NUM_KSPLIT"]
            )
            st["y_pp"] = torch.empty((actual_ksplit, m, n), dtype=torch.float32, device=device)

    return st


@torch.no_grad()
def _get_state(device, m, n, k):
    key = ((device.index or 0), m, n, k)
    st = _STATE.get(key)
    if st is None:
        st = _alloc_state(device, m, n, k)
        _STATE[key] = st
    return st


# =========================================================================
# QUANT INTO BUFFERS (for fallback path only)
# =========================================================================

@torch.no_grad()
def _quant_into_raw(x, x_fp4, raw_scales):
    if not _HAS_LL_QUANT:
        q, s = dynamic_mxfp4_quant(x)
        x_fp4.copy_(q.view(torch.uint8) if q.dtype != torch.uint8 else q)
        raw_scales.copy_(s.view(torch.uint8) if s.dtype != torch.uint8 else s)
        return

    M, N = x.shape
    if M <= 32:
        num_iter, block_size_m, block_size_n = 1, triton.next_power_of_2(M), 32
        nw, ns = 1, 1
    else:
        num_iter, block_size_m, block_size_n = 4, 64, 64
        nw, ns = 4, 2
        if N <= 16384:
            block_size_m, block_size_n = 32, 128

    if N <= 1024:
        num_iter, ns, nw = 1, 1, 4
        block_size_n = max(32, min(256, triton.next_power_of_2(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))
    _dynamic_mxfp4_quant_kernel[grid](
        x, x_fp4, raw_scales,
        *x.stride(), *x_fp4.stride(), *raw_scales.stride(),
        M=M, N=N, MXFP4_QUANT_BLOCK_SIZE=32, SCALING_MODE=0, NUM_ITER=num_iter,
        BLOCK_SIZE_M=block_size_m, BLOCK_SIZE_N=block_size_n,
        NUM_STAGES=ns, num_warps=nw, waves_per_eu=0, num_stages=1,
    )


# =========================================================================
# QUANT INTO PRESHUFFLED (for M >= 64 path) — from v014a
# =========================================================================

@torch.no_grad()
def _quant_into_preshuffled(x, x_fp4, sh_scales):
    """Quantize A and write scales directly in e8m0_shuffle layout (one kernel)."""
    if not _HAS_LL_QUANT:
        q, s = dynamic_mxfp4_quant(x)
        x_fp4.copy_(q.view(torch.uint8) if q.dtype != torch.uint8 else q)
        s_u8 = s.view(torch.uint8) if s.dtype != torch.uint8 else s
        padded = torch.full_like(sh_scales, 127)
        padded[: s_u8.shape[0], : s_u8.shape[1]].copy_(s_u8)
        sh_scales.copy_(e8m0_shuffle(padded))
        return

    M, N = x.shape
    block_m, block_n, num_iter, num_warps = _pick_quant_cfg(M, N)
    grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n * num_iter))

    _dynamic_mxfp4_quant_preshuffle_kernel[grid](
        x,
        x_fp4,
        sh_scales.view(-1),
        *x.stride(),
        *x_fp4.stride(),
        sh_scales.shape[1],
        M=M,
        N=N,
        MXFP4_QUANT_BLOCK_SIZE=32,
        NUM_ITER=num_iter,
        BLOCK_SIZE_M=block_m,
        BLOCK_SIZE_N=block_n,
        NUM_STAGES=1 if N <= 1024 else 2,
        num_warps=num_warps,
        waves_per_eu=0,
        num_stages=1,
    )


# =========================================================================
# DISPATCH: FUSED PATH (M < 64) — from v015c
# =========================================================================

@torch.no_grad()
def _launch_fused(A, B_shuffle, B_scale_sh, st, m, n, k):
    k_packed = k // 2

    w = B_shuffle.view(torch.uint8).reshape(n // 16, k_packed * 16)
    w_scales = B_scale_sh.view(torch.uint8).reshape(
        B_scale_sh.shape[0] // 32, B_scale_sh.shape[1] * 32
    )

    # Config is pre-computed in _alloc_state
    cfg = st["fused_cfg"]

    y = st["out"]
    y_pp = st.get("y_pp")
    if cfg["NUM_KSPLIT"] > 1 and y_pp is not None:
        y_pp = y_pp[: cfg["NUM_KSPLIT"]]
    else:
        y_pp = None

    out_t = y if cfg["NUM_KSPLIT"] == 1 else y_pp

    grid = lambda META: (
        META["NUM_KSPLIT"]
        * triton.cdiv(m, META["BLOCK_SIZE_M"])
        * triton.cdiv(n, META["BLOCK_SIZE_N"]),
    )

    _fused_quant_gemm_kernel[grid](
        A, w, out_t, w_scales,
        m, n, k_packed,
        A.stride(0), A.stride(1),
        w.stride(0), w.stride(1),
        0 if cfg["NUM_KSPLIT"] == 1 else y_pp.stride(0),
        y.stride(0) if cfg["NUM_KSPLIT"] == 1 else y_pp.stride(1),
        y.stride(1) if cfg["NUM_KSPLIT"] == 1 else y_pp.stride(2),
        w_scales.stride(0), w_scales.stride(1),
        **cfg,
    )

    if cfg["NUM_KSPLIT"] > 1 and _HAS_LL_GEMM:
        _gemm_afp4wfp4_reduce_kernel[
            (triton.cdiv(m, 16), triton.cdiv(n, 64))
        ](
            y_pp, y, m, n,
            y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
            y.stride(0), y.stride(1),
            16, 64, st["actual_ksplit"],
            triton.next_power_of_2(cfg["NUM_KSPLIT"]),
        )

    return y


# =========================================================================
# DISPATCH: ASM PATH (M >= 64) — from v014a
# =========================================================================

@torch.no_grad()
def _run_vendor_gemm(st, b_sh, b_scale_sh):
    gemm_a4w4_asm(
        st["aq_fp4"],
        b_sh,
        st["asc_sh_e8m0"],
        b_scale_sh,
        st["out"],
        st["kernel_name"],
        None,
        1.0,
        0.0,
        True,
        log2_k_split=st["split_k"],
    )
    return st["out"]


# =========================================================================
# DISPATCH: FALLBACK PATH (inline quant import failed, M < 64)
# =========================================================================

@torch.no_grad()
def _launch_fallback(A, B_shuffle, B_scale_sh, st, m, n, k):
    a_ptr = A.data_ptr()
    if st["quant_ptr"] != a_ptr:
        _quant_into_raw(A, st["aq"], st["asc_raw"])
        st["quant_ptr"] = a_ptr

    w = B_shuffle.view(torch.uint8).reshape(n // 16, (k // 2) * 16)
    w_scales = B_scale_sh.view(torch.uint8).reshape(
        B_scale_sh.shape[0] // 32, B_scale_sh.shape[1] * 32
    )

    if m == 32 and _HAS_LL_GEMM:
        cfg = dict(_CFG_M32_FALLBACK)
        cfg["SPLITK_BLOCK_SIZE"] = 2 * (k // 2)
        cfg["BLOCK_SIZE_N"] = max(cfg["BLOCK_SIZE_N"], 32)

        grid = lambda META: (
            META["NUM_KSPLIT"] * triton.cdiv(m, META["BLOCK_SIZE_M"]) * triton.cdiv(n, META["BLOCK_SIZE_N"]),
        )
        _gemm_afp4wfp4_preshuffle_kernel[grid](
            st["aq"], w, st["out"], st["asc_raw"], w_scales,
            m, n, k // 2,
            st["aq"].stride(0), st["aq"].stride(1),
            w.stride(0), w.stride(1),
            0, st["out"].stride(0), st["out"].stride(1),
            st["asc_raw"].stride(0), st["asc_raw"].stride(1),
            w_scales.stride(0), w_scales.stride(1),
            **cfg,
        )
        return st["out"]

    if m < 32 and k >= 4096 and _HAS_LL_GEMM:
        k_packed = k // 2
        cfg = dict(_CFG_SPLITK_FALLBACK)
        sbs, bsk, nks = _get_splitk(k_packed, cfg["BLOCK_SIZE_K"], cfg["NUM_KSPLIT"])
        cfg["SPLITK_BLOCK_SIZE"] = sbs
        cfg["BLOCK_SIZE_K"] = bsk
        cfg["NUM_KSPLIT"] = nks
        cfg["BLOCK_SIZE_N"] = max(cfg["BLOCK_SIZE_N"], 32)

        y_pp = st["y_pp"][: cfg["NUM_KSPLIT"]]
        grid = lambda META: (
            META["NUM_KSPLIT"] * triton.cdiv(m, META["BLOCK_SIZE_M"]) * triton.cdiv(n, META["BLOCK_SIZE_N"]),
        )
        _gemm_afp4wfp4_preshuffle_kernel[grid](
            st["aq"], w, y_pp, st["asc_raw"], w_scales,
            m, n, k_packed,
            st["aq"].stride(0), st["aq"].stride(1),
            w.stride(0), w.stride(1),
            y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
            st["asc_raw"].stride(0), st["asc_raw"].stride(1),
            w_scales.stride(0), w_scales.stride(1),
            **cfg,
        )
        actual_ksplit = triton.cdiv(k_packed, (cfg["SPLITK_BLOCK_SIZE"] // 2))
        _gemm_afp4wfp4_reduce_kernel[
            (triton.cdiv(m, 16), triton.cdiv(n, 64))
        ](
            y_pp, st["out"], m, n,
            y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
            st["out"].stride(0), st["out"].stride(1),
            16, 64, actual_ksplit, triton.next_power_of_2(cfg["NUM_KSPLIT"]),
        )
        return st["out"]

    # Generic fallback
    return gemm_afp4wfp4_preshuffle(
        st["aq"], w, st["asc_raw"], w_scales,
        dtype=torch.bfloat16, y=st["out"],
    )


# =========================================================================
# UNIFIED DISPATCH
# =========================================================================

@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    A, _B, _B_q, B_shuffle, B_scale_sh = data
    m, k = A.shape
    n = B_shuffle.shape[0]

    st = _get_state(A.device, m, n, k)

    if m >= 64:
        # Inline quant+preshuffle + ASM GEMM (avoid function call overhead)
        qcfg = st["quant_cfg"]
        _dynamic_mxfp4_quant_preshuffle_kernel[st["quant_grid"]](
            A,
            st["aq_u8"],
            st["asc_sh_flat"],
            *A.stride(),
            *st["aq_u8"].stride(),
            st["asc_sh_scale_n"],
            M=m,
            N=k,
            MXFP4_QUANT_BLOCK_SIZE=32,
            NUM_ITER=qcfg["num_iter"],
            BLOCK_SIZE_M=qcfg["block_m"],
            BLOCK_SIZE_N=qcfg["block_n"],
            NUM_STAGES=qcfg["num_stages_triton"],
            num_warps=qcfg["num_warps"],
            waves_per_eu=0,
            num_stages=1,
        )
        gemm_a4w4_asm(
            st["aq_fp4"], B_shuffle, st["asc_sh_e8m0"], B_scale_sh,
            st["out"], st["kernel_name"],
            None, 1.0, 0.0, True,
            log2_k_split=st["split_k"],
        )
        return st["out_slice"]
    elif _HAS_INLINE_QUANT:
        # v015c path: fused quant+GEMM kernel (ONE kernel, zero quant memory traffic)
        return _launch_fused(A, B_shuffle, B_scale_sh, st, m, n, k)
    else:
        # Fallback: separate quant + Triton GEMM
        return _launch_fallback(A, B_shuffle, B_scale_sh, st, m, n, k)