Nalu-2.0-Base / modeling_luau_coder.py
Upload folder using huggingface_hub
1cc867a verified about 8 hours ago
86.2 kB
from __future__ import annotations

import copy
import inspect
import math
import warnings

import torch
import torch.nn as nn
import torch.nn.functional as F
import transformers
from torch import Tensor
from transformers import GenerationMixin, PreTrainedModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.masking_utils import create_causal_mask
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

from .configuration_luau_coder import LuauCoderConfig

try:
    from transformers import initialization as _init  # v5: skips weights loaded from a checkpoint
except ImportError:  # pragma: no cover
    _init = nn.init

try:
    import triton
    import triton.language as tl
    HAVE_TRITON = True
except ImportError:  # pragma: no cover
    HAVE_TRITON = False

__all__ = ["LuauCoderCache", "LuauCoderForCausalLM", "LuauCoderModel", "LuauCoderPreTrainedModel"]

_V5 = int(transformers.__version__.split(".")[0]) >= 5
IGNORE_INDEX = -100
_HW_TANH = torch.version.hip is None  # PTX tanh.approx in the fused CE; ROCm falls back to sigmoid


# fused elementwise


class _Fused:
    """``fn`` through ``torch.compile`` on CUDA (one generated kernel per
    direction); eager on CPU, inside an outer compile, when disabled, or after
    a compilation failure (warned once)."""

    def __init__(self, fn):
        self.fn = fn
        self.compiled = None
        self.failed = False

    def __call__(self, *args, fuse: bool = True):
        if (not fuse or self.failed or not args[0].is_cuda or torch.compiler.is_compiling()):
            return self.fn(*args)
        if self.compiled is None:
            self.compiled = torch.compile(self.fn, dynamic=None)
        try:
            return self.compiled(*args)
        except torch.OutOfMemoryError:
            raise
        except Exception as exc:  # no working Inductor / Triton toolchain here
            self.failed = True
            warnings.warn(f"torch.compile of {self.fn.__name__} failed ({exc!r}); running it eagerly",
                          stacklevel=2)
            return self.fn(*args)


_CONSTS: dict = {}


def _cached(tensors: tuple, fn):
    """``fn()``, a value derived from the parameters ``tensors``. Without grad
    (inference, outside ``torch.compile``) it is cached until one of them changes:
    decoding would otherwise rebuild these fp32 weights at every token."""
    if torch.is_grad_enabled() or torch.compiler.is_compiling():
        return fn()
    key = tuple((t.data_ptr(), t._version) for t in tensors)
    hit = _CONSTS.get(ids := tuple(id(t) for t in tensors))
    if hit is None or hit[0] != key:
        hit = _CONSTS[ids] = (key, fn())
    return hit[1]


def _rms(x: Tensor, w: Tensor, eps: float) -> Tensor:
    return F.rms_norm(x.float(), (x.shape[-1],), _cached((w,), w.float), eps).to(x.dtype)


def _rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    """Rotate-half RoPE; ``x`` ``(B, T, H, d)``, ``cos`` / ``sin`` ``(B|1, T, d/2)`` fp32."""
    d2 = x.shape[-1] // 2
    xf = x.float()
    x1, x2 = xf[..., :d2], xf[..., d2:]
    c, s = cos.unsqueeze(2), sin.unsqueeze(2)
    return torch.cat((x1 * c - x2 * s, x2 * c + x1 * s), dim=-1).to(x.dtype)


def _rms_norm_fn(x: Tensor, w: Tensor, eps: float) -> Tensor:
    return _rms(x, w, eps)


def _add_rms_norm_fn(x: Tensor, y: Tensor, w: Tensor, eps: float) -> tuple[Tensor, Tensor]:
    h = x + y
    return h, _rms(h, w, eps)


def _qk_rope_fn(q: Tensor, k: Tensor, qw: Tensor | None, kw: Tensor | None, cos: Tensor | None,
                sin: Tensor | None, eps: float) -> tuple[Tensor, Tensor]:
    """QK-norm, then RoPE (``cos`` None: NoPE layer, no rotation)."""
    if qw is not None:
        q, k = _rms(q, qw, eps), _rms(k, kw, eps)
    if cos is None:
        return q, k
    return _rope(q, cos, sin), _rope(k, cos, sin)


