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 , p = softmax(s) # out = sum_i p_i v_i # # Backward, with g = dL/dout: # # c_i = p_i ( - sum_j p_j ) # dv_i = p_i g + c_i r_i w - (c_i 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: 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)