"""Triton kernels for ``luau_infer`` (fused decoding of the LuauCoder HF port).

Every kernel works on *rows*: one row per token of the call (``B`` rows when
decoding, ``B * T`` for a prefill or a speculative verify), each with its own
document position ``pos`` (RoPE, conv / value-shift masks), cache ``slot``
(where its K/V go) and ``start`` (the first valid key slot of its sequence:
left padding). Nothing reads a Python scalar that changes between steps, so a
decode step replays as a CUDA graph.

* ``mix_norm``     Block AttnRes depth mix (a softmax over the sources, one tile)
                   + RMSNorm, or the embedding gather + RMSNorm (layer 0).
* ``gemv``         skinny ``x @ W.T`` (bf16 or rowwise-FP8 weights), epilogues:
                   store, residual add, SwiGLU (gate|up rows), softcap logits.
* ``qk_decode``    CCGQA conv + QK-mean + QK-norm + gains + RoPE + value shift,
                   writes K/V to the cache and the conv latents to their ring.
* ``attention``    GQA flash-decoding: each program reads a span of one KV
                   head's keys once for all its query heads and all the
                   sequence's rows (a speculative verify); the last span
                   program merges and finishes DiffAttn V2 + XSA (+ gate).
* ``qk_prefill``   ``qk_decode`` for prompts, in token tiles.
* ``iha_mix``      IHA: pseudo-head mixes of the latent q / k / v, norm, temperature,
                   RoPE at virtual positions, interleaved pseudo-K/V cache write;
                   ``attention`` then runs over pseudo-rows with a sliding window and
                   its finish collapses them (+ Gram-Schmidt XSA, gate).
* ``head_gate``    the gated byte head's ``[h, sigmoid(u.h + b) P h]`` input.
* ``sample``       repetition penalty / temperature / top-k / top-p / min-p + draw
                   (or greedy), one row per program.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

import torch

try:   # (remote code: transformers only requires imports outside try blocks)
    import triton
    import triton.language as tl
except ImportError as exc:  # pragma: no cover
    raise ImportError("the fused LuauCoder decoding kernels need triton") from exc

from .luau_gemv_tuned import TABLE as _TABLE

LOG2E = 1.4426950408889634

# programmatic dependent launch (SM90+); LUAU_INFER_PDL=0 turns it off
PDL = (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9
       and os.environ.get("LUAU_INFER_PDL", "1") != "0")

# L2 prefetch of the next GEMV's weights (SM90+); LUAU_INFER_PREFETCH=0 turns it off
BULK = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9
PREFETCH = BULK and os.environ.get("LUAU_INFER_PREFETCH", "1") != "0"


def _pf(pf) -> dict:
    """Launch kwargs of an optional prefetch target ``pf`` (a tensor)."""
    if pf is None or not PREFETCH:
        return dict(pf_ptr=None, pf_bytes=0, PF=False, BULK=False)
    return dict(pf_ptr=pf, pf_bytes=pf.numel() * pf.element_size(), PF=True, BULK=BULK)


@triton.jit
def _pdl(ON: tl.constexpr):
    """Wait for the previous kernel's writes, then let the next kernel launch."""
    if ON:
        tl.extra.cuda.gdc_wait()
        tl.extra.cuda.gdc_launch_dependents()


@triton.jit
def _prefetch(ptr, nbytes, PF: tl.constexpr, BULK: tl.constexpr):
    """Fire-and-forget L2 prefetch of ``nbytes`` at ``ptr``, split over the kernel's programs."""
    if PF:
        pid = tl.program_id(0) + tl.num_programs(0) * (tl.program_id(1) + tl.num_programs(1) * tl.program_id(2))
        npg = tl.num_programs(0) * tl.num_programs(1) * tl.num_programs(2)
        CH: tl.constexpr = 4096 if BULK else 128
        n = tl.cdiv(nbytes, CH)
        per = tl.cdiv(n, npg)
        lo = pid * per
        hi = tl.minimum(n, lo + per)
        base = ptr.to(tl.pointer_type(tl.int8))
        for i0 in range(lo, hi, 256):
            i = i0 + tl.arange(0, 256)
            m = (i < hi).to(tl.int32)
            a = base + i.to(tl.int64) * CH
            if BULK:
                sz = tl.minimum(nbytes - i.to(tl.int64) * CH, CH).to(tl.int32)
                tl.inline_asm_elementwise(
                    "{ .reg .pred p; setp.ne.b32 p, $3, 0; @p cp.async.bulk.prefetch.L2.global [$1], $2; "
                    "mov.u32 $0, 0; }", "=r,l,r,r", [a, sz, m], dtype=tl.int32, is_pure=False, pack=1)
            else:
                tl.inline_asm_elementwise(
                    "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p prefetch.global.L2 [$1]; mov.u32 $0, 0; }",
                    "=r,l,r", [a, m], dtype=tl.int32, is_pure=False, pack=1)