def _diff_xsa_fn(y: Tensor, lam: Tensor | None, v: Tensor, nh: int, nkv: int, xsa: bool,
                 gate: Tensor | None = None) -> Tensor:
    """``y`` ``(B, T, nq, d)`` attention output, ``lam`` ``(B, T, nh)`` lambda
    logits (None: no DiffAttn), ``v`` ``(B, T, nkv, d)`` the tokens' own values,
    ``gate`` output-gate logits ``(B, T, nh * d)`` per channel or ``(B, T, nh)`` per head
    (None: no gate), applied after XSA."""
    B, T, _, hd = y.shape
    if lam is not None:
        y = y.view(B, T, nh, 2, hd).float()
        y = y[..., 0, :] - torch.sigmoid(lam.float()).unsqueeze(-1) * y[..., 1, :]
    else:
        y = y.float()
    if xsa:
        y = y.view(B, T, nkv, nh // nkv, hd)
        vn = F.normalize(v.float(), dim=-1).unsqueeze(3)
        y = y - (y * vn).sum(-1, keepdim=True) * vn
    if gate is not None:
        y = y.reshape(B, T, nh, hd) * torch.sigmoid(gate.float()).view(B, T, nh, -1)
    return y.reshape(B, T, nh * hd).to(v.dtype)


def _unit(x: Tensor, tol: float) -> Tensor:
    """``x / |x|`` over the last dim; 0 where ``|x| <= tol``."""
    n = x.norm(dim=-1, keepdim=True)
    return torch.where(n > tol, x / n.clamp_min(tol), 0.0)


def _iha_diff_xsa(y: Tensor, lam: Tensor | None, vp: Tensor, nh: int, nkv: int, xsa: bool,
                  gate: Tensor | None = None) -> Tensor:
    """``_diff_xsa_fn`` for IHA's collapsed heads (training: ``_iha_diff_xsa_gate``): XSA removes the
    span of the token's own pseudo-values ``vp`` ``(B, T, P, nkv, d)``, by modified Gram-Schmidt."""
    B, T, _, hd = y.shape
    if lam is not None:
        y = y.view(B, T, nh, 2, hd).float()
        y = y[..., 0, :] - torch.sigmoid(lam.float()).unsqueeze(-1) * y[..., 1, :]
    else:
        y = y.float()
    if xsa:
        y = y.view(B, T, nkv, nh // nkv, hd)
        basis: list[Tensor] = []
        for p in range(vp.shape[2]):
            w = _unit(vp[:, :, p].float(), 1e-12)
            for e in basis:
                w = w - (w * e).sum(-1, keepdim=True) * e
            basis.append(_unit(w, 1e-2))
            e = basis[-1].unsqueeze(3)
            y = y - (y * e).sum(-1, keepdim=True) * e
    if gate is not None:
        y = y.reshape(B, T, nh, hd) * torch.sigmoid(gate.float()).view(B, T, nh, -1)
    return y.reshape(B, T, nh * hd).to(vp.dtype)


def _iha_attention(q: Tensor, k: Tensor, v: Tensor, mask: Tensor, scaling: float) -> Tensor:
    """GQA-native attention with an explicit mask: ``q`` ``(B, nq, Tq, d)``, ``k`` / ``v`` ``(B, nkv, Tk, d)``,
    ``mask`` ``(B|1, 1, 1, Tq, Tk)`` bool -> ``(B, Tq, nq, d)``; fp32 scores and softmax."""
    B, nq, Tq, d = q.shape
    nkv = k.shape[1]
    s = torch.matmul(q.reshape(B, nkv, nq // nkv, Tq, d).float(), k.float().unsqueeze(2).transpose(-1, -2)) * scaling
    p = s.masked_fill(~mask, float("-inf")).softmax(-1)              # (B, nkv, G, Tq, Tk)
    return torch.matmul(p.to(v.dtype), v.unsqueeze(2)).reshape(B, nq, Tq, d).transpose(1, 2)


def _iha_mix_init(alpha: Tensor, stride: int) -> None:
    """Training's init: pseudo-head ``p`` of head ``h`` starts as base head ``(h + p * stride) % H``
    plus a small Gaussian (``alpha`` ``(H, H, P)``; ``stride`` 2 for DiffAttn's query heads)."""
    H, _, P = alpha.shape
    if getattr(alpha, "_is_hf_initialized", False):   # loaded from a checkpoint
        return
    with torch.no_grad():
        alpha.normal_(std=0.1 / H ** 0.5)
        h = torch.arange(H, device=alpha.device)
        for p in range(P):
            alpha[(h + p * stride) % H, h, p] += 1.0


def _swiglu_fn(gu: Tensor) -> Tensor:
    gate, up = gu.chunk(2, dim=-1)
    return F.silu(gate) * up


def _ccgqa_fn(qk: Tensor, tail: Tensor, v: Tensor, before: Tensor | None, pos: Tensor, W: Tensor,
              b: Tensor, q_gain: Tensor, k_gain: Tensor, cos: Tensor | None, sin: Tensor | None, nq: int,
              nkv: int, eps: float, qk_mean: bool) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor | None]:
    """CCGQA ``q`` ``(B, T, nq, d)``, ``k`` / ``v`` ``(B, T, nkv, d)`` from the q|k latents
    ``qk`` ``(B, T, nxq+nkv, d)`` (``nxq < nq``: VHA's physical query heads) and the raw
    values, in fp32 like ``ccgqa_qk_reference``. ``tail`` holds the ``K-1`` latents
    before ``qk`` and ``before`` the previous token's value-shift half (None: no value
    shift); both come back updated. ``W`` / ``b`` / gains: ``_ccgqa_consts``. Tap ``m``
    reads token ``t-m`` only if ``pos[t] >= m``."""
    B, T, _, hd = qk.shape
    K, nxq = tail.shape[1] + 1, qk.shape[2] - nkv
    ext = torch.cat((tail, qk), 1)                            # (B, K-1+T, nxq+nkv, d)

    def virtual(x: Tensor) -> Tensor:   # VHA: query head h reads physical head h % nxq
        return x if nxq == nq else torch.cat((x[:, :, :nxq].repeat(1, 1, nq // nxq, 1), x[:, :, nxq:]), 2)

    win = virtual(ext).unfold(1, K, 1)                        # (B, T, H, d, K): slot j = lag K-1-j
    lag = torch.arange(K - 1, -1, -1, device=qk.device)
    win = torch.where((pos[..., None] >= lag)[:, :, None, None, :], win.float(), 0.0)
    H = win.shape[2]
    win = win.permute(2, 0, 1, 3, 4).reshape(H, B * T, hd * K)
    s = torch.bmm(win, W).view(H, B, T, hd).permute(1, 2, 0, 3) + b
    if qk_mean:
        xf = virtual(qk).float()
        mu_q = 0.5 * (xf[:, :, :nq] + xf[:, :, nq:].repeat_interleave(nq // nkv, 2))
        s = s + torch.cat((mu_q, mu_q.view(B, T, nkv, nq // nkv, hd).mean(3)), 2)
    n = F.rms_norm(s, (hd,), eps=eps)
    q, k = n[:, :, :nq] * q_gain, n[:, :, nq:] * k_gain
    if cos is not None:
        q, k = _rope(q, cos, sin), _rope(k, cos, sin)
    q, k, tail = q.to(qk.dtype), k.to(qk.dtype), ext[:, ext.shape[1] - (K - 1):]
    if before is None:
        return q, k, v.contiguous(), tail, None
    flat = v.reshape(B, T, nkv * hd)
    half = nkv * hd // 2          # the W_V-bar half of the channels is the previous token's
    shifted = torch.where((pos > 0)[..., None], torch.cat((before, flat[:, :, half:]), 1)[:, :T], 0)
    v = torch.cat((flat[:, :, :half], shifted), -1).view(B, T, nkv, hd)
    return q, k, v, tail, flat[:, -1:, half:]


_rms_norm = _Fused(_rms_norm_fn)
_add_rms_norm = _Fused(_add_rms_norm_fn)
_qk_rope = _Fused(_qk_rope_fn)
_diff_xsa = _Fused(_diff_xsa_fn)
_swiglu = _Fused(_swiglu_fn)
_ccgqa = _Fused(_ccgqa_fn)


# Block AttnRes kernels
#
# One mix, per token, over v = [S_0 .. S_{k-1}, x] (k sources plus the current
# block's running sum x) with score weight w = norm * query:
#
#     r_i = (mean(v_i^2) + eps)^-1/2,   s_i = r_i <v_i, w>,   p = softmax(s)
#     out = sum_i p_i v_i
#
# Backward, with g = dL/dout:
#
#     c_i  = p_i (<g, v_i> - sum_j p_j <g, v_j>)
#     dv_i = p_i g + c_i r_i w - (c_i <v_i, w> r_i^3 / D) v_i
#     dw   = sum_tokens sum_i c_i r_i v_i
#
# A block's sources are read by every mix until the next block starts. Their
# gradient is linear in per-mix vectors with per-token scalars,
# dS_i = sum_c a_ci g_c + sum_c b_ci w_c - (sum_c e_ci) S_i, so each mix
# backward stores g_c and the scalars and one kernel builds dS (and the score
# weights' S part) once per block, after its last mix backward.
#
# Delta Block (attn_res_delta_block, arXiv:2605.18855 App. C): the mix is added
# to the residual stream sum_i v_i, so out = sum_i (1 + p_i) v_i and dv_i gains g.
# The output is the plain stream. Without attn_res_partial_source x isn't routed:
# its score is -inf, so p_x = 0 and dx is just g.

_MAX_HIDDEN = 8192     # whole row in registers
_MAX_SOURCES = 32
_SPAN_BWD = dict(R=64, num_warps=2)
_TPP = 8               # tokens per program in the per-token backward

if HAVE_TRITON:

    @triton.jit
    def _ar_fwd_kernel(s_ptr, x_ptr, w_ptr, out_ptr, prob_ptr, r_ptr, dot_ptr,
                       s_st, s_sk, x_st, o_st, st_st, eps,
                       D: tl.constexpr, NB: tl.constexpr, NB_PAD: tl.constexpr, BLOCK: tl.constexpr,
                       ADD: tl.constexpr, PART: tl.constexpr):
        """One token: online softmax over the NB sources + x, one read of each
        (``ADD``: plus their sum, the residual stream; ``PART`` off: x isn't routed)."""
        t = tl.program_id(0).to(tl.int64)
        offs = tl.arange(0, BLOCK)
        msk = offs < D
        jb = tl.arange(0, NB_PAD)
        w = tl.load(w_ptr + offs, mask=msk, other=0.0)
        acc = tl.zeros([BLOCK], dtype=tl.float32)
        stream = tl.zeros([BLOCK], dtype=tl.float32)
        m_run = tl.full([], float("-inf"), tl.float32)
        l_run = tl.zeros([], dtype=tl.float32)
        svec = tl.zeros([NB_PAD], dtype=tl.float32)
        rvec = tl.zeros([NB_PAD], dtype=tl.float32)
        dvec = tl.zeros([NB_PAD], dtype=tl.float32)
        for j in tl.static_range(NB + 1):
            if j < NB:
                v = tl.load(s_ptr + t * s_st + j * s_sk + offs, mask=msk, other=0.0).to(tl.float32)
            else:
                v = tl.load(x_ptr + t * x_st + offs, mask=msk, other=0.0).to(tl.float32)
            r = tl.rsqrt(tl.sum(v * v, 0) / D + eps)
            dot = tl.sum(v * w, 0)
            s = dot * r if j < NB or PART else tl.full([], float("-inf"), tl.float32)
            m_new = tl.maximum(m_run, s)
            alpha = tl.exp(m_run - m_new)
            e = tl.exp(s - m_new)
            acc = acc * alpha + e * v
            if ADD:
                stream += v
            l_run = l_run * alpha + e
            m_run = m_new
            svec = tl.where(jb == j, s, svec)
            rvec = tl.where(jb == j, r, rvec)
            dvec = tl.where(jb == j, dot, dvec)
        tl.store(out_ptr + t * o_st + offs, (stream + acc / l_run).to(out_ptr.dtype.element_ty), mask=msk)
        valid = jb <= NB
        p = tl.exp(svec - m_run) / l_run
        tl.store(prob_ptr + t * st_st + jb, p, mask=valid)
        tl.store(r_ptr + t * st_st + jb, rvec, mask=valid)
        tl.store(dot_ptr + t * st_st + jb, dvec, mask=valid)

    @triton.jit
    def _ar_bwd_token_kernel(s_ptr, x_ptr, g_ptr, res_ptr, w_ptr, prob_ptr, r_ptr, dot_ptr,
                             dx_ptr, gbuf_ptr, coef_ptr, part_ptr,
                             s_st, s_sk, x_st, g_st, res_st, dx_st, st_st, coef_q, N,
                             D: tl.constexpr, NB: tl.constexpr, NB_PAD: tl.constexpr,
                             HAS_RES: tl.constexpr, TPP: tl.constexpr, BLOCK: tl.constexpr,
                             ADD: tl.constexpr):
        """TPP tokens: <g, v_i> for every entry (one read each), the full
        gradient of x (plus the residual-add path's, ``HAS_RES``), the scalars
        the deferred source gradient needs, and the x part of dw."""
        pid = tl.program_id(0)
        offs = tl.arange(0, BLOCK)
        msk = offs < D
        jb = tl.arange(0, NB_PAD)
        valid = jb <= NB
        src = jb < NB
        is_x = jb == NB
        w = tl.load(w_ptr + offs, mask=msk, other=0.0)
        dwx = tl.zeros([BLOCK], dtype=tl.float32)
        for i in range(TPP):
            t = pid.to(tl.int64) * TPP + i
            ok = t < N
            m = msk & ok
            g = tl.load(g_ptr + t * g_st + offs, mask=m, other=0.0).to(tl.float32)
            tl.store(gbuf_ptr + t * D + offs, g.to(gbuf_ptr.dtype.element_ty), mask=m)
            xv = tl.load(x_ptr + t * x_st + offs, mask=m, other=0.0).to(tl.float32)
            gv = tl.zeros([NB_PAD], dtype=tl.float32)
            for j in tl.static_range(NB):
                v = tl.load(s_ptr + t * s_st + j * s_sk + offs, mask=m, other=0.0).to(tl.float32)
                gv = tl.where(jb == j, tl.sum(g * v, 0), gv)
            gv = tl.where(jb == NB, tl.sum(g * xv, 0), gv)
            p = tl.load(prob_ptr + t * st_st + jb, mask=valid & ok, other=0.0)
            r = tl.load(r_ptr + t * st_st + jb, mask=valid & ok, other=0.0)
            dot = tl.load(dot_ptr + t * st_st + jb, mask=valid & ok, other=0.0)
            c = p * (gv - tl.sum(p * gv, 0))
            b = c * r
            e = b * dot * r * r / D
            tl.store(coef_ptr + t * NB + jb, p + ADD, mask=src & ok)
            tl.store(coef_ptr + coef_q + t * NB + jb, b, mask=src & ok)
            tl.store(coef_ptr + 2 * coef_q + t * NB + jb, e, mask=src & ok)
            p_x = tl.sum(tl.where(is_x, p, 0.0), 0)
            b_x = tl.sum(tl.where(is_x, b, 0.0), 0)
            e_x = tl.sum(tl.where(is_x, e, 0.0), 0)
            dx = (p_x + ADD) * g + b_x * w - e_x * xv
            if HAS_RES:
                dx += tl.load(res_ptr + t * res_st + offs, mask=m, other=0.0).to(tl.float32)
            tl.store(dx_ptr + t * dx_st + offs, dx.to(dx_ptr.dtype.element_ty), mask=m)
            dwx += b_x * xv
        tl.store(part_ptr + pid.to(tl.int64) * D + offs, dwx, mask=msk)

    @triton.jit
    def _ar_span_bwd_kernel(s_ptr, gbuf_ptr, w_ptr, coef_ptr, ds_ptr, part_ptr,
                            s_st, s_sk, ds_st, ds_sk, N, coef_c, coef_q,
                            D: tl.constexpr, NB: tl.constexpr, NC: tl.constexpr, NC_PAD: tl.constexpr,
                            R: tl.constexpr, CD: tl.constexpr):
        """dS for every source of one block from all NC mixes that read it (one
        read + one write per source), plus the S part of each mix's dw."""
        pid_t = tl.program_id(0)
        cols = tl.program_id(1) * CD + tl.arange(0, CD)
        mc = cols < D
        cc = tl.arange(0, NC_PAD)
        mcc = cc < NC
        m2 = mcc[:, None] & mc[None, :]
        W = tl.load(w_ptr + cc[:, None] * D + cols[None, :], mask=m2, other=0.0)
        dW = tl.zeros([NC_PAD, CD], dtype=tl.float32)
        for i in range(R):
            t = (pid_t * R + i).to(tl.int64)
            ok = t < N
            G = tl.load(gbuf_ptr + cc[:, None] * N * D + t * D + cols[None, :], mask=m2 & ok,
                        other=0.0).to(tl.float32)
            for j in range(NB):
                v = tl.load(s_ptr + t * s_st + j * s_sk + cols, mask=mc & ok, other=0.0).to(tl.float32)
                base = coef_ptr + cc * coef_c + t * NB + j
                a = tl.load(base, mask=mcc & ok, other=0.0)
                b = tl.load(base + coef_q, mask=mcc & ok, other=0.0)
                e = tl.load(base + 2 * coef_q, mask=mcc & ok, other=0.0)
                dv = tl.sum(a[:, None] * G, 0) + tl.sum(b[:, None] * W, 0) - tl.sum(e, 0) * v
                tl.store(ds_ptr + t * ds_st + j * ds_sk + cols, dv.to(ds_ptr.dtype.element_ty),
                         mask=mc & ok)
                dW += b[:, None] * v[None, :]
        tl.store(part_ptr + pid_t * NC * D + cc[:, None] * D + cols[None, :], dW, mask=m2)


def _ar_block(d: int) -> tuple[int, int]:
    block = triton.next_power_of_2(d)
    return block, (4 if block <= 2048 else 8 if block <= 4096 else 16)


def _pad2(n: int) -> int:
    return max(2, triton.next_power_of_2(n))


def _ar_forward(x: Tensor, S: Tensor, w: Tensor, eps: float, additive: bool,
                partial: bool) -> tuple[Tensor, Tensor]:
    """Launch the mix forward; returns ``out`` and the per-token stats."""
    N, k, D = S.shape
    out = torch.empty_like(x)
    stats = torch.empty((3, N, k + 1), device=x.device, dtype=torch.float32)
    block, warps = _ar_block(D)
    _ar_fwd_kernel[(N,)](S, x, w, out, stats[0], stats[1], stats[2],
                         S.stride(0), S.stride(1), x.stride(0), out.stride(0), stats.stride(1),
                         float(eps), D=D, NB=k, NB_PAD=_pad2(k + 1), BLOCK=block, ADD=int(additive),
                         PART=int(partial), num_warps=warps)
    return out, stats


def attn_res_reference(x: Tensor, sources: Tensor, norm_weight: Tensor, query: Tensor,
                       eps: float, additive: bool = False, partial: bool = True) -> Tensor:
    """Eager fp32 mix; ``x`` ``(N, D)``, ``sources`` ``(N, k, D)``. ``additive``: Delta
    Block, the mix added to the stream ``x + sum(sources)``. ``partial=False``: ``x``
    takes no part in the softmax."""
    values = torch.cat((sources, x.unsqueeze(1)), dim=1).float()
    keys = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + eps)
    logits = (keys * (norm_weight.float() * query.float())).sum(-1)
    if not partial:
        logits = torch.cat((logits[:, :-1], logits[:, -1:] - float("inf")), dim=1)
    probs = logits.softmax(-1)
    if additive:
        probs = probs + 1.0
    return (probs.unsqueeze(-1) * values).sum(1).to(x.dtype)


class _SpanState:
    """Shared by one block's sources and every mix reading them."""

    def __init__(self, n_calls: int, additive: bool = False, partial: bool = True):
        self.n_calls = n_calls
        self.additive, self.partial = additive, partial   # Delta Block; x is routed
        self.w: Tensor | None = None          # (n_calls, D) fp32 score weights
        self.gbuf: Tensor | None = None       # (n_calls, N, D) output grads, filled in backward
        self.coef: Tensor | None = None       # (n_calls, 3, N, k) fp32 a / b / e
        self.dw_x: Tensor | None = None       # (n_calls, D) fp32 x part of dw
        self.res: dict[int, Tensor] = {}      # call -> residual-add gradient of its x
        self.mixed: set[int] = set()          # calls whose mix backward has run
        self.done = 0


class _SpanSources(torch.autograd.Function):
    """Identity on the sources ``S``; its backward (after every mix of the
    block) builds ``dS`` and the score-weight gradients in one pass."""

    @staticmethod
    def forward(ctx, S: Tensor, state: _SpanState, *weights: Tensor) -> Tensor:
        norms, queries = weights[0::2], weights[1::2]
        state.w = torch.stack([n.float() * q.float() for n, q in zip(norms, queries)])
        ctx.state = state
        ctx.save_for_backward(S, *weights)
        ctx.set_materialize_grads(False)
        return S.view_as(S)

    @staticmethod
    def backward(ctx, grad_S):
        S, *weights = ctx.saved_tensors
        st: _SpanState = ctx.state
        N, k, D = S.shape
        nc = st.n_calls
        if st.done != nc:
            raise RuntimeError(f"AttnRes span: {st.done}/{nc} mix backwards ran before the block's")
        dS = torch.empty_like(S)
        R, nw = _SPAN_BWD["R"], _SPAN_BWD["num_warps"]
        CD = min(1024, 4096 // _pad2(nc))
        grid = (triton.cdiv(N, R), triton.cdiv(D, CD))
        part = torch.empty((grid[0], nc, D), device=S.device, dtype=torch.float32)
        _ar_span_bwd_kernel[grid](S, st.gbuf, st.w, st.coef, dS, part,
                                  S.stride(0), S.stride(1), dS.stride(0), dS.stride(1), N,
                                  st.coef.stride(0), st.coef.stride(1),
                                  D=D, NB=k, NC=nc, NC_PAD=_pad2(nc), R=R, CD=CD, num_warps=nw)
        dw = part.sum(0) + st.dw_x
        if grad_S is not None:
            dS = dS + grad_S
        grads = []
        for c in range(nc):
            norm, query = weights[2 * c], weights[2 * c + 1]
            grads += [(dw[c] * query.float()).to(norm.dtype), (dw[c] * norm.float()).to(query.dtype)]
        st.gbuf = st.coef = st.dw_x = None
        return (dS, None, *grads)


class _ResidualAdd(torch.autograd.Function):
    """``x + y`` where ``x`` is also read by mix ``call``: the backward hands
    the gradient of ``x`` to that mix's backward kernel, which sums it into
    ``dx``, instead of autograd summing the two in a separate pass."""

    @staticmethod
    def forward(ctx, x: Tensor, y: Tensor, state: _SpanState, call: int) -> Tensor:
        ctx.state, ctx.call = state, call
        return x + y

    @staticmethod
    def backward(ctx, g: Tensor):
        st = ctx.state
        if ctx.call in st.mixed:  # the mix already ran its backward: let autograd sum
            return g, g, None, None
        st.res[ctx.call] = g
        return None, g, None, None


class _DepthMix(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: Tensor, S: Tensor, state: _SpanState, call: int, eps: float) -> Tensor:
        x = x if x.stride(-1) == 1 else x.contiguous()
        out, stats = _ar_forward(x, S, state.w[call], eps, state.additive, state.partial)
        ctx.save_for_backward(x, S, stats)
        ctx.state, ctx.call = state, call
        return out

    @staticmethod
    def backward(ctx, g: Tensor):
        x, S, stats = ctx.saved_tensors
        st: _SpanState = ctx.state
        N, k, D = S.shape
        g = g if g.stride(-1) == 1 else g.contiguous()
        if st.gbuf is None:
            st.gbuf = torch.empty((st.n_calls, N, D), device=x.device, dtype=g.dtype)
            st.coef = torch.empty((st.n_calls, 3, N, k), device=x.device, dtype=torch.float32)
            st.dw_x = torch.zeros((st.n_calls, D), device=x.device, dtype=torch.float32)
        c = ctx.call
        res = st.res.pop(c, None)
        if res is not None and res.stride(-1) != 1:
            res = res.contiguous()
        dx = torch.empty_like(x)
        block, warps = _ar_block(D)
        grid = (triton.cdiv(N, _TPP),)
        part = torch.empty((grid[0], D), device=x.device, dtype=torch.float32)
        _ar_bwd_token_kernel[grid](S, x, g, g if res is None else res, st.w[c], stats[0], stats[1],
                                   stats[2], dx, st.gbuf[c], st.coef[c], part,
                                   S.stride(0), S.stride(1), x.stride(0), g.stride(0),
                                   0 if res is None else res.stride(0), dx.stride(0),
                                   stats.stride(1), st.coef.stride(1), N,
                                   D=D, NB=k, NB_PAD=_pad2(k + 1), HAS_RES=res is not None, TPP=_TPP,
                                   BLOCK=block, ADD=int(st.additive), num_warps=warps)
        st.dw_x[c] = part.sum(0)
        st.mixed.add(c)
        st.done += 1
        return dx, None, None, None, None


class AttnResSpan:
    """The sources ``S`` ``(N, k, D)`` of one block and the mixes reading them.
    ``sites[c]`` is the ``(norm, query)`` pair of the ``c``-th mix that reads
    ``S``. With grad enabled every site must be called exactly once per
    forward. Paths: Triton with autograd (training), Triton forward only (no
    grad: inference / generation), eager fp32 reference (CPU, no Triton,
    ``fused_kernels=False``). ``additive``: Delta Block mixes; ``partial``: ``x`` is routed too.
    """

    def __init__(self, S: Tensor, sites: list[tuple[Tensor, Tensor]], eps: float, fused: bool,
                 additive: bool = False, partial: bool = True):
        self.eps, self.sites = eps, sites
        self.additive, self.partial = additive, partial
        fast = (fused and HAVE_TRITON and S.is_cuda and S.shape[-1] <= _MAX_HIDDEN
                and S.shape[1] <= _MAX_SOURCES)
        self.train = fast and torch.is_grad_enabled()
        self.infer = fast and not self.train
        self.S = S
        if self.train:
            self.state = _SpanState(len(sites), additive, partial)
            self.S = _SpanSources.apply(S, self.state, *[t for site in sites for t in site])
        elif self.infer:
            self.w = _cached(tuple(t for site in sites for t in site),
                             lambda: torch.stack([n.float() * q.float() for n, q in sites]))

    @torch.compiler.disable
    def mix(self, x: Tensor, call: int) -> Tensor:
        if self.train:
            return _DepthMix.apply(x, self.S, self.state, call, self.eps)
        if self.infer:
            return _ar_forward(x if x.stride(-1) == 1 else x.contiguous(), self.S, self.w[call],
                               self.eps, self.additive, self.partial)[0]
        norm, query = self.sites[call]
        return attn_res_reference(x, self.S, norm, query, self.eps, self.additive, self.partial)

    @torch.compiler.disable
    def stream(self, x: Tensor) -> Tensor:
        """The residual stream ``sum(S) + x`` (Delta Block's output)."""
        return (self.S.float().sum(1) + x.float()).to(x.dtype)

    @torch.compiler.disable
    def residual(self, x: Tensor, y: Tensor, call: int) -> Tensor:
        """``x + y`` for an ``x`` that only mix ``call`` also reads."""
        if self.train:
            return _ResidualAdd.apply(x, y, self.state, call)
        return x + y


# fused LM head + CE

def softcap_logits(logits: Tensor, cap: float | None) -> Tensor:
    """``cap * tanh(logits / cap)`` (``final_logit_softcapping``); identity for None / 0."""
    return cap * torch.tanh(logits / cap) if cap else logits


if HAVE_TRITON:

    @triton.jit
    def _tanh(x, HW: tl.constexpr):
        if HW:  # one MUFU op (sm_75+); rel. error 2^-11, under the bf16 rounding of the logits
            return tl.inline_asm_elementwise("tanh.approx.f32 $0, $1;", "=r,r", [x],
                                             dtype=tl.float32, is_pure=True, pack=1)
        return 2.0 * tl.sigmoid(2.0 * x) - 1.0

    @triton.jit
    def _ce_fwd_bwd_kernel(logits_ptr, stride_row, labels_ptr, loss_ptr, scale_ptr, V, cap, inv_cap,
                           IGNORE: tl.constexpr, SOFTCAP: tl.constexpr, HW_TANH: tl.constexpr,
                           BLOCK: tl.constexpr):
        """NLL of one row, then its logits overwritten by ``(softmax - onehot) * scale``
        (with ``SOFTCAP``: of ``cap * tanh(logits / cap)``, times ``1 - tanh^2``).
        Base 2 throughout so ``log2(e)`` folds into the ``ex2`` argument."""
        LOG2E: tl.constexpr = 1.4426950408889634
        LN2: tl.constexpr = 0.6931471805599453
        row = tl.program_id(0).to(tl.int64)
        ptr = logits_ptr + row * stride_row
        label = tl.load(labels_ptr + row).to(tl.int64)
        valid = label != IGNORE
        cap2 = cap * LOG2E
        safe_label = tl.where(valid, label, 0)
        # the target stays out of the sum-exp so 1 - p_t is s_nt / s, not
        # 1 - round(p_t) (Softmax Collapse, Barron et al. 2026)
        m = -float("inf")
        s_nt = 0.0
        for off in range(0, V, BLOCK):
            cols = off + tl.arange(0, BLOCK)
            msk = cols < V
            if SOFTCAP:
                x = tl.load(ptr + cols, mask=msk, other=0.0).to(tl.float32)
                x = tl.where(msk, cap2 * _tanh(x * inv_cap, HW_TANH), -float("inf"))
            else:
                x = tl.load(ptr + cols, mask=msk, other=-float("inf")).to(tl.float32) * LOG2E
            m_new = tl.maximum(m, tl.max(x, axis=0))
            e = tl.where(cols == safe_label, 0.0, tl.exp2(x - m_new))
            s_nt = s_nt * tl.exp2(m - m_new) + tl.sum(e, axis=0)
            m = m_new
        target = tl.load(ptr + safe_label).to(tl.float32)
        if SOFTCAP:
            target = cap2 * _tanh(target * inv_cap, HW_TANH)
        else:
            target = target * LOG2E
        s = tl.exp2(target - m) + s_nt
        u = 1.0 + s_nt  # target is the max: loss = log1p(s_nt) (Goldberg's log1p)
        log1p = tl.where(u == 1.0, s_nt, tl.log(u) * (s_nt / (u - 1.0)))
        loss = tl.where(target >= m, log1p, (m + tl.log2(s) - target) * LN2)
        tl.store(loss_ptr + row, tl.where(valid, loss, 0.0))
        r = tl.where(valid, tl.load(scale_ptr), 0.0) / s
        g_t = -s_nt * r  # -(1 - p_t) * scale: the row sums to zero
        for off in range(0, V, BLOCK):
            cols = off + tl.arange(0, BLOCK)
            msk = cols < V
            x = tl.load(ptr + cols, mask=msk, other=0.0).to(tl.float32)
            if SOFTCAP:
                t = _tanh(x * inv_cap, HW_TANH)
                g = tl.exp2(cap2 * t - m) * r
            else:
                g = tl.exp2(x * LOG2E - m) * r
            g = tl.where(cols == safe_label, g_t, g)
            if SOFTCAP:
                g = g - g * t * t
            tl.store(ptr + cols, g.to(logits_ptr.dtype.element_ty), mask=msk)


def _accumulate_dw(dw: Tensor, dlogits: Tensor, h: Tensor) -> None:
    """``dw += dlogits.T @ h`` with an fp32 output from low-precision operands."""
    try:
        torch.addmm(dw, dlogits.T, h, out_dtype=torch.float32, out=dw)
    except (TypeError, RuntimeError):
        dw.add_(dlogits.T @ h)


class _LinearCrossEntropy(torch.autograd.Function):
    """Chunked ``CE(h @ W.T, y)`` (sum over valid targets times ``scale``):
    per chunk one logits GEMM, one Triton pass that writes ``dlogits`` over
    the logits, and the two gradient GEMMs, so the backward only rescales."""

    @staticmethod
    def forward(ctx, hidden: Tensor, weight: Tensor, labels: Tensor, scale: Tensor, chunk: int,
                softcap: float | None) -> Tensor:
        h = hidden.reshape(-1, hidden.shape[-1])
        y = labels.reshape(-1)
        n_rows, V = h.shape[0], weight.shape[0]
        w = weight.to(h.dtype)
        losses = torch.empty(n_rows, device=h.device, dtype=torch.float32)
        need_dh = hidden.requires_grad
        dh = torch.empty_like(h) if need_dh else None
        dw = torch.zeros(weight.shape, device=h.device, dtype=torch.float32) if weight.requires_grad else None
        block = min(4096, triton.next_power_of_2(V))
        cap = float(softcap or 1.0)
        for s in range(0, n_rows, chunk):
            hc = h[s:s + chunk]
            logits = hc @ w.T
            _ce_fwd_bwd_kernel[(hc.shape[0],)](logits, logits.stride(0), y[s:s + chunk],
                                               losses[s:s + chunk], scale, V, cap, 1.0 / cap,
                                               IGNORE=IGNORE_INDEX, SOFTCAP=bool(softcap),
                                               HW_TANH=_HW_TANH, BLOCK=block, num_warps=8 if block >= 2048 else 4)
            if dh is not None:
                torch.mm(logits, w, out=dh[s:s + chunk])
            if dw is not None:
                _accumulate_dw(dw, logits, hc)
        ctx.save_for_backward(dh if dh is not None else h.new_empty(0),
                              dw if dw is not None else h.new_empty(0))
        ctx.hidden_shape, ctx.has_dh, ctx.has_dw = hidden.shape, dh is not None, dw is not None
        ctx.weight_dtype = weight.dtype
        return losses.sum() * scale

    @staticmethod
    def backward(ctx, grad_out: Tensor):
        dh, dw = ctx.saved_tensors
        dh = dh.mul_(grad_out.to(dh.dtype)).view(ctx.hidden_shape) if ctx.has_dh else None
        dw = dw.mul_(grad_out).to(ctx.weight_dtype) if ctx.has_dw else None
        return dh, dw, None, None, None, None


def _loss_scale(labels: Tensor, num_items_in_batch) -> Tensor:
    """``1 / num_items_in_batch`` if given (HF Trainer), else ``1 / #valid targets``."""
    if num_items_in_batch is None:
        n = (labels != IGNORE_INDEX).sum()
    else:
        n = torch.as_tensor(num_items_in_batch, device=labels.device)
    return 1.0 / n.clamp_min(1).to(device=labels.device, dtype=torch.float32)


@torch.compiler.disable
def fused_linear_cross_entropy(hidden: Tensor, weight: Tensor, labels: Tensor,
                               num_items_in_batch=None, chunk: int = 4096,
                               softcap: float | None = None) -> Tensor:
    return _LinearCrossEntropy.apply(hidden, weight, labels, _loss_scale(labels, num_items_in_batch),
                                     chunk, softcap)


# CCGQA


def compose_conv(w1: Tensor, b1: Tensor, w2: Tensor, b2: Tensor) -> tuple[Tensor, Tensor]:
    """Depthwise taps ``w1`` ``(H*hd, ks)`` + ``b1`` and per-head taps ``w2``
    ``(H, kc, hd_in, hd_out)`` + ``b2`` ``(H, hd)`` (taps in lag order, 0 = current
    token) -> one masked causal conv: fp32 ``W`` ``(H, ks+kc-1, hd, hd)``, ``b``
    ``(H, hd)``. Exact inside a document (``lm/kernels/ccgqa.py``)."""
    H, kc, hd, _ = w2.shape
    ks = w1.shape[1]
    a, w2f = w1.float().view(H, hd, ks, 1), w2.float()
    W = torch.stack([sum(a[:, :, i] * w2f[:, m - i] for i in range(ks) if 0 <= m - i < kc)
                     for m in range(ks + kc - 1)], 1)
    b = (b1.float().view(H, 1, hd, 1) * w2f).sum((1, 2)) + b2.float()
    return W, b


class LuauCoderCache(DynamicCache):
    """``DynamicCache`` plus CCGQA's per-layer decode state: the last ``K-1``
    pre-conv q|k latents and the previous token's raw value-shift half.
    Rewinding (``crop``: assisted / speculative decoding) is not supported
    once that state exists."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.luau_conv_states: dict[int, tuple[Tensor, Tensor | None]] = {}

    def _map(self, fn) -> None:
        self.luau_conv_states = {i: tuple(None if t is None else fn(t) for t in st)
                                 for i, st in self.luau_conv_states.items()}

    def reorder_cache(self, beam_idx: Tensor) -> None:
        super().reorder_cache(beam_idx)
        self._map(lambda t: t.index_select(0, beam_idx.to(t.device)))

    def batch_repeat_interleave(self, repeats: int) -> None:
        super().batch_repeat_interleave(repeats)
        self._map(lambda t: t.repeat_interleave(repeats, 0))

    def batch_select_indices(self, indices: Tensor) -> None:
        super().batch_select_indices(indices)
        self._map(lambda t: t[indices])

    def crop(self, max_length: int) -> None:
        if self.luau_conv_states and (max_length < 0 or max_length < self.get_seq_length()):
            raise NotImplementedError("the CCGQA conv state can't be rewound (no assisted decoding)")
        super().crop(max_length)

    def reset(self) -> None:
        if hasattr(super(), "reset"):
            super().reset()
        self.luau_conv_states.clear()


class _InPlaceStates(dict):
    """A conv-state dict whose entries, once made, are updated in place (fixed addresses)."""

    def __setitem__(self, i, new):
        old = self.get(i)
        if old is None:
            super().__setitem__(i, tuple(None if t is None else t.clone() for t in new))
        else:
            for o, n in zip(old, new):
                if o is not None:
                    o.copy_(n)


class LuauCoderStaticCache:
    """Preallocated cache for CUDA-graph decoding (``generate.py``): K/V buffers
    ``(B, nkv, L, d)`` per layer, CCGQA's conv / value-shift state, and which
    buffer slots hold real keys. Every tensor keeps its address and the write
    slots live on the device, so a captured decode step replays as is.
    Before each forward the caller sets ``idx`` (the slots of this call's
    tokens, ``(T,)``) and marks them in ``key_valid``, and passes
    ``position_ids`` (each token's position in its document; a left-padded
    row starts at its first real token). Query ``t`` attends to the valid
    slots ``<= idx[t]``. Only ``generate.py``'s loop drives it: not for
    ``generate()``.
    """

    is_static = True

    def __init__(self, config, batch: int, max_len: int, device, dtype):
        nkv, hd = config.num_key_value_heads, config.head_dim
        self.max_len = max_len
        self.layers = [tuple(torch.zeros(batch, nkv, max_len, hd, device=device, dtype=dtype)
                             for _ in range(2)) for _ in range(config.num_hidden_layers)]
        self.key_valid = torch.zeros(batch, max_len, dtype=torch.bool, device=device)
        self.slots = torch.arange(max_len, device=device)
        self.idx = self.slots[:0]
        self.luau_conv_states = _InPlaceStates()

    def update(self, k: Tensor, v: Tensor, layer_idx: int, *args, **kwargs) -> tuple[Tensor, Tensor]:
        K, V = self.layers[layer_idx]
        K.index_copy_(2, self.idx, k)
        V.index_copy_(2, self.idx, v)
        return K, V

    def mask(self) -> Tensor:
        """``(B, 1, T, L)`` bool: valid slots up to each query's own."""
        return self.key_valid[:, None, None, :] & (self.slots <= self.idx[:, None])[None, None]

    def reset(self) -> None:
        """Empty (same addresses): the conv state at zero acts as no state."""
        self.key_valid.zero_()
        for st in self.luau_conv_states.values():
            for t in st:
                if t is not None:
                    t.zero_()

    def get_seq_length(self, layer_idx: int = 0) -> int:
        return int(self.key_valid.sum(1).max())


def _conv_states(cache) -> dict:
    """The CCGQA state dict of any cache (attached to a plain ``DynamicCache`` too)."""
    states = getattr(cache, "luau_conv_states", None)
    if states is None:
        states = cache.luau_conv_states = {}
    return states


# RoPE


def _rope_settings(config: LuauCoderConfig) -> tuple[float, dict | None]:
    """``(theta, yarn params or None)`` from a v4 (``rope_theta`` + ``rope_scaling``)
    or v5 (``rope_parameters``) config."""
    params = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {}
    theta = float(params.get("rope_theta", getattr(config, "rope_theta", 500000.0)))
    rtype = params.get("rope_type", params.get("type"))
    if rtype == "yarn" and float(params.get("factor", 1.0)) > 1.0:
        return theta, params
    return theta, None


def rope_inv_freq(head_dim: int, theta: float, yarn: dict | None) -> tuple[Tensor, float]:
    """Inverse frequencies ``(head_dim/2,)`` fp64 and the cos/sin scale; YaRN
    "NTK-by-parts" as in the training stack and transformers' ``rope_type: yarn``."""
    dim = head_dim
    pos_freqs = theta ** (torch.arange(0, dim, 2, dtype=torch.float64, device="cpu") / dim)
    extrap = 1.0 / pos_freqs
    if yarn is None:
        return extrap, 1.0
    factor = float(yarn["factor"])
    L = int(yarn.get("original_max_position_embeddings", yarn.get("original_max_position", 4096)))
    beta_fast, beta_slow = float(yarn.get("beta_fast", 32.0)), float(yarn.get("beta_slow", 1.0))
    interp = 1.0 / (factor * pos_freqs)

    def correction_dim(rotations: float) -> float:
        return (dim * math.log(L / (rotations * 2 * math.pi))) / (2 * math.log(theta))

    low, high = correction_dim(beta_fast), correction_dim(beta_slow)
    if yarn.get("truncate", True):
        low, high = math.floor(low), math.ceil(high)
    low, high = max(low, 0), min(high, dim - 1)
    if low == high:
        high += 0.001
    ramp = ((torch.arange(dim // 2, dtype=torch.float64, device="cpu") - low) / (high - low)).clamp(0, 1)
    keep = 1.0 - ramp
    mscale = yarn.get("attention_factor")
    mscale = float(mscale) if mscale is not None else 0.1 * math.log(factor) + 1.0
    return interp * (1 - keep) + extrap * keep, mscale


def lerope_inv_freq(head_dim: int, theta: float, alpha: list[float], yarn: dict | None) -> tuple[Tensor, float]:
    """LeRoPE frequencies fp64 and the cos/sin scale, as the training stack's ``lerope_freqs``: the
    RoPE (+ YaRN) frequencies times ``exp(alpha)`` (transformers' Mistral Large 4 convention)."""
    inv, mscale = rope_inv_freq(head_dim, theta, yarn)
    return inv * torch.tensor(alpha, dtype=torch.float64, device="cpu").exp(), mscale


class LuauCoderRotaryEmbedding(nn.Module):
    """fp32 cos/sin tables (computed in fp64 like training), gathered by
    ``position_ids``. The tables are plain attributes, not buffers, so a
    ``model.to(bfloat16)`` never rounds them."""

    def __init__(self, config: LuauCoderConfig):
        super().__init__()
        theta, yarn = _rope_settings(config)
        alpha = getattr(config, "lerope_alpha", None)
        self.inv_freq, self.mscale = (lerope_inv_freq(config.head_dim, theta, alpha, yarn) if alpha
                                      else rope_inv_freq(config.head_dim, theta, yarn))
        self.min_len = config.max_position_embeddings
        self._cos: Tensor | None = None
        self._sin: Tensor | None = None

    def forward(self, position_ids: Tensor, needed: int) -> tuple[Tensor, Tensor]:
        dev = position_ids.device
        if self._cos is None or self._cos.shape[0] < needed or self._cos.device != dev:
            n = max(needed, self.min_len, 0 if self._cos is None else self._cos.shape[0])
            freqs = torch.outer(torch.arange(n, dtype=torch.float64, device="cpu"), self.inv_freq)
            self._cos = (freqs.cos() * self.mscale).float().to(dev)
            self._sin = (freqs.sin() * self.mscale).float().to(dev)
        return self._cos[position_ids], self._sin[position_ids]


# modules


class LuauCoderRMSNorm(nn.Module):
    """``learnable=False`` (q/k norm): the gain is a persistent buffer, fixed at 1 as in training
    (same key and shape, so checkpoints and the fused kernels are unchanged; not a parameter, so it
    is neither counted nor trained)."""

    def __init__(self, dim: int, eps: float, fuse: bool = True, learnable: bool = True):
        super().__init__()
        if learnable:
            self.weight = nn.Parameter(torch.ones(dim))
        else:
            self.register_buffer("weight", torch.ones(dim))
        self.eps, self.fuse = eps, fuse

    def forward(self, x: Tensor) -> Tensor:
        return _rms_norm(x, self.weight, self.eps, fuse=self.fuse)

    def extra_repr(self) -> str:
        return f"{tuple(self.weight.shape)}, eps={self.eps}"


def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
    B, H, T, d = x.shape
    if n_rep == 1:
        return x
    return x[:, :, None].expand(B, H, n_rep, T, d).reshape(B, H * n_rep, T, d)


def eager_attention_forward(module: nn.Module, query: Tensor, key: Tensor, value: Tensor,
                            attention_mask: Tensor | None, scaling: float, dropout: float = 0.0,
                            **kwargs) -> tuple[Tensor, Tensor]:
    key = repeat_kv(key, module.num_key_value_groups)
    value = repeat_kv(value, module.num_key_value_groups)
    scores = torch.matmul(query, key.transpose(2, 3)) * scaling
    if attention_mask is not None:
        if attention_mask.dtype == torch.bool:
            scores = scores.masked_fill(~attention_mask, torch.finfo(scores.dtype).min)
        else:
            scores = scores + attention_mask[..., :key.shape[-2]]
    probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype)
    probs = F.dropout(probs, p=dropout, training=module.training)
    return torch.matmul(probs, value).transpose(1, 2).contiguous(), probs


def _decode_attention(q: Tensor, k: Tensor, v: Tensor, mask: Tensor, scaling: float) -> Tensor:
    """One query token against a static cache, GQA-native: ``q`` ``(B, nq, 1, d)``,
    ``k`` / ``v`` ``(B, nkv, L, d)``, ``mask`` ``(B, 1, 1, L)`` bool -> ``(B, 1, nq, d)``.
    Each KV head is read once (the backends' masked path repeats K/V per query head,
    which dominates a decode step at a long cache); fp32 scores and softmax."""
    B, nq, _, d = q.shape
    nkv = k.shape[1]
    s = torch.matmul(q.reshape(B, nkv, nq // nkv, d).float(), k.float().transpose(-1, -2)) * scaling
    p = s.masked_fill(~mask, float("-inf")).softmax(-1)              # (B, nkv, G, L)
    return torch.matmul(p.to(v.dtype), v).reshape(B, 1, nq, d)


class LuauCoderAttention(nn.Module):
    """DiffAttn V2 + XSA. Query heads ``2i`` / ``2i+1`` are adjacent, so both
    fall in the same GQA group and share K and V."""

    def __init__(self, config: LuauCoderConfig, layer_idx: int):
        super().__init__()
        self.config, self.layer_idx = config, layer_idx
        self.nh, self.nkv, self.hd = config.num_attention_heads, config.num_key_value_heads, config.head_dim
        self.nq = config.num_query_heads
        # VHA: qkv makes nxq physical q heads; query head h reads physical h % nxq
        self.nxq = getattr(config, "num_q_latent_heads", self.nq)
        self.diff, self.xsa = config.diff_v2, config.xsa
        self.num_key_value_groups = self.nq // self.nkv
        self.scaling = self.hd ** -0.5
        self.is_causal = True
        self.eps = config.rms_norm_eps
        self.fuse = config.compile_elementwise
        self.rope = layer_idx not in (getattr(config, "nope_layers", None) or ())
        self.iha = config.iha_layer(layer_idx)
        D = config.hidden_size
        # output gate: per channel, its W_g rows follow q|k|v in the fused qkv weight;
        # per head (LFM3), a separate dim -> nh gate_proj
        self.g_rows = self.nh * self.hd if config.output_gate == "elementwise" else 0
        self.qkv = nn.Linear(D, (self.nxq + 2 * self.nkv) * self.hd + self.g_rows, bias=False)
        self.lambda_proj = nn.Linear(D, self.nh, bias=False) if self.diff else None
        self.gate_proj = nn.Linear(D, self.nh, bias=False) if config.output_gate == "headwise" else None
        self.o_proj = nn.Linear(self.nh * self.hd, D, bias=False)
        # gains fixed at 1 like training (lm/model.py: RMSNorm(learnable=False))
        self.q_norm = LuauCoderRMSNorm(self.hd, self.eps, learnable=False) if config.qk_norm else None
        self.k_norm = LuauCoderRMSNorm(self.hd, self.eps, learnable=False) if config.qk_norm else None
        self.ccgqa = getattr(config, "ccgqa", False)
        if self.ccgqa:
            # q|k latents (q heads, then k heads): depthwise causal conv taps
            # conv_seq[c, i] and per-head dense taps conv_head[h, j, c_in, c_out], lag order
            H, ks, kc = self.nq + self.nkv, config.ccgqa_seq_kernel, config.ccgqa_head_kernel
            self.conv_taps = ks + kc - 1
            self.qk_mean, self.value_shift = config.ccgqa_qk_mean, config.ccgqa_value_shift
            self.conv_seq = nn.Parameter(torch.empty(H * self.hd, ks))
            self.conv_seq_bias = nn.Parameter(torch.zeros(H * self.hd))
            self.conv_head = nn.Parameter(torch.empty(H, kc, self.hd, self.hd))
            self.conv_head_bias = nn.Parameter(torch.zeros(H, self.hd))
            self.k_temp = nn.Parameter(torch.zeros(self.nkv))  # key scale exp(t), per KV head
        if self.iha:
            P = self.iha_p = config.iha_pseudo_heads
            self.iha_q = nn.Parameter(torch.empty(self.nq, self.nq, P))
            self.iha_k = nn.Parameter(torch.empty(self.nkv, self.nkv, P))
            self.iha_v = nn.Parameter(torch.empty(self.nkv, self.nkv, P))
            self.iha_r = nn.Parameter(torch.empty(self.nq, P))

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor, attention_mask, position_ids: Tensor,
                past_key_values: Cache | None = None, **kwargs) -> tuple[Tensor, Tensor | None]:
        B, T, _ = x.shape
        qkv, gate = self.qkv(x), None
        if self.g_rows:
            qkv, gate = qkv.split([qkv.shape[-1] - self.g_rows, self.g_rows], dim=-1)
        qkv = qkv.view(B, T, self.nxq + 2 * self.nkv, self.hd)
        conv_pos = kwargs.pop("conv_pos", None)
        iha_ctx = kwargs.pop("iha_ctx", None)
        rope = self.rope
        if not self.rope or self.iha:   # IHA layers rotate later, at the virtual positions
            cos = sin = None
        if self.ccgqa:
            qk, v = qkv.split([self.nxq + self.nkv, self.nkv], dim=2)
            q, k, v = self._ccgqa_qkv(qk, v, cos, sin, conv_pos, past_key_values)
        else:
            q, k, v = qkv.split([self.nxq, self.nkv, self.nkv], dim=2)
            q = self._virtual(q)
            qw = self.q_norm.weight if self.q_norm is not None else None
            kw = self.k_norm.weight if self.k_norm is not None else None
            q, k = _qk_rope(q, k, qw, kw, cos, sin, self.eps, fuse=self.fuse)
        if self.iha:
            y, vp = self._iha_attend(q, k, v, iha_ctx, rope, past_key_values)
            probs = None
        else:
            # (B, H, T, d) views: flash transposes them straight back, sdpa takes the strides
            q, k_all, v_all = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
            if past_key_values is not None:
                k_all, v_all = past_key_values.update(k_all, v_all, self.layer_idx)
            if T == 1 and getattr(past_key_values, "is_static", False):
                y, probs = _decode_attention(q, k_all, v_all, attention_mask, self.scaling), None
            else:
                impl = self.config._attn_implementation
                fn = eager_attention_forward if impl in (None, "eager") else ALL_ATTENTION_FUNCTIONS[impl]
                y, probs = fn(self, q, k_all, v_all, attention_mask, dropout=0.0, scaling=self.scaling,
                              position_ids=position_ids, **kwargs)       # (B, T, nq, d)
        if self.diff and self.gate_proj is not None:   # one GEMM for both nh-wide projections
            lam, gate = F.linear(x, torch.cat((self.lambda_proj.weight, self.gate_proj.weight))).split(self.nh, -1)
        else:
            lam = self.lambda_proj(x) if self.diff else None
            if self.gate_proj is not None:
                gate = self.gate_proj(x)
        y = y.view(B, T, self.nq, self.hd)
        if self.iha:
            y = _iha_diff_xsa(y, lam, vp, self.nh, self.nkv, self.xsa, gate)
        else:
            y = _diff_xsa(y, lam, v, self.nh, self.nkv, self.xsa, gate, fuse=self.fuse)
        return self.o_proj(y), probs

    def _iha_attend(self, q: Tensor, k: Tensor, v: Tensor, ctx: dict, rope: bool,
                    cache: Cache | None) -> tuple[Tensor, Tensor]:
        """Training's ``Attention._iha_attend`` on latent, unrotated q / k / v: pseudo-heads, QK-norm,
        temperature, RoPE at virtual positions (``ctx``: ``LuauCoderModel._iha_context``), then
        attention over the interleaved tokens, with the pseudo-keys in the KV cache as extra heads
        ``(B, nkv * P, T, d)``. Returns the collapsed ``(B, T, nq, d)`` and the pseudo-values
        ``(B, T, P, nkv, d)``."""
        B, T, _, hd = q.shape
        P, nq, nkv = self.iha_p, self.nq, self.nkv

        def mix(t: Tensor, alpha: Tensor) -> Tensor:
            return torch.einsum("btmd,mhp->btphd", t, alpha.to(t.dtype))

        qp = mix(q, self.iha_q).reshape(B, T * P, nq, hd)
        kp = mix(k, self.iha_k).reshape(B, T * P, nkv, hd)
        vp = mix(v, self.iha_v)
        if self.q_norm is not None:
            qp, kp = _rms(qp, self.q_norm.weight, self.eps), _rms(kp, self.k_norm.weight, self.eps)
        if self.ccgqa:
            kp = (kp.float() * self.k_temp.float().exp()[:, None]).to(kp.dtype)
        if rope:
            qp, kp = _rope(qp, ctx["cos"], ctx["sin"]), _rope(kp, ctx["cos"], ctx["sin"])
        keys = kp.view(B, T, P, nkv, hd).permute(0, 3, 2, 1, 4).reshape(B, nkv * P, T, hd)
        vals = vp.permute(0, 3, 2, 1, 4).reshape(B, nkv * P, T, hd)
        if cache is not None:
            keys, vals = cache.update(keys, vals, self.layer_idx)

        def interleave(t: Tensor) -> Tensor:   # (B, nkv * P, T', d) -> (B, nkv, Tk * P, d), the tokens in reach
            t = t[:, :, ctx["first"]:]
            return t.reshape(B, nkv, P, -1, hd).transpose(2, 3).reshape(B, nkv, -1, hd)

        y = _iha_attention(qp.transpose(1, 2), interleave(keys), interleave(vals), ctx["mask"], self.scaling)
        return torch.einsum("btphd,hp->bthd", y.view(B, T, P, nq, hd), self.iha_r.to(y.dtype)), vp

    def _virtual(self, q: Tensor) -> Tensor:
        """VHA: physical q heads ``(..., nxq, d)`` -> the ``nq`` virtual heads (``h % nxq``)."""
        return q if self.nxq == self.nq else q.repeat(1, 1, self.nq // self.nxq, 1)

    def _ccgqa_qkv(self, qk: Tensor, v: Tensor, cos: Tensor | None, sin: Tensor | None,
                   pos: Tensor | None, cache: Cache | None) -> tuple[Tensor, Tensor, Tensor]:
        """CCGQA ``q`` / ``k`` / ``v`` (``_ccgqa_fn``), carrying the conv and value-shift
        state in ``cache``. ``pos`` ``(B, T)``: each token's position in its document."""
        B, T = qk.shape[:2]
        if pos is None:
            pos = torch.arange(T, device=qk.device).expand(B, T)
        states = _conv_states(cache) if cache is not None else None
        prev = states.get(self.layer_idx) if states is not None else None
        tail = prev[0] if prev is not None else qk.new_zeros(B, self.conv_taps - 1, *qk.shape[2:])
        before = None
        if self.value_shift:
            before = prev[1] if prev is not None else qk.new_zeros(B, 1, self.nkv * self.hd // 2)
        W, b, q_gain, k_gain = self._ccgqa_consts()
        q, k, v, tail, last_v = _ccgqa(qk, tail, v, before, pos, W, b, q_gain, k_gain, cos, sin,
                                       self.nq, self.nkv, self.eps, self.qk_mean, fuse=self.fuse)
        if states is not None:
            states[self.layer_idx] = (tail, last_v)
        return q, k, v

    def _ccgqa_consts(self) -> tuple[Tensor, Tensor, Tensor, Tensor]:
        """fp32 composed conv ``W`` as ``(H, d*K, d)`` bmm operands (row ``c*K + j`` =
        input channel ``c`` at window slot ``j``, lag ``K-1-j``), its bias, the q gain
        and the k gain times ``exp(temp)``; cached without grad (``_cached``)."""
        params = (self.conv_seq, self.conv_seq_bias, self.conv_head, self.conv_head_bias,
                  self.q_norm.weight, self.k_norm.weight, self.k_temp)

        def build():
            W, b = compose_conv(*params[:4])
            H, K, hd, _ = W.shape
            W = W.flip(1).permute(0, 2, 1, 3).reshape(H, hd * K, hd)
            # IHA applies the temperature after its mix and norm
            kg = params[5].float() if self.iha else params[5].float() * params[6].float().exp()[:, None]
            return W, b, params[4].float(), kg

        return _cached(params, build)


class LuauCoderMLP(nn.Module):
    def __init__(self, config: LuauCoderConfig):
        super().__init__()
        self.hidden = config.intermediate_size
        self.fuse = config.compile_elementwise
        self.w13 = nn.Linear(config.hidden_size, 2 * self.hidden, bias=False)  # fused gate|up
        self.w2 = nn.Linear(self.hidden, config.hidden_size, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        return self.w2(_swiglu(self.w13(x), fuse=self.fuse))


class LuauCoderDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: LuauCoderConfig, layer_idx: int):
        super().__init__()
        D, eps = config.hidden_size, config.rms_norm_eps
        self.layer_idx = layer_idx
        self.attn_norm = LuauCoderRMSNorm(D, eps, config.compile_elementwise)
        self.attn = LuauCoderAttention(config, layer_idx)
        self.ffn_norm = LuauCoderRMSNorm(D, eps, config.compile_elementwise)
        self.ffn = LuauCoderMLP(config)
        self.attn_res = config.attn_res
        if self.attn_res:
            # one pseudo-query and key-norm gain per sublayer input; layer 0's
            # attention reads the embedding alone, so it has nothing to mix
            self.block_start = layer_idx % config.attn_res_block_size == 0
            self.attn_res_norm = nn.Parameter(torch.ones(D)) if layer_idx else None
            self.attn_res_query = nn.Parameter(torch.zeros(D)) if layer_idx else None
            self.mlp_res_norm = nn.Parameter(torch.ones(D))
            self.mlp_res_query = nn.Parameter(torch.zeros(D))

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor, attention_mask=None,
                position_ids: Tensor | None = None, past_key_values: Cache | None = None,
                attn_span: AttnResSpan | None = None, mlp_span: AttnResSpan | None = None,
                **kwargs):
        if self.attn_res:
            return self._forward_attn_res(x, attn_span, mlp_span, cos, sin, attention_mask,
                                          position_ids, past_key_values, **kwargs)
        a, _ = self.attn(self.attn_norm(x), cos, sin, attention_mask, position_ids, past_key_values,
                         **kwargs)
        h, n = _add_rms_norm(x, a, self.ffn_norm.weight, self.ffn_norm.eps, fuse=self.ffn_norm.fuse)
        return h + self.ffn(n)

    def _forward_attn_res(self, partial: Tensor, attn_span: AttnResSpan | None, mlp_span: AttnResSpan,
                          cos, sin, attention_mask, position_ids, past_key_values, **kwargs):
        """``partial`` ``(B, T, D)``: the running sum of the current block's
        sublayer outputs (at a block start: the finished block's, which
        ``mlp_span`` already holds as its newest source). Returns the new
        partial and this layer's attention input."""
        B, T, D = partial.shape
        N = B * T
        flat = partial.reshape(N, D)
        h = flat if attn_span is None else attn_span.mix(flat, self.attn_call)
        a, _ = self.attn(self.attn_norm(h.view(B, T, D)), cos, sin, attention_mask, position_ids,
                         past_key_values, **kwargs)
        a = a.reshape(N, D).to(flat.dtype)
        # each partial is read by exactly one mix and one residual add: the add's
        # gradient is summed inside that mix's backward kernel
        flat = a if self.block_start else attn_span.residual(flat, a, self.attn_call)
        m = mlp_span.mix(flat, self.mlp_call)
        f = self.ffn(self.ffn_norm(m.view(B, T, D))).reshape(N, D).to(flat.dtype)
        return mlp_span.residual(flat, f, self.mlp_call).view(B, T, D), h.view(B, T, D)


# models


class LuauCoderPreTrainedModel(PreTrainedModel):
    config_class = LuauCoderConfig
    config: LuauCoderConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["LuauCoderDecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn = True
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    _supports_flex_attn = True
    _supports_attention_backend = True
    _supports_cache_class = True
    _can_compile_fullgraph = False

    @torch.no_grad()
    def _init_weights(self, module: nn.Module) -> None:
        """The training stack's init: N(0, std) inputs, N(0, std / sqrt(2L))
        residual-branch outputs, lambda logits at half scale, unit norms,
        zero AttnRes pseudo-queries (uniform depth mix)."""
        cfg = self.config
        std = cfg.initializer_range
        out_std = std / math.sqrt(2 * cfg.num_hidden_layers)
        if isinstance(module, LuauCoderAttention):
            _init.normal_(module.qkv.weight, mean=0.0, std=std)
            _init.normal_(module.o_proj.weight, mean=0.0, std=out_std)
            if module.lambda_proj is not None:
                _init.normal_(module.lambda_proj.weight, mean=0.0, std=0.5 * std)
            if module.gate_proj is not None:
                _init.normal_(module.gate_proj.weight, mean=0.0, std=std)
            if module.ccgqa:  # ZAYA init: conv weights N(0, std), zero biases, unit key scale
                _init.normal_(module.conv_seq, mean=0.0, std=std)
                _init.normal_(module.conv_head, mean=0.0, std=std)
                for p in (module.conv_seq_bias, module.conv_head_bias, module.k_temp):
                    _init.zeros_(p)
            if module.iha:
                _iha_mix_init(module.iha_q, 2 if module.diff else 1)
                _iha_mix_init(module.iha_k, 1)
                _iha_mix_init(module.iha_v, 1)
                _init.constant_(module.iha_r, 1.0 / module.iha_p)
        elif isinstance(module, LuauCoderMLP):
            _init.normal_(module.w13.weight, mean=0.0, std=std)
            _init.normal_(module.w2.weight, mean=0.0, std=out_std)
        elif isinstance(module, LuauCoderRMSNorm):
            _init.ones_(module.weight)
        elif isinstance(module, nn.Embedding):
            _init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, LuauCoderForCausalLM):
            if not cfg.tie_word_embeddings:
                _init.normal_(module.lm_head.weight, mean=0.0, std=std)
        for name, p in module.named_parameters(recurse=False):
            if name.endswith("_res_norm"):
                _init.ones_(p)
            elif name.endswith("_res_query"):
                _init.zeros_(p)

    @classmethod
    def _supports_default_dynamic_cache(cls) -> bool:
        # generate() then leaves the cache to forward(), which makes a LuauCoderCache
        # (so CCGQA's conv state follows beam reordering / batch expansion)
        return False

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kwargs):
        # AttnRes sources reach the layers as keyword arguments: reentrant
        # checkpointing would silently drop their gradients
        gc_kwargs = dict(gradient_checkpointing_kwargs or {})
        gc_kwargs["use_reentrant"] = False
        return super().gradient_checkpointing_enable(gradient_checkpointing_kwargs=gc_kwargs, **kwargs)


_MASK_PARAMS = set(inspect.signature(create_causal_mask).parameters)


def _causal_mask(config, embeds: Tensor, attention_mask, past_key_values, position_ids: Tensor,
                 past_len: int):
    """``create_causal_mask`` across the transformers v4 / v5 signatures."""
    kw = dict(config=config, attention_mask=attention_mask, past_key_values=past_key_values,
              position_ids=position_ids)
    kw["inputs_embeds" if "inputs_embeds" in _MASK_PARAMS else "input_embeds"] = embeds
    if "cache_position" in _MASK_PARAMS:
        kw["cache_position"] = torch.arange(past_len, past_len + embeds.shape[1], device=embeds.device)
    return create_causal_mask(**kw)


def _new_cache(config) -> LuauCoderCache:
    try:
        return LuauCoderCache(config=config)
    except TypeError:  # older transformers
        return LuauCoderCache()


class LuauCoderModel(LuauCoderPreTrainedModel):
    def __init__(self, config: LuauCoderConfig):
        super().__init__(config)
        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList(LuauCoderDecoderLayer(config, i) for i in range(config.num_hidden_layers))
        self.norm = LuauCoderRMSNorm(config.hidden_size, config.rms_norm_eps, config.compile_elementwise)
        self.rotary_emb = LuauCoderRotaryEmbedding(config)
        if config.attn_res:
            if not config.attn_res_delta_block:  # final mix over every block before the output norm
                self.out_res_norm = nn.Parameter(torch.ones(config.hidden_size))
                self.out_res_query = nn.Parameter(torch.zeros(config.hidden_size))
            self._plan_attn_res()
        self.gradient_checkpointing = False
        self.post_init()

    def _plan_attn_res(self) -> None:
        """Which block's sources each mix reads, and its call index there.
        Block ``m`` (layers ``mS .. mS+S-1``) is read by its layers' MLP mixes
        (call ``2j``), their attention mixes (``2j-1``, ``j >= 1``) and the
        next block's first attention mix, which runs before that block is
        pushed (``2S-1``); the last block's final slot is the output mix, which
        Delta Block doesn't have. Sites are stored as ``(layer, kind)`` and resolved to parameters at
        every forward (loading can replace the parameter objects).
        """
        S, L = self.config.attn_res_block_size, self.config.num_hidden_layers
        n_spans = (L + S - 1) // S
        self._span_sites: list[list] = [[None] * (2 * min(S, L - m * S)) for m in range(n_spans)]
        for i, blk in enumerate(self.blocks):
            m, j = divmod(i, S)
            blk.mlp_call = 2 * j
            self._span_sites[m][2 * j] = (i, "mlp")
            if i:
                span, call = (m - 1, 2 * S - 1) if j == 0 else (m, 2 * j - 1)
                blk.attn_span_index, blk.attn_call = span, call
                self._span_sites[span][call] = (i, "attn")
        self._out_call = len(self._span_sites[-1]) - 1
        if self.config.attn_res_delta_block:
            self._span_sites[-1].pop()
        else:
            self._span_sites[-1][self._out_call] = (-1, "out")

    def _site(self, layer: int, kind: str) -> tuple[Tensor, Tensor]:
        if kind == "out":
            return self.out_res_norm, self.out_res_query
        blk = self.blocks[layer]
        if kind == "mlp":
            return blk.mlp_res_norm, blk.mlp_res_query
        return blk.attn_res_norm, blk.attn_res_query

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embed

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.embed = value

    def _iha_context(self, position_ids: Tensor, pos: Tensor, attention_mask: Tensor | None,
                     past_len: int) -> dict:
        """What every IHA layer of this forward shares: cos / sin ``(B|1, T * P, d/2)`` at virtual positions
        ``position * P + p``, the first real token still in reach and the mask ``(B|1, 1, 1, T * P, Tk * P)``
        over row indices: causal in virtual order, the window, padding and packed documents (``pos``
        restarting). A row with no valid key (left padding) attends to itself: its softmax stays finite."""
        cfg = self.config
        P, W = cfg.iha_pseudo_heads, cfg.iha_window
        if attention_mask is not None and attention_mask.dim() != 2:
            raise NotImplementedError("IHA layers take a 2-D attention_mask")
        dev, B, T = pos.device, pos.shape[0], pos.shape[1]
        vpos = (position_ids.long().unsqueeze(-1) * P + torch.arange(P, device=dev)).flatten(1)
        needed = (past_len + T) * P
        if not torch.compiler.is_compiling():
            needed = max(needed, int(vpos.max()) + 1)
        cos, sin = self.rotary_emb(vpos, needed)
        first = max(0, past_len - W) if W else 0
        vq = torch.arange(past_len * P, (past_len + T) * P, device=dev)
        vk = torch.arange(first * P, (past_len + T) * P, device=dev)
        d = vq[:, None] - vk[None]
        mask = d >= 0
        if W:
            mask = mask & (d < W * P)
        mask = mask[None]
        if attention_mask is not None:
            keep = attention_mask.bool()[:, -(past_len + T - first):].repeat_interleave(P, 1)
            mask = mask & keep[:, None, :]
        if past_len == 0:
            doc = (pos == 0).long().cumsum(1).repeat_interleave(P, 1)
            mask = mask & (doc[:, :, None] == doc[:, None, :])
        mask = mask | (d == 0)[None]
        return {"cos": cos, "sin": sin, "first": first, "mask": mask[:, None, None]}

    def forward(self, input_ids: Tensor | None = None, attention_mask: Tensor | None = None,
                position_ids: Tensor | None = None, past_key_values: Cache | None = None,
                inputs_embeds: Tensor | None = None, use_cache: bool | None = None,
                output_hidden_states: bool | None = None, return_dict: bool | None = None,
                **kwargs) -> BaseModelOutputWithPast:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("pass exactly one of input_ids and inputs_embeds")
        kwargs.pop("cache_position", None)
        kwargs.pop("output_attentions", None)
        output_hidden_states = (output_hidden_states if output_hidden_states is not None
                                else self.config.output_hidden_states)
        if use_cache is None:  # no cache while training: it only costs memory and disables
            use_cache = self.config.use_cache and not self.training  # packed-sequence masks
        if self.gradient_checkpointing and self.training:
            use_cache = False
        if inputs_embeds is None:
            inputs_embeds = self.embed(input_ids)
        if use_cache and past_key_values is None:
            past_key_values = _new_cache(self.config)
        if not use_cache:
            past_key_values = None
        B, T, D = inputs_embeds.shape
        if getattr(past_key_values, "is_static", False):
            # no host syncs (CUDA-graph capture): the cache holds the slots and valid keys
            if position_ids is None or attention_mask is not None:
                raise ValueError("a LuauCoderStaticCache takes position_ids and no attention_mask")
            if self.config.iha:
                raise NotImplementedError("IHA layers don't run on the static cache (no CUDA-graph decoding)")
            cos, sin = self.rotary_emb(position_ids, past_key_values.max_len)
            mask = past_key_values.mask()
        else:
            past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
            needed = past_len + T
            if position_ids is None:
                position_ids = torch.arange(past_len, past_len + T, device=inputs_embeds.device).unsqueeze(0)
            elif not torch.compiler.is_compiling():
                needed = max(needed, int(position_ids.max()) + 1)
            cos, sin = self.rotary_emb(position_ids, needed)
            mask = _causal_mask(self.config, inputs_embeds, attention_mask, past_key_values,
                                position_ids, past_len)
        layer_kw = dict(cos=cos, sin=sin, attention_mask=mask, position_ids=position_ids,
                        past_key_values=past_key_values, **kwargs)
        if self.config.ccgqa or self.config.iha:
            # document positions for the conv mask: position_ids restart per packed document,
            # and a 2-D padding mask starts each row at its first real token
            pos = position_ids.long().expand(B, T)
            if attention_mask is not None and attention_mask.dim() == 2:
                pos = torch.minimum(pos, (attention_mask.long().cumsum(-1) - 1)[:, -T:])
            layer_kw["conv_pos"] = pos
        if self.config.iha:
            layer_kw["iha_ctx"] = self._iha_context(position_ids, pos, attention_mask, past_len)
        hidden: list[Tensor] = []

        if self.config.attn_res:
            x = inputs_embeds
            if torch.is_autocast_enabled(x.device.type):  # bf16 depth stream, like the sublayers
                x = x.to(torch.get_autocast_dtype(x.device.type))
            partial = x
            sources = x.reshape(B * T, 1, D)[:, :0]
            spans: list[AttnResSpan] = []
            fused, eps = self.config.fused_kernels, self.config.rms_norm_eps
            delta = self.config.attn_res_delta_block
            route_x = self.config.attn_res_partial_source or not delta
            for i, blk in enumerate(self.blocks):
                if blk.block_start:  # the finished block (or the embedding) becomes a source
                    sources = torch.cat((sources, partial.reshape(B * T, 1, D)), dim=1)
                    sites = [self._site(*s) for s in self._span_sites[len(spans)]]
                    spans.append(AttnResSpan(sources, sites, eps, fused, delta, route_x))
                attn_span = spans[blk.attn_span_index] if i else None
                partial, h = blk(partial, attn_span=attn_span, mlp_span=spans[-1], **layer_kw)
                if output_hidden_states:
                    hidden.append(h)
            flat = partial.reshape(B * T, D)
            out = (spans[-1].stream(flat) if delta else spans[-1].mix(flat, self._out_call)).view(B, T, D)
        else:
            out = inputs_embeds
            for blk in self.blocks:
                if output_hidden_states:
                    hidden.append(out)
                out = blk(out, **layer_kw)
        out = self.norm(out)
        if output_hidden_states:
            hidden.append(out)
        return BaseModelOutputWithPast(last_hidden_state=out, past_key_values=past_key_values,
                                       hidden_states=tuple(hidden) if output_hidden_states else None)


# fused generate()

# generation-config fields the fused decoder doesn't implement, with their neutral values
_FUSED_NEUTRAL = {
    "num_beams": (None, 1), "num_beam_groups": (None, 1), "diversity_penalty": (None, 0.0), "penalty_alpha": (None,),
    "dola_layers": (None,), "constraints": (None,), "force_words_ids": (None,), "prompt_lookup_num_tokens": (None,),
    "typical_p": (None, 1.0), "epsilon_cutoff": (None, 0.0), "eta_cutoff": (None, 0.0), "top_h": (None,),
    "no_repeat_ngram_size": (None, 0), "encoder_no_repeat_ngram_size": (None, 0),
    "encoder_repetition_penalty": (None, 1.0), "bad_words_ids": (None,), "suppress_tokens": (None,),
    "begin_suppress_tokens": (None,), "sequence_bias": (None,), "forced_bos_token_id": (None,),
    "forced_eos_token_id": (None,), "exponential_decay_length_penalty": (None,), "min_length": (None, 0),
    "min_new_tokens": (None, 0), "guidance_scale": (None, 1.0), "watermarking_config": (None,),
    "output_scores": (None, False), "output_logits": (None, False), "output_attentions": (None, False),
    "output_hidden_states": (None, False), "return_dict_in_generate": (None, False), "token_healing": (None, False),
    "max_time": (None,), "use_mtp": (None, False),
}
# generate() keyword arguments (model kwargs) the fused path accepts; anything else falls back.
# (``use_cache`` / ``generation_mode`` are accepted and ignored: the fused path always caches.)
_FUSED_MODEL_KWARGS = {"input_ids", "attention_mask", "use_cache", "tokenizer", "generation_mode"}


def _fused_fallback_reason(model, inputs, gc, model_kwargs, logits_processor, stopping_criteria,
                           extra: dict) -> str | None:
    """None if ``generate()`` can run on the fused decoder with these settings, else why not."""
    if not (HAVE_TRITON and model.device.type == "cuda") or model.training:
        return "needs CUDA + triton, eval mode"
    if model.dtype != torch.bfloat16:   # its GEMMs round like bf16 ones: an fp32 model keeps its own precision
        return f"a {model.dtype} model (the fused decoder runs bf16 models)"
    used = [k for k, v in extra.items() if v is not None and v is not False]
    if used:
        return f"unsupported generate() arguments {used}"
    if logits_processor or stopping_criteria:
        return "custom logits processors / stopping criteria"
    if set(model_kwargs) - _FUSED_MODEL_KWARGS:
        return f"model kwargs {sorted(set(model_kwargs) - _FUSED_MODEL_KWARGS)}"
    for name, ok in _FUSED_NEUTRAL.items():
        if getattr(gc, name, None) not in ok:
            return f"generation_config.{name}={getattr(gc, name)!r}"
    if not isinstance(inputs, Tensor) or inputs.dim() != 2 or inputs.is_floating_point():
        return "needs input_ids (B, T)"
    eos = gc.eos_token_id
    if isinstance(eos, (list, tuple)) and len(set(eos)) > 1:
        return "several EOS tokens"
    if gc.stop_strings and model_kwargs.get("tokenizer") is None:
        return "stop_strings without a tokenizer"
    mask = model_kwargs.get("attention_mask")
    if mask is not None and mask.dim() == 2 and bool((mask[:, 1:] < mask[:, :-1]).any()):
        return "not left-padded"
    from .luau_infer import supported
    return supported(model.config)


class _FusedGenerationMixin:
    """``generate()`` through the fused decoding engine (``luau_infer``: ~9 Triton kernels per
    layer, each step one CUDA graph, prompt-lookup speculation with the same output) whenever
    every setting is one it reproduces exactly: greedy or sampling (temperature, top-k, top-p,
    min-p), repetition penalty, ``max_new_tokens`` / ``max_length``, EOS / pad, stop strings,
    ``num_return_sequences``, left-padded batches and a streamer. Anything else (beam search,
    assisted / contrastive decoding, custom logits processors or stopping criteria, scores /
    dict outputs, ...) is transformers' own ``generate()``.
    It runs bf16 models, on the weights as loaded (no quantisation). Extra keyword
    arguments: ``fused_generate=False`` always uses transformers' path; ``fused_spec`` the
    speculative draft length (default 6, used while it pays; 0 = off)."""

    def generate(self, inputs=None, generation_config=None, logits_processor=None, stopping_criteria=None,
                 prefix_allowed_tokens_fn=None, synced_gpus=None, assistant_model=None, streamer=None,
                 negative_prompt_ids=None, negative_prompt_attention_mask=None, custom_generate=None, **kwargs):
        fused = kwargs.pop("fused_generate", True)
        spec = int(kwargs.pop("fused_spec", 6))
        hf = dict(prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, synced_gpus=synced_gpus,
                  assistant_model=assistant_model, negative_prompt_ids=negative_prompt_ids,
                  negative_prompt_attention_mask=negative_prompt_attention_mask, custom_generate=custom_generate)
        if fused:
            out = self._fused_generate(inputs, generation_config, logits_processor, stopping_criteria, streamer,
                                       spec, hf, kwargs)
            if out is not None:
                return out
        return super().generate(inputs, generation_config, logits_processor, stopping_criteria, streamer=streamer,
                                **hf, **kwargs)

    def _fused_generate(self, inputs, generation_config, logits_processor, stopping_criteria, streamer, spec,
                        hf: dict, kwargs: dict):
        """The fused decoder's output, or None when these settings need transformers' loop."""
        try:
            gc, model_kwargs = self._prepare_generation_config(copy.deepcopy(generation_config), **kwargs)
        except Exception:   # another transformers version's signature: its own generate() decides
            return None
        if inputs is None:
            inputs = model_kwargs.get("input_ids")
        why = _fused_fallback_reason(self, inputs, gc, model_kwargs, logits_processor, stopping_criteria, hf)
        self._fused_last_reason = why   # why the last generate() fell back (None: it ran fused)
        if why is not None:
            return None
        eos = gc.eos_token_id
        if isinstance(eos, (list, tuple)):
            eos = eos[0] if eos else None
        pad = gc.pad_token_id if gc.pad_token_id is not None else eos
        if isinstance(pad, (list, tuple)):
            pad = pad[0]
        mask = model_kwargs.get("attention_mask")
        if mask is None:   # as transformers infers it: pad tokens (when pad != eos) are padding
            mask = (inputs != pad).long() if pad is not None and pad != eos and bool((inputs == pad).any()) \
                else torch.ones_like(inputs)
        n_new = gc.max_new_tokens if gc.max_new_tokens is not None else (gc.max_length or 20) - inputs.shape[1]
        stop = gc.stop_strings or []
        stop = [stop] if isinstance(stop, str) else list(stop)
        decoder = self._fused_decoder(spec)
        with torch.inference_mode():
            return decoder.decode(
                inputs.to(self.device), mask.to(self.device), max_new_tokens=max(0, int(n_new)),
                sample=bool(gc.do_sample), temperature=gc.temperature or 1.0, top_k=gc.top_k or 0,
                top_p=1.0 if gc.top_p is None else gc.top_p, min_p=gc.min_p or 0.0,
                repetition_penalty=gc.repetition_penalty or 1.0, num_samples=gc.num_return_sequences or 1,
                eos=eos, pad=pad, stop=stop, tokenizer=model_kwargs.get("tokenizer"), streamer=streamer)

    def _fused_decoder(self, spec: int):
        """The decoder of these weights, kept between calls (rebuilt when a parameter changed)."""
        from .luau_infer import FusedDecoder
        try:
            key = (spec,) + tuple((p.data_ptr(), p._version) for p in self.parameters())
        except RuntimeError:   # inference tensors have no version counter
            key = (spec,) + tuple(p.data_ptr() for p in self.parameters())
        hit = self.__dict__.get("_fused")
        if hit is None or hit[0] != key:
            self.__dict__["_fused"] = None   # free the old engine's caches first
            with torch.inference_mode():
                self.__dict__["_fused"] = (key, FusedDecoder(self, None, spec))
        return self.__dict__["_fused"][1]


class LuauCoderForCausalLM(LuauCoderPreTrainedModel, _FusedGenerationMixin, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed.weight"} if _V5 else ["lm_head.weight"]

    def __init__(self, config: LuauCoderConfig):
        super().__init__(config)
        self.model = LuauCoderModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.model.embed = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_output_embeddings(self, value: nn.Linear) -> None:
        self.lm_head = value

    def forward(self, input_ids: Tensor | None = None, attention_mask: Tensor | None = None,
                position_ids: Tensor | None = None, past_key_values: Cache | None = None,
                inputs_embeds: Tensor | None = None, labels: Tensor | None = None,
                use_cache: bool | None = None, logits_to_keep: int | Tensor = 0,
                output_hidden_states: bool | None = None, return_dict: bool | None = None,
                **kwargs) -> CausalLMOutputWithPast:
        """``labels`` follow the HF convention (unshifted, ``-100`` ignored);
        ``shift_labels`` (already shifted) takes precedence. With labels and
        grad enabled the fused LM-head + CE kernel runs and ``logits`` is None."""
        num_items_in_batch = kwargs.pop("num_items_in_batch", None)
        shift_labels = kwargs.pop("shift_labels", None)
        out = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
                         past_key_values=past_key_values, inputs_embeds=inputs_embeds,
                         use_cache=use_cache, output_hidden_states=output_hidden_states, **kwargs)
        h = out.last_hidden_state
        if labels is not None and shift_labels is None:
            shift_labels = F.pad(labels[:, 1:], (0, 1), value=IGNORE_INDEX)
        loss = logits = None
        fused = (shift_labels is not None and torch.is_grad_enabled() and h.is_cuda and HAVE_TRITON
                 and self.config.fused_kernels)
        if fused:
            if torch.is_autocast_enabled(h.device.type):
                h = h.to(torch.get_autocast_dtype(h.device.type))
            loss = fused_linear_cross_entropy(h, self.lm_head.weight, shift_labels.to(h.device),
                                              num_items_in_batch, self.config.loss_chunk_size,
                                              self.config.final_logit_softcapping)
        else:
            sl = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
            logits = softcap_logits(self.lm_head(h if shift_labels is not None else h[:, sl, :]),
                                    self.config.final_logit_softcapping)
            if shift_labels is not None:
                y = shift_labels.to(logits.device).reshape(-1)
                nll = F.cross_entropy(logits.float().reshape(-1, logits.shape[-1]), y,
                                      ignore_index=IGNORE_INDEX, reduction="sum")
                loss = nll * _loss_scale(y, num_items_in_batch)
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=out.past_key_values,
                                      hidden_states=out.hidden_states)