submission 746162

"""
FP4 quant + FP4 GEMM reference: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Quant logic follows aiter op_tests/test_gemm_a4w4.py (get_triton_quant(QuantType.per_1x32)).
"""
from task import input_t, output_t


import os
import tempfile
import torch
import weakref

try:
    import triton
    import triton.language as tl

    _AFU_TRITON_AVAILABLE = True
except Exception:
    triton = None
    tl = None
    _AFU_TRITON_AVAILABLE = False

# ---------------------------------------------------------------------------
# A4W4 tuned-config override
#
# Runner evidence: default tuned CSV is missing (N,K)=(2880,512) for MI355X,
# causing "not found tuned config ..." and fallback to a slower default path.
#
# Fix: provide a minimal tuned CSV (for the exact task test/bench shapes) that
# forces a known-good ASM kernel (_32x128) on cu_num=256 (MI355X).
# ---------------------------------------------------------------------------

_AFU_A4W4_CU_NUM = 256
_AFU_A4W4_DEFAULT_TUNED_CSV = "/home/runner/aiter/aiter/configs/a4w4_blockscale_tuned_gemm.csv"


def _afu_mangled_f4gemm_bf16_per1x32_bpreshuffle(tile_m: int, tile_n: int) -> str:
    fn = f"f4gemm_bf16_per1x32Fp4_BpreShuffle_{tile_m}x{tile_n}"
    return f"_ZN5aiter{len(fn)}{fn}E"


_AFU_A4W4_KERNELNAME_32x128 = _afu_mangled_f4gemm_bf16_per1x32_bpreshuffle(32, 128)
_AFU_A4W4_KERNELNAME_64x128 = _afu_mangled_f4gemm_bf16_per1x32_bpreshuffle(64, 128)


def _afu_pick_kernel_name(n: int, k: int) -> str:
    # Empirical: perNK (2112,7168)->64x128 tends to improve overall gmean.
    if (n, k) == (2112, 7168):
        return _AFU_A4W4_KERNELNAME_64x128
    return _AFU_A4W4_KERNELNAME_32x128

# Best-effort override: provide entries for all task shapes, but with a very large
# `us` so we *only* fill holes in the default tuned CSV (do not override).
# Shapes from `src/reference-kernels/problems/amd_202602/mxfp4-mm/task.yml`.
_AFU_A4W4_OVERRIDE_SHAPES = (
    # tests
    (8, 2112, 7168),
    (16, 3072, 1536),
    (64, 3072, 1536),
    (256, 2880, 512),
    # benchmarks
    (4, 2880, 512),
    (16, 2112, 7168),
    (32, 4096, 512),
    (32, 2880, 512),
    (64, 7168, 2048),
    (256, 3072, 1536),
)

_AFU_A4W4_OVERRIDE_READY = False

# Cached handles (init-once).
_AFU_GEMM_A4W4 = None
_AFU_GEMM_A4W4_ASM = None
_AFU_DYNAMIC_MXFP4_QUANT = None
_AFU_E8M0_SHUFFLE = None
_AFU_DTYPE_FP4X2 = None
_AFU_DTYPE_FP8_E8M0 = None
_AFU_DTYPE_BF16 = None

# Per-shape (M,N,K) cached best (kernelName, log2_k_split) for asm GEMM.
_AFU_A4W4_AUTOTUNE_CACHE = {}
_AFU_A4W4_AQ_CACHE = {}
# Per-shape cached fastest backend for the full gemm call.
# Values: ("aiter", None, 0) or ("asm", kernelName, log2_k_split)
_AFU_A4W4_FASTEST_CACHE = {}
# Per-shape cached output buffers for asm GEMM (reduces alloc overhead).
_AFU_A4W4_OUT_CACHE = {}
# Per-shape cached quantized activation buffers (fp4 + shuffled scales).
_AFU_QUANT_BUF_CACHE = {}

# f4gemm kernel tile options (gfx950). Source: /home/runner/aiter/hsa/gfx950/f4gemm/f4gemm_bf16_per1x32Fp4.csv
_AFU_F4GEMM_TILE_N_BY_M = {
    32: (128, 256, 384, 512, 640, 768, 896, 1024),
    64: (128, 256, 384, 512, 640, 768, 896, 1024),
    96: (128, 256, 384, 512, 640),
    128: (128, 256, 384, 512),
    160: (128, 256, 384),
    192: (128, 256),
    224: (128, 256),
    256: (128, 256),
}

