submission 712312

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X
"""
v10h: clean, rules-compliant version of v10d.

ARCHITECTURE:
  One custom Triton kernel `_gemm_k` with two modes:
    - FUSED (m≤32): load bf16 A-tile → in-register MXFP4 quant via
      aiter's _mxfp4_quant_op → tl.dot_scaled vs B_q. 1 kernel launch.
    - PREQUANT (m≥64): tiny quant kernel writes Afp4/Asc; GEMM reads
      fp4 A. Amortizes quant cost across N-tiles.
  + split-K (workspace + reduce) for thin grids (m≤16, k≥2048).

KEY TECHNIQUES:
  1. In-register quant fused into GEMM K-loop (no separate quant launch
     or intermediate HBM buffer for small-M).
  2. B_scale_sh read DIRECTLY from its e8m0_shuffle layout via the
     closed-form forward index → no unshuffle preprocessing.
  3. L2-cold autotune (matches eval.py's clear_l2_cache semantics).
  4. Per-(m,n,k) config cache; preallocated C/W/Afp4/Asc.

NOTE: eval.py measures GPU-event time AFTER a 16GB L2-flush, so Python
launch overhead (~15µs) is fully overlapped and never measured → plain
`kernel[grid](...)` is optimal; no low-level launch tricks needed.
"""
import os, sys
os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx950")
import warnings; warnings.filterwarnings("ignore")
import torch
import triton
import triton.language as tl

import aiter
from aiter import dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle
from aiter.ops.triton._triton_kernels.quant.quant import _mxfp4_quant_op

_L = lambda *a: print(*a, file=sys.stderr, flush=True)