@triton.jit
def _mix_norm_kernel(s_ptr, s_row, s_ent, ids_ptr, emb_ptr, w_ptr, g_ptr, out_ptr, o_row, eps, D,
                     NE: tl.constexpr, NE_PAD: tl.constexpr, MODE: tl.constexpr, BLOCK: tl.constexpr,
                     ADD: tl.constexpr = 0, PART: tl.constexpr = 1, PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
    """MODE 0: h = S[r, 0]; 1: embedding gather; 2: AttnRes mix; 3: sum of entries. out = rmsnorm(h) * g."""
    r = tl.program_id(0).to(tl.int64)
    offs = tl.arange(0, BLOCK)
    msk = offs < D
    g = tl.load(g_ptr + offs, mask=msk, other=0.0).to(tl.float32)
    if MODE == 2:
        w = tl.load(w_ptr + offs, mask=msk, other=0.0)
    _prefetch(pf_ptr, pf_bytes, PF, BULK)
    _pdl(PDL)
    if MODE == 1:
        tok = tl.load(ids_ptr + r).to(tl.int64)
        hb = tl.load(emb_ptr + tok * D + offs, mask=msk, other=0.0)
        tl.store(s_ptr + r * s_row + offs, hb, mask=msk)
        h = hb.to(tl.float32)
    elif MODE == 0:
        h = tl.load(s_ptr + r * s_row + offs, mask=msk, other=0.0).to(tl.float32)
    else:
        j = tl.arange(0, NE_PAD)
        jm = j < NE
        v = tl.load(s_ptr + r * s_row + j[:, None] * s_ent + offs[None, :], mask=jm[:, None] & msk[None, :],
                    other=0.0).to(tl.float32)
        if MODE == 3:
            h = tl.sum(v, 0)
        else:
            sc = tl.sum(v * w[None, :], 1) * tl.rsqrt(tl.sum(v * v, 1) / D + eps)
            sc = tl.where(j < NE - 1 + PART, sc, float("-inf"))
            e = tl.exp(sc - tl.max(sc, 0))
            p = e / tl.sum(e, 0)
            h = tl.sum((p + ADD)[:, None] * v, 0)
        h = h.to(out_ptr.dtype.element_ty).to(tl.float32)   # the mix output is bf16
    y = h * tl.rsqrt(tl.sum(h * h, 0) / D + eps) * g
    tl.store(out_ptr + r * o_row + offs, y.to(out_ptr.dtype.element_ty), mask=msk)


def mix_norm(S: torch.Tensor, n_ent: int, mode: int, out: torch.Tensor, g: torch.Tensor, eps: float,
             w: torch.Tensor | None = None, ids: torch.Tensor | None = None,
             emb: torch.Tensor | None = None, add: bool = False, partial: bool = True, pf=None) -> None:
    """``S`` ``(R, NS, D)`` sources (row-major per row), ``out`` ``(R, D)``."""
    R, _, D = S.shape
    block = triton.next_power_of_2(D)
    ne_pad = triton.next_power_of_2(n_ent) if mode >= 2 else 1
    warps = max(4, min(16, ne_pad * block // 1024))
    _mix_norm_kernel[(R,)](S, S.stride(0), S.stride(1), ids if ids is not None else S,
                           emb if emb is not None else S, w if w is not None else g, g, out,
                           out.stride(0), eps, D, NE=n_ent, NE_PAD=ne_pad, MODE=mode, BLOCK=block,
                           ADD=int(add), PART=int(partial), num_warps=warps, PDL=PDL, launch_pdl=PDL, **_pf(pf))


_TUNED_PATH = Path(__file__).with_name("luau_gemv_tuned.py")
_TUNED: dict | None = None


def _tuned() -> dict:
    """Pinned GEMV configs for this GPU (``benchmark/tune_gemv.py``), else {}."""
    global _TUNED
    if _TUNED is None:
        name = torch.cuda.get_device_name()
        p = torch.cuda.get_device_properties(0)
        arch = [p.major, p.minor, p.multi_processor_count]
        hit = _TABLE.get(name) or next((e for e in _TABLE.values() if e.get("arch") == arch), None)
        _TUNED = dict(hit["gemv"]) if hit else {}
    return _TUNED


def gemv_key(m_pad: int, n: int, k: int, fp8: bool, epi: int) -> str:
    return f"{m_pad},{n},{k},{int(fp8)},{int(epi == 2)}"


def gemv_config(m_pad: int, n: int, k: int, fp8: bool, epi: int) -> tuple[int, int, int, int, int]:
    """``(BLOCK_N, BLOCK_K, num_warps, num_stages, DOT)``: pinned, else a default."""
    hit = _tuned().get(gemv_key(m_pad, n, k, fp8, epi))
    if hit:
        return tuple(hit) + (0,) * (5 - len(hit))
    if m_pad >= 4:
        return 32, 128, 4, 3, 1
    bk = 512 if fp8 else 256
    bn = max(1, min(16, 2048 // (bk * m_pad) * 4))
    return bn, bk, 4, 2, 0


@triton.jit
def _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn, nm, K, M_PAD: tl.constexpr, BLOCK_N: tl.constexpr,
               BLOCK_K: tl.constexpr, DOT: tl.constexpr):
    if DOT:
        acc = tl.zeros([M_PAD, BLOCK_N], dtype=tl.float32)
        for k0 in range(0, K, BLOCK_K):
            rk = k0 + tl.arange(0, BLOCK_K)
            km = rk < K
            x = tl.load(x_ptr + rm[:, None] * x_row + rk[None, :], mask=mm[:, None] & km[None, :], other=0.0)
            w = tl.load(w_ptr + rn[:, None] * w_row + rk[None, :], mask=nm[:, None] & km[None, :], other=0.0)
            acc = tl.dot(x.to(tl.bfloat16), tl.trans(w.to(tl.bfloat16)), acc)
        return acc
    acc3 = tl.zeros([M_PAD, BLOCK_N, BLOCK_K], dtype=tl.float32)
    for k0 in range(0, K, BLOCK_K):
        rk = k0 + tl.arange(0, BLOCK_K)
        km = rk < K
        x = tl.load(x_ptr + rm[:, None] * x_row + rk[None, :], mask=mm[:, None] & km[None, :], other=0.0)
        w = tl.load(w_ptr + rn[:, None] * w_row + rk[None, :], mask=nm[:, None] & km[None, :], other=0.0)
        acc3 += x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :]
    return tl.sum(acc3, 2)


@triton.jit
def _gemv_kernel(x_ptr, x_row, w_ptr, w_row, ws_ptr, y_ptr, y_row, M, N, K, up_off, cap,
                 M_PAD: tl.constexpr, EPI: tl.constexpr, FP8: tl.constexpr,
                 BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, DOT: tl.constexpr = False,
                 PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
    """y[m, n] = x[m] . W[n]. EPI 0: store; 1: residual add; 2: SwiGLU; 3: softcapped fp32 logits."""
    _prefetch(pf_ptr, pf_bytes, PF, BULK)
    _pdl(PDL)
    pid = tl.program_id(0)
    rn = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    nm = rn < N
    rm = tl.arange(0, M_PAD)
    mm = rm < M
    acc = _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn, nm, K, M_PAD, BLOCK_N, BLOCK_K, DOT)
    if FP8:
        acc = acc * tl.load(ws_ptr + rn, mask=nm, other=0.0)[None, :]
    ym = mm[:, None] & nm[None, :]
    yp = y_ptr + rm[:, None] * y_row + rn[None, :]
    if EPI == 2:
        up = _gemv_tile(x_ptr, x_row, w_ptr, w_row, rm, mm, rn + up_off, nm, K, M_PAD, BLOCK_N, BLOCK_K, DOT)
        if FP8:
            up = up * tl.load(ws_ptr + up_off + rn, mask=nm, other=0.0)[None, :]
        gate = acc.to(tl.bfloat16).to(tl.float32)
        up = up.to(tl.bfloat16).to(tl.float32)
        tl.store(yp, (gate * tl.sigmoid(gate) * up).to(y_ptr.dtype.element_ty), mask=ym)
    elif EPI == 1:
        res = tl.load(yp, mask=ym, other=0.0).to(tl.float32)
        tl.store(yp, (acc.to(tl.bfloat16).to(tl.float32) + res).to(y_ptr.dtype.element_ty), mask=ym)
    elif EPI == 3:
        z = acc.to(tl.bfloat16).to(tl.float32)
        if cap > 0:
            t = z / cap
            e = tl.exp(-2.0 * tl.abs(t))
            th = (1.0 - e) / (1.0 + e)
            z = cap * tl.where(t < 0, -th, th)
        tl.store(yp, z, mask=ym)
    else:
        tl.store(yp, acc.to(y_ptr.dtype.element_ty), mask=ym)


def gemv(x: torch.Tensor, w: torch.Tensor, y: torch.Tensor, epi: int = 0, scale: torch.Tensor | None = None,
         n: int | None = None, up_off: int = 0, cap: float = 0.0, config: tuple | None = None, pf=None) -> None:
    """``x`` ``(M, K)`` (row stride free), ``w`` ``(N_w, K)``, ``y`` ``(M, n)``; ``n`` output
    columns (SwiGLU: the hidden size, ``up_off`` = its up rows' offset)."""
    M, K = x.shape
    N = n if n is not None else w.shape[0]
    m_pad = max(1, triton.next_power_of_2(M))
    fp8 = scale is not None
    cfg = tuple(config or gemv_config(m_pad, N, K, fp8, epi))
    bn, bk, warps, stages, dot = cfg + (0,) * (5 - len(cfg))
    _gemv_kernel[(triton.cdiv(N, bn),)](x, x.stride(0), w, w.stride(0), scale if fp8 else w, y, y.stride(0),
                                        M, N, K, up_off, float(cap), M_PAD=max(16, m_pad) if dot else m_pad,
                                        EPI=epi, FP8=fp8, BLOCK_N=bn, BLOCK_K=bk, DOT=bool(dot),
                                        num_warps=warps, num_stages=stages, PDL=PDL, launch_pdl=PDL, **_pf(pf))


# CCGQA conv-latent / value-shift ring slots; a verify of up to RING - 3 drafts never overwrites
RING: tl.constexpr = tl.constexpr(32)


@triton.jit
def _lat(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, r, t, b, slot, pos, col, cidx, LAG: tl.constexpr):
    """The pre-conv latent of the row's token ``LAG`` back (zero across a document start)."""
    if LAG == 0:
        v = tl.load(qkv_ptr + r * qkv_row + col + cidx).to(tl.float32)
    elif t >= LAG:
        v = tl.load(qkv_ptr + (r - LAG) * qkv_row + col + cidx, mask=pos >= LAG, other=0.0).to(tl.float32)
    else:
        v = tl.load(lat_ptr + b * lat_b + ((slot - LAG) % RING) * lat_l + col + cidx, mask=pos >= LAG,
                    other=0.0).to(tl.float32)
    return tl.where(pos >= LAG, v, 0.0)


@triton.jit
def _conv(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, cs_ptr, csb_ptr, ch_ptr, chb_ptr, h, r, t, b, slot, pos, col,
          HD: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr):
    """Virtual head ``h``'s causal conv."""
    d = tl.arange(0, HD)
    b1 = tl.load(csb_ptr + h * HD + d).to(tl.float32)
    acc = tl.load(chb_ptr + h * HD + d).to(tl.float32)
    for j in tl.static_range(KC):
        u = b1
        for i in tl.static_range(KS):
            a = tl.load(cs_ptr + (h * HD + d) * KS + i).to(tl.float32)
            u += a * _lat(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, r, t, b, slot, pos, col, d, i + j)
        w2 = tl.load(ch_ptr + ((h * KC + j) * HD + d[:, None]) * HD + d[None, :]).to(tl.float32)
        acc += tl.sum(w2 * u[:, None], 0)
    return acc


@triton.jit
def _conv_pre(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, b1, acc, a0, a1, w0, w1, r, t, b, slot, pos, col,
              HD: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr):
    """``_conv`` (``KS``, ``KC`` <= 2) on taps already in registers (loaded before the PDL wait)."""
    d = tl.arange(0, HD)
    for j in tl.static_range(KC):
        u = b1
        for i in tl.static_range(KS):
            a = a0 if i == 0 else a1
            u += a * _lat(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, r, t, b, slot, pos, col, d, i + j)
        w2 = w0 if j == 0 else w1
        acc += tl.sum(w2 * u[:, None], 0)
    return acc


@triton.jit
def _qk_decode_kernel(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, vh_ptr, vh_b, vh_l,
                      pos_ptr, slot_ptr, cos_ptr, sin_ptr, cs_ptr, csb_ptr, ch_ptr, chb_ptr, qg_ptr, kg_ptr,
                      q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h, T, eps,
                      NQ: tl.constexpr, NKV: tl.constexpr, NXQ: tl.constexpr, HD: tl.constexpr,
                      V_OFF: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr, CCGQA: tl.constexpr,
                      QK_MEAN: tl.constexpr, VSHIFT: tl.constexpr, ROPE: tl.constexpr,
                      PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
    """One (row, virtual head): queries go to ``q``, keys / values to the caches."""
    h = tl.program_id(1)
    PRE: tl.constexpr = CCGQA and KS <= 2 and KC <= 2
    if PRE:
        dd = tl.arange(0, HD)
        cb1 = tl.load(csb_ptr + h * HD + dd).to(tl.float32)
        cacc = tl.load(chb_ptr + h * HD + dd).to(tl.float32)
        ca0 = tl.load(cs_ptr + (h * HD + dd) * KS).to(tl.float32)
        ca1 = tl.load(cs_ptr + (h * HD + dd) * KS + (KS - 1)).to(tl.float32)
        cw0 = tl.load(ch_ptr + ((h * KC) * HD + dd[:, None]) * HD + dd[None, :]).to(tl.float32)
        cw1 = tl.load(ch_ptr + ((h * KC + (KC - 1)) * HD + dd[:, None]) * HD + dd[None, :]).to(tl.float32)
    r = tl.program_id(0).to(tl.int64)
    b = r // T
    t = r % T
    pos = tl.load(pos_ptr + r).to(tl.int64)
    slot = tl.load(slot_ptr + r).to(tl.int64)
    _prefetch(pf_ptr, pf_bytes, PF, BULK)
    _pdl(PDL)
    HALF: tl.constexpr = HD // 2
    G: tl.constexpr = NQ // NKV
    d = tl.arange(0, HD)
    is_k = h >= NQ
    g = h - NQ
    ph = tl.where(is_k, NXQ + g, h % NXQ)
    col = ph * HD
    base = qkv_ptr + r * qkv_row
    keep = t + RING >= T   # only a call's last RING tokens go to the rings (no write races)
    if CCGQA:
        if PRE:
            x = _conv_pre(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, cb1, cacc, ca0, ca1, cw0, cw1, r, t, b, slot,
                          pos, col, HD, KS, KC)
        else:
            x = _conv(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, cs_ptr, csb_ptr, ch_ptr, chb_ptr, h, r, t, b, slot,
                      pos, col, HD, KS, KC)
        if QK_MEAN:
            if is_k:
                qm = tl.zeros([HD], dtype=tl.float32)
                for jh in tl.static_range(G):
                    qm += tl.load(base + ((g * G + jh) % NXQ) * HD + d).to(tl.float32)
                x += 0.5 * (qm / G + tl.load(base + col + d).to(tl.float32))
            else:
                x += 0.5 * (tl.load(base + col + d).to(tl.float32)
                            + tl.load(base + (NXQ + h // G) * HD + d).to(tl.float32))
    else:
        x = tl.load(base + col + d).to(tl.float32)
    gain = tl.load(tl.where(is_k, kg_ptr + g * HD, qg_ptr) + d)
    x = x * tl.rsqrt(tl.sum(x * x, 0) / HD + eps) * gain
    if not CCGQA:   # plain QK-norm rounds to bf16 before RoPE
        x = x.to(tl.bfloat16).to(tl.float32)
    lo, hi = tl.split(tl.permute(tl.reshape(x, [2, HALF]), (1, 0)))
    if ROPE:
        i2 = tl.arange(0, HALF)
        c = tl.load(cos_ptr + pos * HALF + i2)
        s = tl.load(sin_ptr + pos * HALF + i2)
        lo, hi = lo * c - hi * s, hi * c + lo * s
    lo_i = tl.arange(0, HALF)
    if is_k:
        dst = kc_ptr + b * kc_b + g * kc_h + slot * HD
    else:
        dst = q_ptr + r * q_row + h * HD
    tl.store(dst + lo_i, lo.to(q_ptr.dtype.element_ty))
    tl.store(dst + HALF + lo_i, hi.to(q_ptr.dtype.element_ty))
    if is_k:
        ch = g * HD + d
        own = tl.load(base + V_OFF + ch)
        v = own
        if VSHIFT:
            VH: tl.constexpr = NKV * HD // 2
            upper = ch >= VH
            if t >= 1:
                prev = tl.load(base - qkv_row + V_OFF + ch, mask=upper, other=0.0)
            else:
                prev = tl.load(vh_ptr + b * vh_b + ((slot - 1) % RING) * vh_l + ch - VH,
                               mask=upper & (slot >= 1), other=0.0)
            v = tl.where(upper, tl.where(pos > 0, prev, 0.0), own)
            tl.store(vh_ptr + b * vh_b + (slot % RING) * vh_l + ch - VH, own, mask=upper & keep)
        tl.store(vc_ptr + b * kc_b + g * kc_h + slot * HD + d, v.to(vc_ptr.dtype.element_ty))
    if CCGQA:
        if ((h < NXQ) | is_k) & keep:
            tl.store(lat_ptr + b * lat_b + (slot % RING) * lat_l + col + d, tl.load(base + col + d))


def qk_decode(qkv, lat, vh, pos, slot, cos, sin, conv, qg, kg, q, kc, vc, T: int, eps: float, *, nq, nkv,
              nxq, hd, v_off, ks, kc_taps, ccgqa, qk_mean, vshift, rope, pf=None) -> None:
    R = qkv.shape[0]
    cs, csb, ch, chb = conv if conv is not None else (qkv,) * 4
    _qk_decode_kernel[(R, nq + nkv)](
        qkv, qkv.stride(0), lat if lat is not None else qkv, lat.stride(0) if lat is not None else 0,
        lat.stride(1) if lat is not None else 0, vh if vh is not None else qkv,
        vh.stride(0) if vh is not None else 0, vh.stride(1) if vh is not None else 0,
        pos, slot, cos, sin, cs, csb, ch, chb, qg, kg, q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1),
        T, eps, NQ=nq, NKV=nkv, NXQ=nxq, HD=hd, V_OFF=v_off, KS=ks, KC=kc_taps, CCGQA=ccgqa,
        QK_MEAN=qk_mean, VSHIFT=vshift, ROPE=rope, num_warps=8, PDL=PDL, launch_pdl=PDL, **_pf(pf))


@triton.jit
def _tile_lat(qkv_ptr, qkv_row, r0, t, tm, pos, col, d, T, LAG: tl.constexpr):
    """``[BT, HD]`` pre-conv latents ``LAG`` tokens back of a prompt tile."""
    ok = tm & (pos >= LAG) & (t >= LAG)
    return tl.load(qkv_ptr + (r0 + t - LAG)[:, None] * qkv_row + col + d[None, :], mask=ok[:, None],
                   other=0.0).to(tl.float32)


@triton.jit
def _qk_prefill_kernel(qkv_ptr, qkv_row, lat_ptr, lat_b, lat_l, vh_ptr, vh_b, vh_l, pos_ptr, cos_ptr, sin_ptr,
                       cs_ptr, csb_ptr, ch_ptr, chb_ptr, qg_ptr, kg_ptr, q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h,
                       T, eps, NQ: tl.constexpr, NKV: tl.constexpr, NXQ: tl.constexpr, HD: tl.constexpr,
                       V_OFF: tl.constexpr, KS: tl.constexpr, KC: tl.constexpr, QK_MEAN: tl.constexpr,
                       VSHIFT: tl.constexpr, ROPE: tl.constexpr, BT: tl.constexpr):
    """``qk_decode`` for a prompt, ``BT`` tokens of one virtual head per program (taps on tensor cores)."""
    tt = tl.program_id(0)
    h = tl.program_id(1)
    b = tl.program_id(2).to(tl.int64)
    HALF: tl.constexpr = HD // 2
    G: tl.constexpr = NQ // NKV
    t = tt * BT + tl.arange(0, BT)
    tm = t < T
    r0 = b * T
    rows = r0 + t
    pos = tl.load(pos_ptr + rows, mask=tm, other=0).to(tl.int64)
    d = tl.arange(0, HD)
    is_k = h >= NQ
    g = h - NQ
    ph = tl.where(is_k, NXQ + g, h % NXQ)
    col = ph * HD
    b1 = tl.load(csb_ptr + h * HD + d).to(tl.float32)
    x = tl.zeros([BT, HD], dtype=tl.float32) + tl.load(chb_ptr + h * HD + d).to(tl.float32)[None, :]
    for j in tl.static_range(KC):
        u = tl.zeros([BT, HD], dtype=tl.float32) + b1[None, :]
        for i in tl.static_range(KS):
            a = tl.load(cs_ptr + (h * HD + d) * KS + i).to(tl.float32)
            u += a[None, :] * _tile_lat(qkv_ptr, qkv_row, r0, t, tm, pos, col, d, T, i + j)
        w2 = tl.load(ch_ptr + ((h * KC + j) * HD + d[:, None]) * HD + d[None, :]).to(tl.bfloat16)
        u_hi = u.to(tl.bfloat16)
        x = tl.dot(u_hi, w2, x)
        x = tl.dot((u - u_hi.to(tl.float32)).to(tl.bfloat16), w2, x)
    base = qkv_ptr + rows[:, None] * qkv_row
    if QK_MEAN:
        if is_k:
            qm = tl.zeros([BT, HD], dtype=tl.float32)
            for jh in tl.static_range(G):
                qm += tl.load(base + ((g * G + jh) % NXQ) * HD + d[None, :], mask=tm[:, None], other=0.0).to(
                    tl.float32)
            x += 0.5 * (qm / G + tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0).to(tl.float32))
        else:
            x += 0.5 * (tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0).to(tl.float32)
                        + tl.load(base + (NXQ + h // G) * HD + d[None, :], mask=tm[:, None], other=0.0).to(
                            tl.float32))
    gain = tl.load(tl.where(is_k, kg_ptr + g * HD, qg_ptr) + d)
    x = x * tl.rsqrt(tl.sum(x * x, 1) / HD + eps)[:, None] * gain[None, :]
    lo, hi = tl.split(tl.permute(tl.reshape(x, [BT, 2, HALF]), (0, 2, 1)))
    i2 = tl.arange(0, HALF)
    if ROPE:
        c = tl.load(cos_ptr + pos[:, None] * HALF + i2[None, :])
        sn = tl.load(sin_ptr + pos[:, None] * HALF + i2[None, :])
        lo, hi = lo * c - hi * sn, hi * c + lo * sn
    if is_k:
        dst = kc_ptr + b * kc_b + g * kc_h + t[:, None] * HD
    else:
        dst = q_ptr + rows[:, None] * q_row + h * HD
    tl.store(dst + i2[None, :], lo.to(q_ptr.dtype.element_ty), mask=tm[:, None])
    tl.store(dst + HALF + i2[None, :], hi.to(q_ptr.dtype.element_ty), mask=tm[:, None])
    keep = tm & (t + RING >= T)
    if is_k:
        ch = g * HD + d
        own = tl.load(base + V_OFF + ch[None, :], mask=tm[:, None], other=0.0)
        v = own
        if VSHIFT:
            VH: tl.constexpr = NKV * HD // 2
            upper = (ch >= VH)[None, :]
            prev = tl.load(base - qkv_row + V_OFF + ch[None, :], mask=tm[:, None] & upper & (t >= 1)[:, None],
                           other=0.0)
            v = tl.where(upper, tl.where((pos > 0)[:, None], prev, 0.0), own)
            tl.store(vh_ptr + b * vh_b + (t % RING)[:, None] * vh_l + ch[None, :] - VH, own,
                     mask=keep[:, None] & upper)
        tl.store(vc_ptr + b * kc_b + g * kc_h + t[:, None] * HD + d[None, :], v.to(vc_ptr.dtype.element_ty),
                 mask=tm[:, None])
    if (h < NXQ) | is_k:
        tl.store(lat_ptr + b * lat_b + (t % RING)[:, None] * lat_l + col + d[None, :],
                 tl.load(base + col + d[None, :], mask=tm[:, None], other=0.0), mask=keep[:, None])


def qk_prefill(qkv, lat, vh, pos, cos, sin, conv, qg, kg, q, kc, vc, T: int, eps: float, *, nq, nkv, nxq, hd,
               v_off, ks, kc_taps, qk_mean, vshift, rope) -> None:
    """CCGQA ``qk_decode`` for ``B`` prompts of ``T`` tokens starting at slot 0."""
    B = qkv.shape[0] // T
    cs, csb, ch, chb = conv
    bt = 32 if hd <= 128 else 16  # 64-token tiles spill
    _qk_prefill_kernel[(triton.cdiv(T, bt), nq + nkv, B)](
        qkv, qkv.stride(0), lat, lat.stride(0), lat.stride(1), vh if vh is not None else qkv,
        vh.stride(0) if vh is not None else 0, vh.stride(1) if vh is not None else 0, pos, cos, sin,
        cs, csb, ch, chb, qg, kg, q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1), T, eps, NQ=nq, NKV=nkv,
        NXQ=nxq, HD=hd, V_OFF=v_off, KS=ks, KC=kc_taps, QK_MEAN=qk_mean, VSHIFT=vshift, ROPE=rope, BT=bt,
        num_warps=4)


# grid (span, KV head, sequence); the last span program merges the partials and finishes the rows


@triton.jit
def _spans(start, slot, NSPL: tl.constexpr, BL: tl.constexpr, MINB: tl.constexpr = 1):
    """``(span, n_spans)`` for keys ``[start, slot]``, sized on the device."""
    n = slot + 1 - start
    span = tl.maximum(tl.cdiv(tl.cdiv(n, NSPL), BL), MINB) * BL
    return span, tl.where(n > 0, tl.cdiv(n, span), 0)


@triton.jit
def _merge(ml_ptr, acc_ptr, base0, n_sp, hrow, HD: tl.constexpr, G_PAD: tl.constexpr, GH_PAD: tl.constexpr,
           NSPL: tl.constexpr, CH: tl.constexpr):
    """Merge the span partials of the query heads ``hrow`` -> bf16-rounded output."""
    d = tl.arange(0, HD)
    si = tl.arange(0, NSPL)
    m = tl.load(ml_ptr + (base0 + si)[:, None] * 2 * G_PAD + hrow[None, :], mask=(si < n_sp)[:, None],
                other=float("-inf"))
    l = tl.load(ml_ptr + (base0 + si)[:, None] * 2 * G_PAD + G_PAD + hrow[None, :], mask=(si < n_sp)[:, None],
                other=0.0)
    mg = tl.max(m, 0)
    mg = tl.where(mg == float("-inf"), 0.0, mg)
    lsum = tl.sum(tl.exp2(m - mg[None, :]) * l, 0)
    o = tl.zeros([GH_PAD, HD], dtype=tl.float32)
    for c0 in range(0, n_sp, CH):
        cs = c0 + tl.arange(0, CH)
        cm = cs < n_sp
        mc = tl.load(ml_ptr + (base0 + cs)[:, None] * 2 * G_PAD + hrow[None, :], mask=cm[:, None],
                     other=float("-inf"))
        a = tl.load(acc_ptr + ((base0 + cs)[:, None, None] * G_PAD + hrow[None, :, None]) * HD
                    + d[None, None, :], mask=cm[:, None, None], other=0.0)
        o += tl.sum(tl.exp2(mc - mg[None, :])[:, :, None] * a, 0)
    o = o / tl.maximum(lsum, 1e-30)[:, None]
    return o.to(tl.bfloat16).to(tl.float32)


@triton.jit
def _finish(ml_ptr, acc_ptr, r, g, b, n_sp, slot, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row,
            lam_off, gate_off, NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr,
            HD: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr, GATE: tl.constexpr):
    """Merge row ``r``'s spans for KV head ``g`` and finish DiffAttn V2 + XSA (+ gate) into ``y``."""
    base0 = (r * NKV + g) * NSPL
    d = tl.arange(0, HD)
    if DIFF:
        GH: tl.constexpr = G // 2
        GH_PAD: tl.constexpr = G_PAD // 2
        i = tl.arange(0, GH_PAD)
        im = i < GH
        y0 = _merge(ml_ptr, acc_ptr, base0, n_sp, 2 * i, HD, G_PAD, GH_PAD, NSPL, CH)
        y1 = _merge(ml_ptr, acc_ptr, base0, n_sp, 2 * i + 1, HD, G_PAD, GH_PAD, NSPL, CH)
        lam = tl.load(qkv_ptr + r * qkv_row + lam_off + g * GH + i, mask=im, other=0.0).to(tl.float32)
        y = y0 - tl.sigmoid(lam)[:, None] * y1
    else:
        GH: tl.constexpr = G
        GH_PAD: tl.constexpr = G_PAD
        i = tl.arange(0, GH_PAD)
        im = i < GH
        y = _merge(ml_ptr, acc_ptr, base0, n_sp, i, HD, G_PAD, GH_PAD, NSPL, CH)
    if XSA:
        v = tl.load(vc_ptr + b * kc_b + g * kc_h + slot * HD + d).to(tl.float32)
        vn = v / tl.maximum(tl.sqrt(tl.sum(v * v, 0)), 1e-12)
        y = y - tl.sum(y * vn[None, :], 1)[:, None] * vn[None, :]
    oh = g * GH + i
    if GATE == 1:   # per channel
        gt = tl.load(qkv_ptr + r * qkv_row + gate_off + oh[:, None] * HD + d[None, :], mask=im[:, None],
                     other=0.0).to(tl.float32)
        y = y * tl.sigmoid(gt)
    elif GATE == 2:   # per head
        gt = tl.load(qkv_ptr + r * qkv_row + gate_off + oh, mask=im, other=0.0).to(tl.float32)
        y = y * tl.sigmoid(gt)[:, None]
    tl.store(y_ptr + r * y_row + oh[:, None] * HD + d[None, :], y.to(y_ptr.dtype.element_ty), mask=im[:, None])


@triton.jit
def _fence():
    """``fence.acq_rel.gpu`` in every thread (CUDA's ``__threadfence``)."""
    tl.inline_asm_elementwise("fence.acq_rel.gpu; mov.u32 $0, 0;", "=r", [], dtype=tl.int32, is_pure=False,
                              pack=1)


@triton.jit
def _seq_keys(start_ptr, slot_ptr, r0, T, win, WIN: tl.constexpr):
    """A sequence's key range ``[start, last]``."""
    start = tl.load(start_ptr + r0)
    if WIN:
        start = tl.maximum(start, tl.load(slot_ptr + r0) - win + 1)
    return start, tl.load(slot_ptr + r0 + T - 1)


@triton.jit
def _unit(x, tol):
    """``x / |x|``; 0 where ``|x| <= tol`` (the port's ``_unit``)."""
    n = tl.sqrt(tl.sum(x * x, 0))
    return tl.where(n > tol, x / tl.maximum(n, tol), 0.0)


@triton.jit
def _iha_finish(ml_ptr, acc_ptr, rt, g, b, n_sp, vs0, r_ptr, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr,
                y_row, gate_off, P: tl.constexpr, NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr,
                G_PAD: tl.constexpr, HD: tl.constexpr, CH: tl.constexpr, XSA: tl.constexpr, GATE: tl.constexpr):
    """IHA finish: merge, collapse with ``R``, Gram-Schmidt XSA and the gate for one token row."""
    d = tl.arange(0, HD)
    i = tl.arange(0, G_PAD)
    im = i < G
    oh = g * G + i
    y = tl.zeros([G_PAD, HD], dtype=tl.float32)
    for p in tl.static_range(P):
        yp = _merge(ml_ptr, acc_ptr, ((rt * P + p) * NKV + g) * NSPL, n_sp, i, HD, G_PAD, G_PAD, NSPL, CH)
        rr = tl.load(r_ptr + oh * P + p, mask=im, other=0.0).to(tl.float32)
        y += rr[:, None] * yp
    y = y.to(tl.bfloat16).to(tl.float32)
    if XSA:
        vb = vc_ptr + b * kc_b + g * kc_h + vs0 * HD
        e0 = tl.zeros([HD], dtype=tl.float32)
        e1 = tl.zeros([HD], dtype=tl.float32)
        e2 = tl.zeros([HD], dtype=tl.float32)
        for p in tl.static_range(P):
            w = _unit(tl.load(vb + p * HD + d).to(tl.float32), 1e-12)
            w = w - tl.sum(w * e0, 0) * e0
            w = w - tl.sum(w * e1, 0) * e1
            w = w - tl.sum(w * e2, 0) * e2
            e = _unit(w, 1e-2)
            y = y - tl.sum(y * e[None, :], 1)[:, None] * e[None, :]
            if p == 0:
                e0 = e
            elif p == 1:
                e1 = e
            elif p == 2:
                e2 = e
    if GATE == 1:
        gt = tl.load(qkv_ptr + rt * qkv_row + gate_off + oh[:, None] * HD + d[None, :], mask=im[:, None],
                     other=0.0).to(tl.float32)
        y = y * tl.sigmoid(gt)
    elif GATE == 2:
        gt = tl.load(qkv_ptr + rt * qkv_row + gate_off + oh, mask=im, other=0.0).to(tl.float32)
        y = y * tl.sigmoid(gt)[:, None]
    tl.store(y_ptr + rt * y_row + oh[:, None] * HD + d[None, :], y.to(y_ptr.dtype.element_ty), mask=im[:, None])


@triton.jit
def _attn_kernel(q_ptr, q_row, kc_ptr, vc_ptr, kc_b, kc_h, slot_ptr, start_ptr, ml_ptr, acc_ptr, cnt_ptr,
                 qkv_ptr, qkv_row, y_ptr, y_row, T, qk_scale, lam_off, gate_off, win, r_ptr,
                 NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr, QR: tl.constexpr,
                 HD: tl.constexpr, BL: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr,
                 GATE: tl.constexpr, FINISH: tl.constexpr, WIN: tl.constexpr = False, IHA: tl.constexpr = 0,
                 MINB: tl.constexpr = 1, PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
    """Flash-decoding attention; ``FINISH``: the last span program finishes the row."""
    sp = tl.program_id(0)
    g = tl.program_id(1)
    b = tl.program_id(2).to(tl.int64)
    r0 = b * T
    start, last = _seq_keys(start_ptr, slot_ptr, r0, T, win, WIN)
    span, n_sp = _spans(start, last, NSPL, BL, MINB)
    lo = start + sp * span
    hi = tl.minimum(last + 1, lo + span)
    qi = tl.arange(0, QR)
    t = qi // G
    j = qi % G
    qm = t < T
    lim = tl.load(slot_ptr + r0 + t, mask=qm, other=-1)
    _prefetch(pf_ptr, pf_bytes, PF, BULK)
    _pdl(PDL)
    if lo < hi:
        d = tl.arange(0, HD)
        q = tl.load(q_ptr + (r0 + t)[:, None] * q_row + (g * G + j)[:, None] * HD + d[None, :],
                    mask=qm[:, None], other=0.0)
        m_i = tl.full([QR], float("-inf"), tl.float32)
        l_i = tl.zeros([QR], dtype=tl.float32)
        acc = tl.zeros([QR, HD], dtype=tl.float32)
        kbase = b * kc_b + g * kc_h
        for k0 in range(lo, hi, BL):
            kk = k0 + tl.arange(0, BL)
            km = kk < hi
            k = tl.load(kc_ptr + kbase + kk[:, None] * HD + d[None, :], mask=km[:, None], other=0.0)
            s = tl.dot(q, tl.trans(k)) * qk_scale                       # base-2 logits
            ok = km[None, :] & (kk[None, :] <= lim[:, None])
            if WIN:
                ok = ok & (kk[None, :] > lim[:, None] - win)
            s = tl.where(ok, s, float("-inf"))
            m_new = tl.maximum(m_i, tl.max(s, 1))
            m_ok = tl.where(m_new == float("-inf"), 0.0, m_new)         # a row with no key yet
            alpha = tl.exp2(m_i - m_ok)
            p = tl.exp2(s - m_ok[:, None])
            l_i = l_i * alpha + tl.sum(p, 1)
            v = tl.load(vc_ptr + kbase + kk[:, None] * HD + d[None, :], mask=km[:, None], other=0.0)
            acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
            m_i = m_new
        base = ((r0 + t) * NKV + g) * NSPL + sp
        tl.store(ml_ptr + base * 2 * G_PAD + j, m_i, mask=qm)
        tl.store(ml_ptr + base * 2 * G_PAD + G_PAD + j, l_i, mask=qm)
        tl.store(acc_ptr + (base * G_PAD + j)[:, None] * HD + d[None, :], acc, mask=qm[:, None])
        if FINISH:
            _fence()
            tl.debug_barrier()
            old = tl.atomic_add(cnt_ptr + b * NKV + g, 1, sem="acq_rel", scope="gpu")
            if old == n_sp - 1:
                _fence()
                tl.atomic_xchg(cnt_ptr + b * NKV + g, 0, sem="relaxed", scope="gpu")
                if IHA:
                    _iha_finish(ml_ptr, acc_ptr, b, g, b, n_sp, tl.load(slot_ptr + r0), r_ptr, qkv_ptr, qkv_row,
                                vc_ptr, kc_b, kc_h, y_ptr, y_row, gate_off, IHA, NKV, NSPL, G, G_PAD, HD, CH,
                                XSA, GATE)
                else:
                    _finish(ml_ptr, acc_ptr, r0, g, b, n_sp, last, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr,
                            y_row, lam_off, gate_off, NKV, NSPL, G, G_PAD, HD, CH, DIFF, XSA, GATE)


@triton.jit
def _finish_kernel(ml_ptr, slot_ptr, start_ptr, acc_ptr, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row, T,
                   lam_off, gate_off, NKV: tl.constexpr, NSPL: tl.constexpr, G: tl.constexpr, G_PAD: tl.constexpr,
                   HD: tl.constexpr, BL: tl.constexpr, CH: tl.constexpr, DIFF: tl.constexpr, XSA: tl.constexpr,
                   GATE: tl.constexpr, MINB: tl.constexpr = 1, PDL: tl.constexpr = False):
    """The finish of a multi-row (verify / short prompt) attention, a program per (row, KV head)."""
    _pdl(PDL)
    r = tl.program_id(0).to(tl.int64)
    g = tl.program_id(1)
    b = r // T
    _, n_sp = _spans(tl.load(start_ptr + b * T), tl.load(slot_ptr + b * T + T - 1), NSPL, BL, MINB)
    _finish(ml_ptr, acc_ptr, r, g, b, n_sp, tl.load(slot_ptr + r), qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr,
            y_row, lam_off, gate_off, NKV, NSPL, G, G_PAD, HD, CH, DIFF, XSA, GATE)


@triton.jit
def _iha_finish_kernel(ml_ptr, slot_ptr, start_ptr, acc_ptr, qkv_ptr, qkv_row, vc_ptr, kc_b, kc_h, y_ptr, y_row,
                       Tt, win, r_ptr, gate_off, P: tl.constexpr, NKV: tl.constexpr, NSPL: tl.constexpr,
                       G: tl.constexpr, G_PAD: tl.constexpr, HD: tl.constexpr, BL: tl.constexpr, CH: tl.constexpr,
                       XSA: tl.constexpr, GATE: tl.constexpr, WIN: tl.constexpr, MINB: tl.constexpr = 1,
                       PDL: tl.constexpr = False):
    """IHA's finish of a multi-token call, a program per (token row, KV head)."""
    _pdl(PDL)
    rt = tl.program_id(0).to(tl.int64)
    g = tl.program_id(1)
    b = rt // Tt
    start, last = _seq_keys(start_ptr, slot_ptr, b * Tt * P, Tt * P, win, WIN)
    _, n_sp = _spans(start, last, NSPL, BL, MINB)
    _iha_finish(ml_ptr, acc_ptr, rt, g, b, n_sp, tl.load(slot_ptr + rt * P), r_ptr, qkv_ptr, qkv_row, vc_ptr,
                kc_b, kc_h, y_ptr, y_row, gate_off, P, NKV, NSPL, G, G_PAD, HD, CH, XSA, GATE)


# tests set 32 / False so a 1-row step and a verify match bitwise
QR_MIN = 16
FUSED_FINISH = True
# smallest span, in 64-key blocks
ATTN_MINB = 1


def attention(q, kc, vc, slot, start, qkv, y, ml, acc, cnt, T: int, *, nq, nkv, hd, nspl, lam_off,
              gate_off, diff, xsa, gate, win: int = 0, iha_p: int = 0, iha_r=None, pf=None, bl: int = 64,
              warps: int | None = None) -> None:
    """Flash-decoding attention for ``R = B * T`` rows (``win``: sliding window; ``iha_p``: pseudo-rows)."""
    if iha_p:
        assert not diff, "IHA with DiffAttn V2: no fused finish"
        T = T * iha_p
    R = q.shape[0]
    G = nq // nkv
    g_pad = max(16, triton.next_power_of_2(G))
    qr = max(QR_MIN, triton.next_power_of_2(T * G))
    ch = 8 if g_pad <= 16 else 4
    _attn_kernel[(nspl, nkv, R // T)](q, q.stride(0), kc, vc, kc.stride(0), kc.stride(1), slot, start, ml, acc,
                                      cnt, qkv, qkv.stride(0), y, y.stride(0), T, hd ** -0.5 * LOG2E, lam_off,
                                      gate_off, int(win), iha_r if iha_p else q,
                                      NKV=nkv, NSPL=nspl, G=G, G_PAD=g_pad, QR=qr, HD=hd, BL=bl,
                                      CH=ch, DIFF=diff, XSA=xsa, GATE=gate,
                                      FINISH=T == max(iha_p, 1) and FUSED_FINISH, WIN=win > 0, IHA=iha_p,
                                      MINB=ATTN_MINB,
                                      num_warps=warps or (4 if qr <= 32 else 8), num_stages=2, PDL=PDL, launch_pdl=PDL,
                                      **_pf(pf))
    if iha_p:
        if T > iha_p or not FUSED_FINISH:
            _iha_finish_kernel[(R // iha_p, nkv)](ml, slot, start, acc, qkv, qkv.stride(0), vc, kc.stride(0),
                                                  kc.stride(1), y, y.stride(0), T // iha_p, int(win), iha_r,
                                                  gate_off, P=iha_p, NKV=nkv, NSPL=nspl, G=G, G_PAD=g_pad,
                                                  HD=hd, BL=bl, CH=ch, XSA=xsa, GATE=gate, WIN=win > 0, MINB=ATTN_MINB,
                                                  num_warps=4, PDL=PDL, launch_pdl=PDL)
        return
    if T > 1 or not FUSED_FINISH:
        _finish_kernel[(R, nkv)](ml, slot, start, acc, qkv, qkv.stride(0), vc, kc.stride(0), kc.stride(1), y,
                                 y.stride(0), T, lam_off, gate_off, NKV=nkv, NSPL=nspl, G=G, G_PAD=g_pad, HD=hd,
                                 BL=bl, CH=ch, DIFF=diff, XSA=xsa, GATE=gate, MINB=ATTN_MINB, num_warps=4, PDL=PDL,
                                 launch_pdl=PDL)


@triton.jit
def _iha_mix_kernel(q_ptr, q_row, kl_ptr, vl_ptr, kl_b, kl_h, aq_ptr, ak_ptr, av_ptr, qg_ptr, kg_ptr, temp_ptr,
                    pos_ptr, slot_ptr, cos_ptr, sin_ptr, qv_ptr, qv_row, kc_ptr, vc_ptr, kc_b, kc_h, T, eps,
                    P: tl.constexpr, NQ: tl.constexpr, NQ_PAD: tl.constexpr, NKV: tl.constexpr, HD: tl.constexpr,
                    ROPE: tl.constexpr, PDL: tl.constexpr = False, pf_ptr=None, pf_bytes=0, PF: tl.constexpr = False, BULK: tl.constexpr = False):
    """One (row, pseudo-head p): IHA mixes, norm, temperature and RoPE, then cache / ``qv`` writes."""
    r = tl.program_id(0).to(tl.int64)
    p = tl.program_id(1)
    b = r // T
    HALF: tl.constexpr = HD // 2
    rslot = tl.load(slot_ptr + r).to(tl.int64)
    pos = tl.load(pos_ptr + r).to(tl.int64) * P + p
    slot = rslot * P + p
    d = tl.arange(0, HD)
    i2 = tl.arange(0, HALF)
    h = tl.arange(0, NQ_PAD)
    hm = h < NQ
    at = tl.load(aq_ptr + (h[None, :] * NQ + h[:, None]) * P + p, mask=hm[:, None] & hm[None, :], other=0.0)
    qgain = tl.load(qg_ptr + d)
    kgain = tl.load(kg_ptr + d)
    if ROPE:
        c = tl.load(cos_ptr + pos * HALF + i2)
        s = tl.load(sin_ptr + pos * HALF + i2)
    _prefetch(pf_ptr, pf_bytes, PF, BULK)
    _pdl(PDL)
    # queries: one tensor-core tile
    q = tl.load(q_ptr + r * q_row + h[:, None] * HD + d[None, :], mask=hm[:, None], other=0.0)
    x = tl.dot(at.to(tl.bfloat16), q.to(tl.bfloat16)).to(tl.bfloat16).to(tl.float32)   # at[h, m]
    x = x * tl.rsqrt(tl.sum(x * x, 1) / HD + eps)[:, None] * qgain[None, :]
    x = x.to(tl.bfloat16).to(tl.float32)
    lo, hi = tl.split(tl.permute(tl.reshape(x, [NQ_PAD, 2, HALF]), (0, 2, 1)))
    if ROPE:
        lo, hi = lo * c[None, :] - hi * s[None, :], hi * c[None, :] + lo * s[None, :]
    dst = qv_ptr + (r * P + p) * qv_row + h[:, None] * HD + i2[None, :]
    tl.store(dst, lo.to(qv_ptr.dtype.element_ty), mask=hm[:, None])
    tl.store(dst + HALF, hi.to(qv_ptr.dtype.element_ty), mask=hm[:, None])
    # keys / values: CUDA-core sums
    src = b * kl_b + rslot * HD
    for g in tl.static_range(NKV):
        k = tl.zeros([HD], dtype=tl.float32)
        v = tl.zeros([HD], dtype=tl.float32)
        for m in tl.static_range(NKV):
            ak = tl.load(ak_ptr + (m * NKV + g) * P + p).to(tl.float32)
            av = tl.load(av_ptr + (m * NKV + g) * P + p).to(tl.float32)
            k += ak * tl.load(kl_ptr + src + m * kl_h + d).to(tl.float32)
            v += av * tl.load(vl_ptr + src + m * kl_h + d).to(tl.float32)
        k = k.to(tl.bfloat16).to(tl.float32)
        k = (k * tl.rsqrt(tl.sum(k * k, 0) / HD + eps) * kgain).to(tl.bfloat16).to(tl.float32)
        k = (k * tl.load(temp_ptr + g)).to(tl.bfloat16).to(tl.float32)
        klo, khi = tl.split(tl.permute(tl.reshape(k, [2, HALF]), (1, 0)))
        if ROPE:
            klo, khi = klo * c - khi * s, khi * c + klo * s
        kd = kc_ptr + b * kc_b + g * kc_h + slot * HD
        tl.store(kd + i2, klo.to(kc_ptr.dtype.element_ty))
        tl.store(kd + HALF + i2, khi.to(kc_ptr.dtype.element_ty))
        tl.store(vc_ptr + b * kc_b + g * kc_h + slot * HD + d, v.to(vc_ptr.dtype.element_ty))


def iha_mix(q, kl, vl, aq, ak, av, qg, kg, temp, pos, slot, cos, sin, qv, kc, vc, T: int, eps: float, *,
            P: int, nq: int, nkv: int, hd: int, rope: bool, pf=None) -> None:
    """``q`` ``(R, nq, hd)`` latent queries, ``kl`` / ``vl`` ``(B, nkv, L, hd)`` latent keys / values at the
    rows' slots, mixes ``(H, H, P)`` bf16, ``qv`` ``(R * P, nq, hd)``, caches ``(B, nkv, L * P, hd)``."""
    R = q.shape[0]
    _iha_mix_kernel[(R, P)](q, q.stride(0), kl, vl, kl.stride(0), kl.stride(1), aq, ak, av, qg, kg, temp, pos,
                            slot, cos, sin, qv, qv.stride(0), kc, vc, kc.stride(0), kc.stride(1), T, eps,
                            P=P, NQ=nq, NQ_PAD=max(16, triton.next_power_of_2(nq)), NKV=nkv, HD=hd, ROPE=rope,
                            num_warps=4, PDL=PDL, launch_pdl=PDL, **_pf(pf))


@triton.jit
def _head_gate_kernel(y_ptr, y_row, b_ptr, h_ptr, h_row, D, RK: tl.constexpr, BLOCK: tl.constexpr,
                      PDL: tl.constexpr = False):
    """Gated byte head input ``[h, sigmoid(u . h + b) * P h]``."""
    _pdl(PDL)
    r = tl.program_id(0).to(tl.int64)
    j = tl.arange(0, BLOCK)
    jm = j < RK
    a = tl.load(y_ptr + r * y_row + RK).to(tl.float32) + tl.load(b_ptr).to(tl.float32)
    gt = tl.sigmoid(a.to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16).to(tl.float32)
    yp = tl.load(y_ptr + r * y_row + j, mask=jm, other=0.0).to(tl.float32)
    tl.store(h_ptr + r * h_row + D + j, (gt * yp).to(h_ptr.dtype.element_ty), mask=jm)


def head_gate(y, bias, h, D: int, rank: int) -> None:
    _head_gate_kernel[(y.shape[0],)](y, y.stride(0), bias, h, h.stride(0), D, RK=rank,
                                     BLOCK=triton.next_power_of_2(rank), num_warps=4, PDL=PDL, launch_pdl=PDL)


@triton.jit
def _sample_kernel(x_ptr, x_row, V, inv_temp, top_k, top_p, min_p, seed_ptr, excl_ptr, y_ptr, pd_ptr, u_ptr,
                   pres_ptr, pres_row, ids_ptr, penalty, T,
                   TOP_K: tl.constexpr, TOP_P: tl.constexpr, MIN_P: tl.constexpr, EXCL: tl.constexpr,
                   REP: tl.constexpr, GREEDY: tl.constexpr, T_MAX: tl.constexpr, BLOCK: tl.constexpr,
                   ITERS: tl.constexpr):
    """One row: repetition penalty, temperature, top-k, top-p, min-p (thresholds by bisection), then a
    draw."""
    r = tl.program_id(0).to(tl.int64)
    offs = tl.arange(0, BLOCK)
    msk = offs < V
    x = tl.load(x_ptr + r * x_row + offs, mask=msk, other=float("-inf"))
    if REP:
        b = r // T
        t = r % T
        seen = tl.load(pres_ptr + b * pres_row + offs, mask=msk, other=0) != 0
        for i in tl.static_range(1, T_MAX):
            if i <= t:
                seen = seen | (offs == tl.load(ids_ptr + b * T + i))
        x = tl.where(seen, tl.where(x < 0, x * penalty, x / penalty), x)
    if GREEDY:
        tl.store(y_ptr + r, tl.argmax(x, 0).to(tl.int64))
    else:
        x = x * inv_temp
        if TOP_K:
            lo = tl.min(tl.where(msk, x, float("inf")), 0)
            hi = tl.max(x, 0)
            for _ in range(ITERS):   # count(x >= lo) >= k > count(x >= hi)
                mid = lo + (hi - lo) * 0.5
                ok = tl.sum((x >= mid).to(tl.int32), 0) >= top_k
                lo = tl.where(ok, mid, lo)
                hi = tl.where(ok, hi, mid)
            kth = tl.min(tl.where(x >= lo, x, float("inf")), 0)
            x = tl.where(x >= kth, x, float("-inf"))
        m = tl.max(x, 0)
        p = tl.exp(x - m)                      # max 1
        z = tl.sum(p, 0)
        if TOP_P:
            lo = tl.zeros([], dtype=tl.float32)
            hi = tl.full([], 1.0, tl.float32)  # mass(p >= lo) >= top_p * z > mass(p >= hi) (or hi = 1)
            goal = top_p * z
            for _ in range(ITERS):
                mid = lo + (hi - lo) * 0.5
                ok = tl.sum(tl.where(p >= mid, p, 0.0), 0) >= goal
                lo = tl.where(ok, mid, lo)
                hi = tl.where(ok, hi, mid)
            tau = tl.min(tl.where(p >= lo, p, float("inf")), 0)
            p = tl.where(p >= tau, p, 0.0)
        if MIN_P:
            p = tl.where(p >= min_p, p, 0.0)   # max is 1: min_p * max = min_p
        seed = tl.load(seed_ptr)
        if EXCL:
            e = tl.load(excl_ptr + r)
            pd = tl.sum(tl.where(offs == e, p, 0.0), 0) / tl.sum(p, 0)
            tl.store(pd_ptr + r, pd)
            tl.store(u_ptr + r, tl.rand(seed + 7919, r.to(tl.int32)))
            p = tl.where(offs == e, 0.0, p)
        cdf = tl.cumsum(p, 0)
        u = tl.rand(seed, r.to(tl.int32)) * tl.max(cdf, 0)
        tok = tl.min(tl.where((cdf > u) & (p > 0), offs, BLOCK), 0)
        tl.store(y_ptr + r, tl.where(tok < V, tok, tl.argmax(p, 0)).to(tl.int64))


def sample(x: torch.Tensor, seed: torch.Tensor, y: torch.Tensor, *, temperature: float, top_k: int, top_p: float,
           min_p: float, excl: torch.Tensor | None = None, pd: torch.Tensor | None = None,
           u: torch.Tensor | None = None, greedy: bool = False, penalty: float = 1.0,
           pres: torch.Tensor | None = None, ids: torch.Tensor | None = None, T: int = 1) -> None:
    """Draw ``y`` ``(R,)`` from fp32 logits ``x`` ``(R, V)``."""
    R, V = x.shape
    block = triton.next_power_of_2(V)
    rep = penalty != 1.0
    _sample_kernel[(R,)](x, x.stride(0), V, 1.0 / temperature, int(top_k), float(top_p), float(min_p), seed,
                         excl if excl is not None else y, y, pd if pd is not None else x,
                         u if u is not None else x, pres.view(torch.uint8) if rep else y,
                         pres.stride(0) if rep else 0, ids if rep else y, float(penalty), T,
                         TOP_K=bool(top_k) and top_k < V, TOP_P=top_p < 1.0, MIN_P=bool(min_p),
                         EXCL=excl is not None, REP=rep, GREEDY=greedy, T_MAX=T, BLOCK=block, ITERS=40,
                         num_warps=32 if block >= 16384 else 8)


kernels = sys.modules[__name__]   # ``from .luau_kernels import kernels as K``: transformers'
                                  # remote code only follows ``from .module import ...`` lines