# Only a subset of f4gemm kernels support split-K (see csv `splitK=1` rows).
_AFU_F4GEMM_SPLITK_CAPABLE = {
    (128, 512),
    (256, 256),
}


def _afu_getenv_bool(name: str, default: bool) -> bool:
    v = os.getenv(name)
    if v is None:
        return default
    v = v.strip().lower()
    return v in ("1", "true", "yes", "y", "on")


# Default-on: the elimination leaderboard measures steady-state per-shape runtime.
# We autotune once during the untimed correctness check and cache the best asm
# kernel per (M,N,K), so timed runs use the chosen MFMA/ASM kernel directly.
#
_AFU_A4W4_AUTOTUNE = _afu_getenv_bool("AFU_A4W4_AUTOTUNE", default=True)
_AFU_A4W4_AUTOTUNE_MAX_CANDIDATES = int(os.getenv("AFU_A4W4_AUTOTUNE_MAX", "12") or "12")
_AFU_A4W4_VALIDATE = _afu_getenv_bool("AFU_A4W4_VALIDATE", default=True)
_AFU_A4W4_VALIDATE_RTOL = float(os.getenv("AFU_A4W4_VALIDATE_RTOL", "1e-2") or "1e-2")
_AFU_A4W4_VALIDATE_ATOL = float(os.getenv("AFU_A4W4_VALIDATE_ATOL", "1e-2") or "1e-2")

_AFU_A4W4_FUSED_SCALE_SHUFFLE = _afu_getenv_bool("AFU_A4W4_FUSED_SCALE_SHUFFLE", default=True)
_AFU_A4W4_CACHE_AQ = _afu_getenv_bool("AFU_A4W4_CACHE_AQ", default=False)

# Activation quant backend:
# - "aiter": dynamic_mxfp4_quant + e8m0_shuffle (reference)
# - "fused": AFU Triton kernel that fuses quant + e8m0_shuffle-layout store
# - "auto":  time both once per (M,K) and cache the faster
_AFU_A4W4_QUANT_BACKEND = os.getenv("AFU_A4W4_QUANT_BACKEND", "auto").strip().lower()
_AFU_A4W4_QUANT_BACKEND_CACHE: dict[tuple[int, int, str], str] = {}


def _afu_clear_l2_cache_best_effort():
    # Match the evaluation harness behavior (clears caches between timed runs)
    # to avoid picking kernels that only win on warm-cache microbenchmarks.
    if not torch.cuda.is_available():
        return
    try:
        # ~256MB write: big enough to evict L2, small enough to be cheap.
        dummy = torch.empty((64 * 1024 * 1024,), device="cuda", dtype=torch.float32)
        dummy.uniform_()
        del dummy
    except Exception:
        pass


def _afu_time_cuda_us(fn, iters: int = 3) -> float | None:
    if not torch.cuda.is_available():
        return None
    best_us = None
    try:
        _ = fn()
        torch.cuda.synchronize()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        for _ in range(max(1, iters)):
            _afu_clear_l2_cache_best_effort()
            start.record()
            out = fn()
            end.record()
            torch.cuda.synchronize()
            us = float(start.elapsed_time(end)) * 1000.0
            best_us = us if best_us is None else min(best_us, us)
            del out
    except Exception:
        return None
    return best_us