# e8m0_shuffle forward flat idx (from aiter/utility/fp4_utils.py):
#   view(sm//32,2,16, sn//8,2,4).permute(0,3,5,2,4,1).reshape(sm,sn)
# Separable: flat = row_part(r) + col_part(c)
@triton.jit
def _sh_row(r, sn8):
    return (r // 32) * (sn8 * 256) + (r % 16) * 4 + (r // 16) % 2


@triton.jit
def _sh_col(c):
    return (c // 8) * 256 + (c % 4) * 64 + (c // 4) % 2 * 2


@triton.jit
def _gemm_k(
    A, Asc, Bq, Bsc, C,
    M, N, K, sA_m, sAsc_m, sBq_n, sC_k, sC_m, sn8,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
    SPLIT_K: tl.constexpr, EVEN_N: tl.constexpr, PREQUANT: tl.constexpr,
):
    pid = tl.program_id(0)
    num_n = tl.cdiv(N, BN)
    num_mn = tl.cdiv(M, BM) * num_n
    pid_k = pid // num_mn
    pid_mn = pid % num_mn
    pid_m = pid_mn // num_n
    pid_n = pid_mn % num_n

    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_n64 = offs_n.to(tl.int64)
    mask_m = offs_m < M
    mask_n = offs_n < N
    rk = tl.arange(0, BK)
    rk2 = tl.arange(0, BK // 2)
    rk32 = tl.arange(0, BK // 32)

    k_per = tl.cdiv(tl.cdiv(K, BK), SPLIT_K) * BK
    k_lo = pid_k * k_per
    k_hi = min(k_lo + k_per, K)

    bq_ptrs = Bq + offs_n64[:, None] * sBq_n + (k_lo // 2 + rk2)[None, :]
    bsc_row = _sh_row(offs_n64, sn8)
    acc = tl.zeros((BM, BN), dtype=tl.float32)

    if PREQUANT:
        a_ptrs = A + offs_m[:, None].to(tl.int64) * sA_m \
                 + (k_lo // 2 + rk2)[None, :]
        asc_ptrs = Asc + offs_m[:, None].to(tl.int64) * sAsc_m \
                   + (k_lo // 32 + rk32)[None, :]
    else:
        a_ptrs = A + offs_m[:, None].to(tl.int64) * sA_m \
                 + (k_lo + rk)[None, :]

    for k in tl.range(k_lo, k_hi, BK):
        if PREQUANT:
            a_fp4 = tl.load(a_ptrs, mask=mask_m[:, None], other=0)
            a_sc = tl.load(asc_ptrs, mask=mask_m[:, None], other=0)
            a_ptrs += BK // 2; asc_ptrs += BK // 32
        else:
            a_bf = tl.load(a_ptrs, mask=mask_m[:, None], other=0.0)
            a_fp4, a_sc = _mxfp4_quant_op(a_bf.to(tl.float32), BK, BM, 32)
            a_ptrs += BK

        if EVEN_N:
            b_fp4_t = tl.load(bq_ptrs)
            b_sc = tl.load(Bsc + bsc_row[:, None]
                           + _sh_col(k // 32 + rk32)[None, :])
        else:
            b_fp4_t = tl.load(bq_ptrs, mask=mask_n[:, None], other=0)
            b_sc = tl.load(Bsc + bsc_row[:, None]
                           + _sh_col(k // 32 + rk32)[None, :],
                           mask=mask_n[:, None], other=0)
        acc = tl.dot_scaled(a_fp4, a_sc, "e2m1",
                            tl.trans(b_fp4_t), b_sc, "e2m1", acc)
        bq_ptrs += BK // 2

    c_off = (pid_k * sC_k + offs_m[:, None].to(tl.int64) * sC_m
             + offs_n[None, :])
    cmask = mask_m[:, None] & mask_n[None, :]
    if SPLIT_K == 1:
        tl.store(C + c_off, acc.to(tl.bfloat16), mask=cmask)
    else:
        tl.store(C + c_off, acc, mask=cmask)


@triton.jit
def _reduce_k(W, C, SK, M, N, sW_k, sW_m, sC_m,
              BLK: tl.constexpr, SKC: tl.constexpr):
    pid = tl.program_id(0)
    off = pid * BLK + tl.arange(0, BLK)
    om = off // N; on = off % N; mask = om < M
    base = om.to(tl.int64) * sW_m + on
    s = tl.zeros((BLK,), dtype=tl.float32)
    for i in tl.static_range(SKC):
        s += tl.load(W + i * sW_k + base, mask=mask & (i < SK), other=0.0)
    tl.store(C + om.to(tl.int64) * sC_m + on, s.to(tl.bfloat16), mask=mask)


@triton.jit
def _quant_a_k(A, Afp4, Asc, M, K, sA_m, sAf_m, sAs_m,
               BM: tl.constexpr, BK: tl.constexpr):
    pid = tl.program_id(0)
    nk = tl.cdiv(K, BK)
    pm = pid // nk; pk = pid % nk
    offs_m = pm * BM + tl.arange(0, BM)
    offs_k = pk * BK + tl.arange(0, BK)
    mask_m = offs_m < M
    a = tl.load(A + offs_m[:, None].to(tl.int64) * sA_m + offs_k[None, :],
                mask=mask_m[:, None], other=0.0).to(tl.float32)
    af, asc = _mxfp4_quant_op(a, BK, BM, 32)
    tl.store(Afp4 + offs_m[:, None].to(tl.int64) * sAf_m
             + (pk * (BK // 2) + tl.arange(0, BK // 2))[None, :],
             af, mask=mask_m[:, None])
    tl.store(Asc + offs_m[:, None].to(tl.int64) * sAs_m
             + (pk * (BK // 32) + tl.arange(0, BK // 32))[None, :],
             asc, mask=mask_m[:, None])


def _cfgs(m, n, k):
    """(BM, BN, BK, SPLIT_K, num_warps, nonK, num_stages, PREQUANT)"""
    out = []
    PQs = (False,) if m <= 32 else (True,) if m >= 128 else (False, True)
    for PQ in PQs:
        BMs = (16,) if not PQ else tuple(b for b in (16, 32, 64, 128) if b <= m)
        for BM in BMs:
            for BN in (32, 64, 128, 256):
                if BN > n: continue
                bt = -(-m // BM) * -(-n // BN)
                for BK in (256, 512):
                    if BK > k: continue
                    nkit = k // BK
                    SKs = [1]
                    if bt < 200 and nkit >= 2:
                        tgt = max(1, 256 // bt)
                        for s in (2, 4, 8):
                            if s <= nkit and s <= tgt * 2: SKs.append(s)
                    for SK in SKs:
                        for nw in (4, 8):
                            for nK in ((16,) if BM == 16 else (16, 32)):
                                out.append((BM, BN, BK, SK, nw, nK, 2, PQ))
    seen, r = set(), []
    for c in out:
        if c not in seen: seen.add(c); r.append(c)
    return r


_L2FLUSH = torch.empty(512 * 1024 * 1024, dtype=torch.int8, device="cuda")


def _gpu_time_cold(fn, n_iter=6):
    for _ in range(2): fn()
    torch.cuda.synchronize()
    evs = [(torch.cuda.Event(True), torch.cuda.Event(True)) for _ in range(n_iter)]
    for e0, e1 in evs:
        _L2FLUSH.zero_()
        e0.record(); fn(); e1.record()
    torch.cuda.synchronize()
    ts = sorted(e0.elapsed_time(e1) for e0, e1 in evs)
    return sum(ts[:n_iter - 1]) * 1000.0 / (n_iter - 1)


def _ref(A, B_shuffle, B_scale_sh):
    Aq, As = dynamic_mxfp4_quant(A)
    return aiter.gemm_a4w4(Aq.view(dtypes.fp4x2), B_shuffle,
                           e8m0_shuffle(As).view(dtypes.fp8_e8m0),
                           B_scale_sh, dtype=dtypes.bf16, bpreshuffle=True)


_STATE: dict = {}


def _build(data):
    A, B, B_q, B_shuffle, B_scale_sh = data
    m, k = A.shape; n = B.shape[0]
    sn = B_scale_sh.shape[1]; sn8 = sn // 8
    dev = A.device

    Bq = B_q.contiguous().view(torch.uint8)
    Bsc = B_scale_sh.contiguous().view(torch.uint8)
    C_bf = torch.empty((m, n), dtype=torch.bfloat16, device=dev)
    W = torch.zeros((8, m, n), dtype=torch.float32, device=dev)
    Afp4 = torch.empty((m, k // 2), dtype=torch.uint8, device=dev)
    Asc = torch.empty((m, k // 32), dtype=torch.uint8, device=dev)

    ref_f = _ref(A, B_shuffle, B_scale_sh).float()
    mag = ref_f.abs().mean().item() + 1e-9

    QBM, QBK = 16, min(256, k)
    q_gx = triton.cdiv(m, QBM) * (k // QBK)
    r_gx = triton.cdiv(m * n, 256)

    def _do_quant(A_in):
        _quant_a_k[(q_gx,)](A_in, Afp4, Asc, m, k, k, k // 2, k // 32,
                            BM=QBM, BK=QBK, num_warps=4)

    def _do_reduce(SK):
        _reduce_k[(r_gx,)](W, C_bf, SK, m, n, m * n, n, n,
                           BLK=256, SKC=8, num_warps=4)

    _do_quant(A); _do_reduce(1); torch.cuda.synchronize()

    cfgs = _cfgs(m, n, k)
    _L(f"\n[v10h m={m} n={n} k={k}] {len(cfgs)} cfgs")

    best, best_t, best_go = None, float("inf"), None
    for cfg in cfgs:
        BM, BN, BK, SK, nw, nK, ns, PQ = cfg
        C_out = W if SK > 1 else C_bf
        sC_k = W.stride(0) if SK > 1 else 0
        sC_m = W.stride(1) if SK > 1 else n
        sAm = (k // 2) if PQ else k
        even_n = (n % BN == 0)
        gx = triton.cdiv(m, BM) * triton.cdiv(n, BN) * SK

        def _go(_A=A, _Bq=Bq, _Bsc=Bsc, _cfg=cfg, _gx=gx, _C=C_out,
                _sAm=sAm, _sCk=sC_k, _sCm=sC_m, _even=even_n):
            BM, BN, BK, SK, nw, nK, ns, PQ = _cfg
            if PQ:
                _do_quant(_A)
                a_src = Afp4
            else:
                a_src = _A
            _gemm_k[(_gx,)](
                a_src, Asc, _Bq, _Bsc, _C, m, n, k,
                _sAm, k // 32, k // 2, _sCk, _sCm, sn8,
                BM=BM, BN=BN, BK=BK, SPLIT_K=SK, EVEN_N=_even,
                PREQUANT=PQ, num_warps=nw, num_stages=ns,
                matrix_instr_nonkdim=nK, waves_per_eu=0)
            if SK > 1:
                _do_reduce(SK)
            return C_bf

        try:
            out = _go()
            torch.cuda.synchronize()
            err = ((out.float() - ref_f).abs().mean() / mag).item()
            if err > 5e-3:
                if best is None: _L(f"  {cfg}: ERR {err:.2%}")
                continue
            t = _gpu_time_cold(_go)
            if t < best_t:
                best_t, best, best_go = t, cfg, _go
                _L(f"  {cfg}: {t:.2f}us grid={gx} *")
        except Exception as e:
            if best is None:
                _L(f"  {cfg}: {type(e).__name__}: {str(e)[:100]}")

    if best is None:
        _L(f"  → fallback"); return {"hot": None}
    _L(f"  → best={best} @ {best_t:.2f}us")
    return {"hot": best_go}


def custom_kernel(data):
    A = data[0]; Bq = data[2]
    key = (A.shape[0], Bq.shape[0], A.shape[1])
    S = _STATE.get(key)
    if S is None:
        S = _build(data); _STATE[key] = S
    hot = S["hot"]
    if hot is None:
        return _ref(A, data[3], data[4])
    return hot(A, Bq.view(torch.uint8), data[4].view(torch.uint8))