if _AFU_TRITON_AVAILABLE:
    @triton.jit
    def _afu_mxfp4_quant_op(
        x,
        BLOCK_SIZE_N,
        BLOCK_SIZE_M,
        MXFP4_QUANT_BLOCK_SIZE,
    ):
        # Copy of ROCm/aiter _mxfp4_quant_op (kept local to avoid importing
        # internal aiter modules inside Triton).
        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)
        # Calculate scale
        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)

        # blockscale_e8m0
        bs_e8m0 = scale_e8m0_unbiased.to(tl.uint8) + 127

        quant_scale = tl.exp2(-scale_e8m0_unbiased)

        # Compute quantized x
        qx = x * quant_scale
        qx = qx.to(tl.uint32, bitcast=True)

        # Extract sign
        s = qx & 0x80000000
        # Set everything to positive, will add sign back at the end
        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)

        # Denormal numbers
        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 numbers
        normal_x = qx
        # resulting mantissa is odd
        mant_odd = (normal_x >> (MBITS_F32 - MBITS_FP4)) & 1
        # update exponent, rounding bias part 1
        val_to_add = ((EXP_BIAS_FP4 - EXP_BIAS_FP32) << MBITS_F32) + (1 << 21) - 1
        normal_x += val_to_add
        # rounding bias part 2
        normal_x += mant_odd
        # take the bits!
        normal_x = normal_x >> (MBITS_F32 - MBITS_FP4)
        normal_x = normal_x.to(tl.uint8)

        # Merge results
        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)
        # add sign back
        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.jit
    def _afu_dynamic_mxfp4_quant_kernel_shuffled(
        x_ptr,
        x_fp4_ptr,
        bs_shuf_ptr,
        stride_x_m_in,
        stride_x_n_in,
        stride_x_fp4_m_in,
        stride_x_fp4_n_in,
        stride_bs_m_in,
        stride_bs_n_in,
        M,
        N,
        BS_NCOLS,  # padded scale cols (= ceil(N/32) padded to 8)
        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)
        stride_bs_m = tl.cast(stride_bs_m_in, tl.int64)
        stride_bs_n = tl.cast(stride_bs_n_in, tl.int64)

        NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_N // MXFP4_QUANT_BLOCK_SIZE
        n_groups = BS_NCOLS // 8  # BS_NCOLS is padded to multiple-of-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 = _afu_mxfp4_quant_op(
                x, BLOCK_SIZE_N, BLOCK_SIZE_M, MXFP4_QUANT_BLOCK_SIZE
            )

            # Store FP4 tensor
            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 scale directly in the e8m0_shuffle layout.
            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)

            r = bs_offs_m[:, None]
            c = bs_offs_n[None, :]

            rg = r >> 5  # //32
            rin = r & 31
            r1 = rin >> 4  # //16
            r2 = rin & 15

            cg = c >> 3  # //8
            cin = c & 7
            c1 = cin >> 2  # //4
            c2 = cin & 3

            idx_in_rg = cg * 256 + c2 * 64 + r2 * 4 + c1 * 2 + r1
            row_in_rg = idx_in_rg // (n_groups * 8)
            col_out = idx_in_rg - row_in_rg * (n_groups * 8)
            row_out = rg * 32 + row_in_rg

            bs_out_offs = row_out * stride_bs_m + col_out * stride_bs_n

            bs_mask = (r < M) & (c < (N // MXFP4_QUANT_BLOCK_SIZE))
            tl.store(bs_shuf_ptr + bs_out_offs, bs_e8m0, mask=bs_mask)


def _afu_dynamic_mxfp4_quant_shuffled(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    # Best-effort: fall back to aiter ops on any issue.
    if not _AFU_TRITON_AVAILABLE or not x.is_cuda or not _AFU_A4W4_FUSED_SCALE_SHUFFLE:
        raise RuntimeError("AFU fused MXFP4 quant is disabled or Triton unavailable")

    if x.ndim != 2:
        raise ValueError("dynamic_mxfp4_quant expects 2D tensor")
    M, N = x.shape
    if (N // 2) % 2 != 0:
        raise ValueError("N must satisfy (N//2)%2==0 for fp4 packing")

    MXFP4_QUANT_BLOCK_SIZE = 32
    bs_cols = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE
    bs_cols_padded = ((bs_cols + 7) // 8) * 8
    bs_rows_padded = ((M + 255) // 256) * 256
    cache_key = (int(M), int(N), int(bs_rows_padded), int(bs_cols_padded), str(x.device))
    cached = _AFU_QUANT_BUF_CACHE.get(cache_key)
    if cached is None:
        x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
        bs_shuf = torch.empty((bs_rows_padded, bs_cols_padded), dtype=torch.uint8, device=x.device)
        if len(_AFU_QUANT_BUF_CACHE) > 32:
            _AFU_QUANT_BUF_CACHE.clear()
        _AFU_QUANT_BUF_CACHE[cache_key] = (x_fp4, bs_shuf)
    else:
        x_fp4, bs_shuf = cached

    NUM_ITER = 1
    BLOCK_SIZE_N = 256
    NUM_STAGES = 2
    # Small-M specializations help the (m,n,k)=(16,2112,7168) ranked case.
    if M <= 4:
        BLOCK_SIZE_M = 4
        NUM_WARPS = 1
    elif M <= 16:
        BLOCK_SIZE_M = 16
        NUM_WARPS = 2
    else:
        # Use a single larger tile for common sizes to keep compile time low.
        BLOCK_SIZE_M = 32
        NUM_WARPS = 4

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

    _afu_dynamic_mxfp4_quant_kernel_shuffled[grid](
        x,
        x_fp4,
        bs_shuf,
        *x.stride(),
        *x_fp4.stride(),
        *bs_shuf.stride(),
        M=M,
        N=N,
        BS_NCOLS=bs_cols_padded,
        MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE,
        NUM_ITER=NUM_ITER,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        NUM_STAGES=NUM_STAGES,
        EVEN_M_N=(M % BLOCK_SIZE_M == 0) and (N % (BLOCK_SIZE_N * NUM_ITER) == 0),
        num_warps=NUM_WARPS,
        waves_per_eu=0,
        num_stages=1,
    )

    return x_fp4, bs_shuf


def _afu_prepare_a4w4_override_and_cache():
    global _AFU_A4W4_OVERRIDE_READY
    global _AFU_GEMM_A4W4
    global _AFU_GEMM_A4W4_ASM
    global _AFU_DYNAMIC_MXFP4_QUANT, _AFU_E8M0_SHUFFLE
    global _AFU_DTYPE_FP4X2, _AFU_DTYPE_FP8_E8M0, _AFU_DTYPE_BF16

    if _AFU_GEMM_A4W4 is not None:
        return

    # 1) Prepare override CSV & env var BEFORE importing aiter.
    if not _AFU_A4W4_OVERRIDE_READY:
        _AFU_A4W4_OVERRIDE_READY = True
        try:
            d = tempfile.mkdtemp(prefix="afu_a4w4_")
            csv_path = os.path.join(d, "a4w4_blockscale_tuned_gemm_override.csv")
            with open(csv_path, "w", encoding="utf-8") as f:
                f.write("cu_num,M,N,K,kernelId,splitK,us,kernelName,tflops,bw,errRatio\n")
                for m, n, k in _AFU_A4W4_OVERRIDE_SHAPES:
                    kernel_name = _afu_pick_kernel_name(n, k)
                    f.write(
                        # us is only used to resolve duplicates during merge; keep it large
                        # so we don't override default tuned entries if they exist.
                        f"{_AFU_A4W4_CU_NUM},{m},{n},{k},0,0,1000000000,"
                        f"{kernel_name},0,0,0\n"
                    )
            base = os.environ.get("AITER_CONFIG_GEMM_A4W4") or _AFU_A4W4_DEFAULT_TUNED_CSV
            os.environ["AITER_CONFIG_GEMM_A4W4"] = base + os.pathsep + csv_path
        except Exception:
            # Best-effort: if override fails, fall back to default aiter behavior.
            pass

    # 2) Import and cache handles (hot path).
    import aiter
    from aiter import dtypes
    from aiter.ops.triton.quant import dynamic_mxfp4_quant
    from aiter.ops.gemm_op_a4w4 import gemm_a4w4_asm
    from aiter.utility.fp4_utils import e8m0_shuffle

    _AFU_GEMM_A4W4 = aiter.gemm_a4w4
    _AFU_GEMM_A4W4_ASM = gemm_a4w4_asm
    _AFU_DYNAMIC_MXFP4_QUANT = dynamic_mxfp4_quant
    _AFU_E8M0_SHUFFLE = e8m0_shuffle
    _AFU_DTYPE_FP4X2 = dtypes.fp4x2
    _AFU_DTYPE_FP8_E8M0 = dtypes.fp8_e8m0
    _AFU_DTYPE_BF16 = dtypes.bf16


def _afu_quant_a_per1x32_shuffled(a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Quantize activation A to MXFP4 (fp4x2 packed) and produce shuffled E8M0 scales.

    Returns:
      (a_fp4_u8, a_scale_shuf_u8)
    """
    a = a.contiguous()
    m, k = int(a.shape[0]), int(a.shape[1])
    cache_key = (m, k, str(a.device))
    forced = _AFU_A4W4_QUANT_BACKEND
    if forced in ("aiter", "ref", "baseline"):
        a_fp4_u8, bs_e8m0 = _AFU_DYNAMIC_MXFP4_QUANT(a)
        bs_shuf_u8 = _AFU_E8M0_SHUFFLE(bs_e8m0)
        return a_fp4_u8, bs_shuf_u8

    if forced in ("fused", "triton"):
        return _afu_dynamic_mxfp4_quant_shuffled(a)

    chosen = _AFU_A4W4_QUANT_BACKEND_CACHE.get(cache_key)
    if chosen == "fused":
        try:
            return _afu_dynamic_mxfp4_quant_shuffled(a)
        except Exception:
            chosen = None
    if chosen == "aiter":
        a_fp4_u8, bs_e8m0 = _AFU_DYNAMIC_MXFP4_QUANT(a)
        bs_shuf_u8 = _AFU_E8M0_SHUFFLE(bs_e8m0)
        return a_fp4_u8, bs_shuf_u8

    # Auto-pick: time both once per (M,K) and cache the faster.
    best = "aiter"
    aiter_us = None
    fused_us = None

    def _call_aiter():
        x_fp4, bs_e8m0 = _AFU_DYNAMIC_MXFP4_QUANT(a)
        bs_sh = _AFU_E8M0_SHUFFLE(bs_e8m0)
        return x_fp4, bs_sh

    aiter_us = _afu_time_cuda_us(_call_aiter, iters=2)

    def _call_fused():
        return _afu_dynamic_mxfp4_quant_shuffled(a)

    try:
        fused_us = _afu_time_cuda_us(_call_fused, iters=2)
    except Exception:
        fused_us = None

    if fused_us is not None and (aiter_us is None or fused_us < aiter_us * 0.99):
        best = "fused"

    if len(_AFU_A4W4_QUANT_BACKEND_CACHE) > 64:
        _AFU_A4W4_QUANT_BACKEND_CACHE.clear()
    _AFU_A4W4_QUANT_BACKEND_CACHE[cache_key] = best

    if best == "fused":
        try:
            return _afu_dynamic_mxfp4_quant_shuffled(a)
        except Exception:
            pass
    return _call_aiter()


def _afu_pick_candidate_kernels(m: int, n: int, k: int) -> list[tuple[str, int]]:
    # Returns a list of (kernelName, log2_k_split) candidates.
    padded_m = ((m + 31) // 32) * 32

    # Tile-M candidates: keep small set to avoid slow warmup.
    if padded_m <= 32:
        tile_ms = (32, 64)
    elif padded_m <= 64:
        tile_ms = (64, 32, 96)
    elif padded_m <= 96:
        tile_ms = (96, 64, 128)
    elif padded_m <= 128:
        tile_ms = (128, 96, 64)
    else:
        tile_ms = (256, 128, 96, 64)

    # Large-K cases are often split-K limited; include splitK-capable tiles even when M is small.
    if k >= 4096:
        if 128 not in tile_ms:
            tile_ms = tuple(tile_ms) + (128,)
        if 256 not in tile_ms:
            tile_ms = tuple(tile_ms) + (256,)

    cand: list[tuple[str, int, int]] = []
    for tm in tile_ms:
        tns = _AFU_F4GEMM_TILE_N_BY_M.get(tm)
        if not tns:
            continue

        divs = [tn for tn in tns if n % tn == 0]
        chosen_tns: list[int]
        if divs:
            chosen_tns = sorted(divs, reverse=True)[:2]
        else:
            le = [tn for tn in tns if tn <= n]
            chosen_tns = sorted(le, reverse=True)[:2] if le else [tns[0]]

        # Always include a smaller tile_n for occupancy if present.
        for extra in (256, 128):
            if extra in tns and extra not in chosen_tns:
                chosen_tns.append(extra)

        for tn in chosen_tns:
            kernel_name = _afu_mangled_f4gemm_bf16_per1x32_bpreshuffle(tm, tn)
            if (tm, tn) in _AFU_F4GEMM_SPLITK_CAPABLE:
                for split in (0, 1, 2, 3):
                    cand.append((kernel_name, tm, split))
            else:
                cand.append((kernel_name, tm, 0))

    # De-dup while keeping order.
    seen = set()
    out: list[tuple[str, int]] = []
    for kn, tm, sk in cand:
        key = (kn, sk)
        if key in seen:
            continue
        seen.add(key)
        out.append((kn, sk))
        if len(out) >= _AFU_A4W4_AUTOTUNE_MAX_CANDIDATES:
            break
    return out


def _afu_gemm_a4w4_asm(
    a_q: torch.Tensor,
    b_shuf: torch.Tensor,
    a_scale: torch.Tensor,
    b_scale: torch.Tensor,
    kernel_name: str,
    log2_k_split: int,
) -> torch.Tensor:
    # Allocate padded output like aiter.gemm_a4w4 does.
    m = a_q.numel() // a_q.shape[-1]
    n = b_shuf.shape[0]
    m_pad = (m + 31) // 32 * 32
    cache_key = (int(m_pad), int(n), str(a_q.device))
    out = _AFU_A4W4_OUT_CACHE.get(cache_key)
    if out is None:
        out = torch.empty((m_pad, n), dtype=torch.bfloat16, device=a_q.device)
        if len(_AFU_A4W4_OUT_CACHE) > 32:
            _AFU_A4W4_OUT_CACHE.clear()
        _AFU_A4W4_OUT_CACHE[cache_key] = out
    _AFU_GEMM_A4W4_ASM(
        a_q.view(m, -1),
        b_shuf,
        a_scale,
        b_scale,
        out,
        kernel_name,
        None,
        1.0,
        0.0,
        True,
        log2_k_split,
    )
    return out[:m]


def _afu_autotune_a4w4_asm(
    m: int,
    n: int,
    k: int,
    a_q: torch.Tensor,
    b_shuf: torch.Tensor,
    a_scale: torch.Tensor,
    b_scale: torch.Tensor,
) -> tuple[str | None, int]:
    # Best-effort: if timing or a kernel fails, skip it and continue.
    # Cache key uses logical M,N,K (not padded).
    key = (int(m), int(n), int(k))
    cached = _AFU_A4W4_AUTOTUNE_CACHE.get(key)
    if cached is not None:
        return cached

    if not _AFU_A4W4_AUTOTUNE or not a_q.is_cuda or _AFU_GEMM_A4W4_ASM is None:
        _AFU_A4W4_AUTOTUNE_CACHE[key] = (None, 0)
        return (None, 0)

    best = (None, 0)
    best_us = None
    candidates = _afu_pick_candidate_kernels(m, n, k)
    timings: list[tuple[float, str, int]] = []
    for kernel_name, split in candidates:
        try:
            # Warmup
            _ = _afu_gemm_a4w4_asm(a_q, b_shuf, a_scale, b_scale, kernel_name, split)
            torch.cuda.synchronize()

            start = torch.cuda.Event(enable_timing=True)
            end = torch.cuda.Event(enable_timing=True)
            us = None
            for _ in range(2):
                _afu_clear_l2_cache_best_effort()
                start.record()
                _ = _afu_gemm_a4w4_asm(a_q, b_shuf, a_scale, b_scale, kernel_name, split)
                end.record()
                torch.cuda.synchronize()
                v = float(start.elapsed_time(end)) * 1000.0
                us = v if us is None else min(us, v)
        except Exception:
            continue

        timings.append((us, kernel_name, split))
        if best_us is None or us < best_us:
            best_us = us
            best = (kernel_name, split)

    # Optional correctness validation: ensure the chosen asm config matches aiter.gemm_a4w4
    # within the competition tolerance (rtol/atol 1e-2). This avoids split-K/kernel corner
    # cases that can drift numerically on some shapes.
    if _AFU_A4W4_VALIDATE and timings:
        try:
            ref = _AFU_GEMM_A4W4(
                a_q,
                b_shuf,
                a_scale,
                b_scale,
                dtype=_AFU_DTYPE_BF16,
                bpreshuffle=True,
            )
            torch.cuda.synchronize()
            timings.sort(key=lambda t: t[0])
            validated = (None, 0)
            for _, kn, sk in timings:
                try:
                    out = _afu_gemm_a4w4_asm(a_q, b_shuf, a_scale, b_scale, kn, sk)
                    if torch.allclose(
                        out,
                        ref,
                        rtol=_AFU_A4W4_VALIDATE_RTOL,
                        atol=_AFU_A4W4_VALIDATE_ATOL,
                    ):
                        validated = (kn, sk)
                        break
                except Exception:
                    continue
            best = validated
        except Exception:
            # If validation fails for any reason, keep the best timing result.
            pass

    _AFU_A4W4_AUTOTUNE_CACHE[key] = best
    return best


def _afu_time_gemm_us(fn, iters: int = 2) -> float | None:
    if not torch.cuda.is_available():
        return None
    best_us = None
    try:
        _ = fn()
        torch.cuda.synchronize()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        for _ in range(max(1, iters)):
            _afu_clear_l2_cache_best_effort()
            start.record()
            out = fn()
            end.record()
            torch.cuda.synchronize()
            us = float(start.elapsed_time(end)) * 1000.0
            best_us = us if best_us is None else min(best_us, us)
            del out
    except Exception:
        return None
    return best_us


def _afu_pick_fastest_a4w4_backend(
    m: int,
    n: int,
    k: int,
    a_q: torch.Tensor,
    b_shuf: torch.Tensor,
    a_scale: torch.Tensor,
    b_scale: torch.Tensor,
) -> tuple[str, str | None, int]:
    """
    Decide per (M,N,K) whether to run:
      - aiter.gemm_a4w4 (may dispatch to non-asm/cktile paths), or
      - explicit gemm_a4w4_asm(kernelName, splitK)

    This prevents regressions when asm autotune picks a slower kernel than aiter's
    internal dispatcher for a given shape.
    """
    key = (int(m), int(n), int(k))
    cached = _AFU_A4W4_FASTEST_CACHE.get(key)
    if cached is not None:
        return cached

    # Default: aiter path.
    chosen: tuple[str, str | None, int] = ("aiter", None, 0)
    if (not _AFU_A4W4_AUTOTUNE) or (not a_q.is_cuda) or (_AFU_GEMM_A4W4_ASM is None):
        _AFU_A4W4_FASTEST_CACHE[key] = chosen
        return chosen

    asm_kernel, asm_split = _afu_autotune_a4w4_asm(m, n, k, a_q, b_shuf, a_scale, b_scale)
    if not asm_kernel:
        _AFU_A4W4_FASTEST_CACHE[key] = chosen
        return chosen

    def _call_aiter():
        return _AFU_GEMM_A4W4(
            a_q,
            b_shuf,
            a_scale,
            b_scale,
            dtype=_AFU_DTYPE_BF16,
            bpreshuffle=True,
        )

    def _call_asm():
        return _afu_gemm_a4w4_asm(a_q, b_shuf, a_scale, b_scale, asm_kernel, asm_split)

    aiter_us = _afu_time_gemm_us(_call_aiter, iters=2)
    asm_us = _afu_time_gemm_us(_call_asm, iters=2)

    # Prefer asm only if it is clearly faster (avoid noisy flips).
    if asm_us is not None and (aiter_us is None or asm_us < aiter_us * 0.99):
        chosen = ("asm", asm_kernel, int(asm_split))

    _AFU_A4W4_FASTEST_CACHE[key] = chosen
    return chosen


def custom_kernel(data: input_t) -> output_t:
    """
    Reference: MXFP4 per-1x32 quant on A; B_shuffle, B_scale_sh from generate_input.
    gemm_a4w4 with bpreshuffle=True.
    """
    _afu_prepare_a4w4_override_and_cache()

    A, _B, _B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()

    A_q = None
    A_scale_sh = None
    if _AFU_A4W4_CACHE_AQ:
        # Optional: speeds up non-ranked benchmarking (reuses same A).
        version = getattr(A, "_version", -1)
        key = id(A)
        cached = _AFU_A4W4_AQ_CACHE.get(key)
        if cached is not None:
            ref, cached_version, A_q, A_scale_sh = cached
            if ref() is not A or cached_version != version:
                A_q = None
                A_scale_sh = None
        if A_q is None:
            cached = None

    if A_q is None:
        x_fp4, bs_shuf = _afu_quant_a_per1x32_shuffled(A)

        A_q = x_fp4.view(_AFU_DTYPE_FP4X2)
        A_scale_sh = bs_shuf.view(_AFU_DTYPE_FP8_E8M0)

        if _AFU_A4W4_CACHE_AQ:
            version = getattr(A, "_version", -1)
            key = id(A)
            if len(_AFU_A4W4_AQ_CACHE) > 128:
                _AFU_A4W4_AQ_CACHE.clear()
            _AFU_A4W4_AQ_CACHE[key] = (weakref.ref(A), version, A_q, A_scale_sh)

    m = int(A.shape[0])
    n = int(B_shuffle.shape[0])
    k = int(A.shape[1])

    backend, kernel_name, split = _afu_pick_fastest_a4w4_backend(
        m, n, k, A_q, B_shuffle, A_scale_sh, B_scale_sh
    )
    if backend == "asm" and kernel_name:
        try:
            return _afu_gemm_a4w4_asm(A_q, B_shuffle, A_scale_sh, B_scale_sh, kernel_name, split)
        except Exception:
            pass

    return _AFU_GEMM_A4W4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=_AFU_DTYPE_BF16,
        bpreshuffle=True,
    )