"""Cortex-A: the production language-model architecture (flax.nnx), TPU-optimized. ONE configurable trunk (ModelConfig flags) implementing the full stack: - GQA attention with RoPE (+ dynamic YaRN) and optional QK-norm - Differential Attention (2410.05258): two-softmax noise cancellation - CLA-N: K/V shared across every ``kv_share_group`` consecutive layers - growing causal KV-memory (gated depth-axis read; O(1) state) - sliding-window attention + always-visible global sinks (StreamingLLM) - Block Attention Residuals (Kimi 2603.15031), SwiGLU FFN - Mixture-of-Depths FFN gate (optional): top-k token routing through the FFN at train/prefill + a causal predictor head for per-token decode (see _ffn_block) - tied embeddings with an optional factorized (learned d x d) output head - Multi-Token Prediction (training-only auxiliary; shared recurrent head) TPU optimizations (semantics-preserving; verified numerically equivalent): - FUSED SwiGLU: gate||up packed into one wide MatMul (split after). - FUSED K/V: wk||wv packed into one MatMul per CLA group. - FLASH attention via jax.nn.dot_product_attention -- causal/sliding-window/sink handled natively (is_causal skips the upper-triangle tiles); GQA via head broadcast (no K/V repeat); DiffAttn = (dpa(q1,k1,v) - lam*dpa(q2,k2,v)) over the two value halves -> never materializes the [B,H,T,T] score matrix. - RoPE tables are built at trace time (T static) -> baked compile-time constants, sliced per position; no per-step trig. - CHUNKED loss: forward_train_hidden() returns pre-head hidden so the head+CE are fused chunk-wise in losses.chunked_lm_loss (no [B,T,vocab] logit array). Flags off (use_diff_attn / use_attn_res / use_growing_memory / use_mtp / use_factorized_head off, kv_share_group=1, sliding_window=None) recovers a vanilla RoPE+SwiGLU+GQA transformer. Require ``layers_per_block % kv_share_group == 0``. """ from __future__ import annotations import math import numpy as np import jax import jax.numpy as jnp from flax import nnx from .config import ModelConfig F32 = jnp.float32 # ---------------------------------------------------------------------------- # RoPE (+ dynamic YaRN). Built with numpy at trace time (T is static) so the YaRN # ramp stays host-side and the cos/sin tables become baked compile-time constants # (no per-step trig; the forward just slices them by position). # ---------------------------------------------------------------------------- def compute_rope(T: int, head_dim: int, base: float, cfg: ModelConfig, dtype) -> tuple[jax.Array, jax.Array]: half = head_dim // 2 inv_freq = 1.0 / (base ** (np.arange(0, head_dim, 2, dtype=np.float64) / head_dim)) attn_factor = 1.0 if cfg.use_yarn: s = max(1.0, float(T) / float(cfg.yarn_orig_context)) if s > 1.0: def corr_dim(num_rot): return head_dim * np.log(cfg.yarn_orig_context / (num_rot * 2 * np.pi)) / (2 * np.log(base)) low = max(np.floor(corr_dim(cfg.yarn_beta_fast)), 0.0) high = min(np.ceil(corr_dim(cfg.yarn_beta_slow)), half - 1.0) denom = max(high - low, 1e-3) ramp = np.clip((np.arange(half) - low) / denom, 0.0, 1.0) extrap = 1.0 - ramp inv_freq = inv_freq * extrap + (inv_freq / s) * (1.0 - extrap) attn_factor = 0.1 * np.log(s) + 1.0 ang = np.arange(T)[:, None] * inv_freq[None, :] cos = np.concatenate([np.cos(ang), np.cos(ang)], axis=-1) * attn_factor sin = np.concatenate([np.sin(ang), np.sin(ang)], axis=-1) * attn_factor return jnp.asarray(cos, dtype), jnp.asarray(sin, dtype) def apply_rope(x: jax.Array, cos: jax.Array, sin: jax.Array) -> jax.Array: half = x.shape[-1] // 2 x1, x2 = x[..., :half], x[..., half:] rot = jnp.concatenate([-x2, x1], axis=-1) return x * cos[None, :, None, :] + rot * sin[None, :, None, :] def causal_bias(T: int, window: int | None, n_sink: int = 0) -> jax.Array: """Additive [T,T] mask (kept for the growing-memory path / reference).""" i = jnp.arange(T)[:, None] j = jnp.arange(T)[None, :] allowed = j <= i if window is not None: local = j > i - window if n_sink: local = local | (j < n_sink) allowed = allowed & local return jnp.where(allowed, 0.0, -1e30) def _sw_mask(T: int, window: int, n_sink: int) -> jax.Array: """Boolean [1,1,T,T] mask: causal AND (within window OR a global sink).""" i = jnp.arange(T)[:, None] j = jnp.arange(T)[None, :] allowed = (j <= i) & ((j > i - window) | (j < n_sink)) return allowed[None, None] # ---------------------------------------------------------------------------- # helpers # ---------------------------------------------------------------------------- def _linear(din: int, dout: int, cfg: ModelConfig, rngs: nnx.Rngs, scale: float = 1.0) -> nnx.Linear: return nnx.Linear( din, dout, use_bias=False, kernel_init=nnx.initializers.normal(stddev=0.02 * scale), dtype=cfg.compute_dtype, param_dtype=cfg.param_dtype, rngs=rngs, ) def _rmsnorm(dim: int, cfg: ModelConfig, rngs: nnx.Rngs) -> nnx.RMSNorm: return nnx.RMSNorm(dim, epsilon=cfg.norm_eps, dtype=cfg.compute_dtype, param_dtype=cfg.param_dtype, rngs=rngs) def _identity_init(key, shape, dtype=jnp.float32): return jnp.eye(shape[0], shape[1], dtype=dtype) def _causal_cummean(v): # [B,T,H,d] -> causal running mean over T T = v.shape[1] return jnp.cumsum(v, axis=1) / jnp.arange(1, T + 1, dtype=v.dtype).reshape(1, T, 1, 1) # ---------------------------------------------------------------------------- # Fused SwiGLU MLP: gate||up packed into a single wide MatMul. # ---------------------------------------------------------------------------- class MLP(nnx.Module): def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs): resid = 1.0 / math.sqrt(2 * cfg.n_layers) self.d_ff = cfg.d_ff self.gate_up = _linear(cfg.d_model, 2 * cfg.d_ff, cfg, rngs) # [gate | up] self.down = _linear(cfg.d_ff, cfg.d_model, cfg, rngs, scale=resid) def __call__(self, x, sink=None): gu = self.gate_up(x) g, u = gu[..., :self.d_ff], gu[..., self.d_ff:] a = jax.nn.silu(g) * u if sink is not None: # dormant-neuron probe: mean |activation| per unit sink.append(jnp.abs(a).mean(axis=(0, 1))) # [d_ff] return self.down(a) # ---------------------------------------------------------------------------- # Shared K/V for one CLA group (fused wk||wv), + memory write. Returns UN-repeated # K/V (dpa broadcasts kv-heads -> q-heads, so no GQA repeat is materialized). # ---------------------------------------------------------------------------- class _BlockKV(nnx.Module): def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs, use_mem: bool): self.cfg = cfg self.use_mem = use_mem self.diff = cfg.use_diff_attn hd = cfg.head_dim self.hkv = cfg.diff_kv_heads if self.diff else cfg.n_kv_heads self.vd = 2 * hd if self.diff else hd # value width per head self.k_out = self.hkv * (2 * hd if self.diff else hd) self.v_out = self.hkv * self.vd self.wkv = _linear(cfg.d_model, self.k_out + self.v_out, cfg, rngs) # fused K||V if cfg.use_qk_norm: self.k_norm = _rmsnorm(hd, cfg, rngs) if use_mem: self.wmk = _linear(self.hkv * self.vd, self.hkv * hd, cfg, rngs) self.learned_sink = cfg.learned_sink if self.learned_sink: # learnable, RoPE-free sink K/V (per group) init = nnx.initializers.normal(stddev=0.02); ns = self.learned_sink if self.diff: self.sink_k1 = nnx.Param(init(rngs.params(), (ns, self.hkv, hd), cfg.param_dtype)) self.sink_k2 = nnx.Param(init(rngs.params(), (ns, self.hkv, hd), cfg.param_dtype)) self.sink_v = nnx.Param(init(rngs.params(), (ns, self.hkv, 2 * hd), cfg.param_dtype)) else: self.sink_k = nnx.Param(init(rngs.params(), (ns, self.hkv, hd), cfg.param_dtype)) self.sink_v = nnx.Param(init(rngs.params(), (ns, self.hkv, hd), cfg.param_dtype)) def with_sinks(self, out_kv, B): """Append the learned (RoPE-free, always-attendable) sink K/V after the real K/V.""" if not self.learned_sink: return out_kv dt = out_kv[-1].dtype bc = lambda p: jnp.broadcast_to(p.value.astype(dt)[None], (B,) + p.value.shape) if self.diff: k1, k2, v = out_kv return (jnp.concatenate([k1, bc(self.sink_k1)], 1), jnp.concatenate([k2, bc(self.sink_k2)], 1), jnp.concatenate([v, bc(self.sink_v)], 1)) k, v = out_kv return (jnp.concatenate([k, bc(self.sink_k)], 1), jnp.concatenate([v, bc(self.sink_v)], 1)) def __call__(self, h, rope): cfg = self.cfg B, T, _ = h.shape cos, sin = rope hd = cfg.head_dim kv = self.wkv(h) kp, vp = kv[..., :self.k_out], kv[..., self.k_out:] if self.diff: k = kp.reshape(B, T, self.hkv, 2, hd) v = vp.reshape(B, T, self.hkv, 2 * hd) k1, k2 = k[:, :, :, 0], k[:, :, :, 1] if cfg.use_qk_norm: k1, k2 = self.k_norm(k1), self.k_norm(k2) k1, k2 = apply_rope(k1, cos, sin), apply_rope(k2, cos, sin) out_kv = (k1, k2, v) # un-repeated (hkv heads) v_raw = v else: k = kp.reshape(B, T, self.hkv, hd) v = vp.reshape(B, T, self.hkv, hd) if cfg.use_qk_norm: k = self.k_norm(k) k = apply_rope(k, cos, sin) out_kv = (k, v) v_raw = v new_mem = None if self.use_mem: mem_v = _causal_cummean(v_raw) # [B,T,hkv,vd], causal mk = self.wmk(mem_v.reshape(B, T, -1)).reshape(B, T, self.hkv, hd) new_mem = (mk, mem_v) return self.with_sinks(out_kv, B), new_mem # sinks appended after the real (RoPE'd) K/V def kv_decode(self, h, rope_p, mem_sum, pos): """Incremental decode: K/V for a single new token at position `pos` (matching __call__'s out_kv), plus the growing-memory updated in O(1) -- `mem_sum` is the running sum of raw V over past positions, so the causal cummean is mem_sum/(pos+1).""" cfg = self.cfg; B = h.shape[0]; hd = cfg.head_dim cos, sin = rope_p kv = self.wkv(h) kp, vp = kv[..., :self.k_out], kv[..., self.k_out:] if self.diff: k = kp.reshape(B, 1, self.hkv, 2, hd); v = vp.reshape(B, 1, self.hkv, 2 * hd) k1, k2 = k[:, :, :, 0], k[:, :, :, 1] if cfg.use_qk_norm: k1, k2 = self.k_norm(k1), self.k_norm(k2) k1, k2 = apply_rope(k1, cos, sin), apply_rope(k2, cos, sin) kv_p = (k1, k2, v); v_raw = v else: k = kp.reshape(B, 1, self.hkv, hd); v = vp.reshape(B, 1, self.hkv, hd) if cfg.use_qk_norm: k = self.k_norm(k) k = apply_rope(k, cos, sin) kv_p = (k, v); v_raw = v mem_p = None; mem_sum_new = mem_sum if self.use_mem: mem_sum_new = mem_sum + v_raw[:, 0].astype(F32) # [B,hkv,vd] mem_v = (mem_sum_new / jnp.asarray(pos + 1, F32)).astype(v_raw.dtype)[:, None] mk = self.wmk(mem_v.reshape(B, 1, -1)).reshape(B, 1, self.hkv, hd) mem_p = (mk, mem_v) return kv_p, mem_p, mem_sum_new # ---------------------------------------------------------------------------- # Per-layer sublayer: own Q + FLASH attention over the group's shared K/V, + gated # growing-memory read, + SwiGLU. Standard and Differential attention. # ---------------------------------------------------------------------------- class _Sublayer(nnx.Module): def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs, use_mem: bool, layer_idx: int): self.cfg = cfg self.use_mem = use_mem self.diff = cfg.use_diff_attn self.window = cfg.sliding_window self.n_sink = cfg.n_sink self.learned_sink = cfg.learned_sink self.scale = 1.0 / math.sqrt(cfg.head_dim) resid = 1.0 / math.sqrt(2 * cfg.n_layers) hd = cfg.head_dim self.norm1 = _rmsnorm(cfg.d_model, cfg, rngs) self.norm2 = _rmsnorm(cfg.d_model, cfg, rngs) self.mlp = MLP(cfg, rngs) self.use_ffn_gate = cfg.use_ffn_gate self.ffn_keep = cfg.ffn_keep if cfg.use_ffn_gate: # Mixture-of-Depths (FFN-only) self.ffn_router = _linear(cfg.d_model, 1, cfg, rngs) # routing score: top-k + weight self.ffn_predictor = _linear(cfg.d_model, 1, cfg, rngs) # causal head: mimics top-k at decode if cfg.use_qk_norm: self.q_norm = _rmsnorm(hd, cfg, rngs) if self.diff: self.hq = cfg.diff_q_heads self.wq = _linear(cfg.d_model, self.hq * 2 * hd, cfg, rngs) self.wo = _linear(self.hq * 2 * hd, cfg.d_model, cfg, rngs, scale=resid) self.head_norm = _rmsnorm(2 * hd, cfg, rngs) self.lambda_init = 0.8 - 0.6 * math.exp(-0.3 * layer_idx) init = nnx.initializers.normal(stddev=0.1) self.lq1 = nnx.Param(init(rngs.params(), (hd,), cfg.param_dtype)) self.lk1 = nnx.Param(init(rngs.params(), (hd,), cfg.param_dtype)) self.lq2 = nnx.Param(init(rngs.params(), (hd,), cfg.param_dtype)) self.lk2 = nnx.Param(init(rngs.params(), (hd,), cfg.param_dtype)) else: self.hq = cfg.n_q_heads self.wq = _linear(cfg.d_model, cfg.q_dim, cfg, rngs) self.wo = _linear(cfg.q_dim, cfg.d_model, cfg, rngs, scale=resid) if use_mem: self.mem_gate = nnx.Param(jnp.zeros((), cfg.param_dtype)) def _lam(self, dtype): l1 = jnp.exp(jnp.sum(self.lq1.value * self.lk1.value)) l2 = jnp.exp(jnp.sum(self.lq2.value * self.lk2.value)) return (l1 - l2 + self.lambda_init).astype(dtype) def _dpa(self, q, k, v, T): """Dense single-softmax attention (causal/window/sink) via dot_product_attention. q:[B,Tq,hq,hd] k:[B,Tk,hkv,hd] v:[B,Tk,hkv,vd] (dpa broadcasts hkv->hq).""" if self.learned_sink: # appended sinks -> explicit mask Tq, Tk = q.shape[1], k.shape[1]; real_k = Tk - self.learned_sink qpos, kpos = jnp.arange(Tq), jnp.arange(Tk) allowed = (kpos[None, :] <= qpos[:, None]) & (kpos[None, :] < real_k) if self.window is not None: allowed = allowed & (kpos[None, :] > qpos[:, None] - self.window) if self.n_sink: allowed = allowed | ((kpos[None, :] < self.n_sink) & (kpos[None, :] <= qpos[:, None])) allowed = allowed | (kpos[None, :] >= real_k) # learned sinks always attendable return jax.nn.dot_product_attention(q, k, v, mask=allowed[None, None]) if self.window is None and self.n_sink == 0: return jax.nn.dot_product_attention(q, k, v, is_causal=True) if self.n_sink == 0: return jax.nn.dot_product_attention(q, k, v, is_causal=True, local_window_size=(self.window - 1, 0)) return jax.nn.dot_product_attention(q, k, v, mask=_sw_mask(T, self.window, self.n_sink)) def _flash(self, q, k, v, T): """Memory-efficient attention: split queries into blocks of cfg.attn_block and scan, so the [B,H,T,T] score matrix is NEVER materialized -- each step holds only [B,H,blk,*] and it is recomputed (not stored) in the backward pass (checkpoint). Numerically identical to _dpa. Falls back to the dense path for short/non-divisible sequences. SWA fast path (window-cropped): when a sliding window is set, each query block can only see `window + blk` consecutive real keys, so we dynamic-slice exactly that span of K/V (+ the appended learned sinks) instead of masking the full key axis. The out-of-window key blocks are structurally REMOVED -- they vanish from the forward AND from every checkpoint recompute -- which is the tile-skipping the (broken on this jax/TPU) splash kernel was meant to provide, on the court-certified XLA path.""" blk = self.cfg.attn_block B, Tq, hq, hd = q.shape real_k = k.shape[1] - self.learned_sink # OUR fused flash kernel wins ONLY when out-of-window tiles dominate -- v5e-measured # 9.1x vs dense at the writer (seq 4095 / window 256), but a 2.4x LOSS at the short # planner (seq 819 / window 410, window ~ seq -> nothing to skip). Gate it on the same # "window small vs sequence" test the dense crop below uses, so only the writer path # (and any long-context SWA layer) takes it; the planner stays on dense attention. if (self.cfg.use_pallas_attn and v.shape[-1] == hd and self.n_sink == 0 and self.window is not None and (self.window + blk) <= 0.75 * real_k): from .flash_swa import flash_swa # fwd+dq/dkv at tile speed; Tq padded return flash_swa(q, k, v, window=self.window, n_sink=0, learned_sink=self.learned_sink) if blk <= 0 or Tq <= blk or Tq % blk != 0: return self._dpa(q, k, v, T) nb = Tq // blk Tk = k.shape[1] # T + learned_sink (sinks appended) real_k = Tk - self.learned_sink # real (causal) key count qblocks = q.reshape(B, nb, blk, hq, hd).swapaxes(0, 1) # [nb,B,blk,hq,hd] W = (self.window + blk) if self.window is not None else 0 # max real keys a block can see # Crop only when it removes a MATERIAL key fraction (>=25%): measured on v5e, the # dynamic-slice + unaligned-width overhead costs ~5% at seq 3072 / window 2048 # (W/real = 0.83) while the masked path tiles better; at Phase-2 geometry # (seq >> window) the crop removes 60%+ of the work and wins decisively. if self.window is not None and self.n_sink == 0 and W <= 0.75 * real_k: k_real, v_real = k[:, :real_k], v[:, :real_k] ks, vs = k[:, real_k:], v[:, real_k:] # learned sinks (may be empty) ls = self.learned_sink def body(_, xb): i, qb = xb # i:[] qb:[B,blk,hq,hd] start = jnp.clip((i + 1) * blk - W, 0, real_k - W) # window span for this block kw = jax.lax.dynamic_slice_in_dim(k_real, start, W, axis=1) vw = jax.lax.dynamic_slice_in_dim(v_real, start, W, axis=1) qabs = i * blk + jnp.arange(blk) kabs = start + jnp.arange(W) allowed = (kabs[None, :] <= qabs[:, None]) & (kabs[None, :] > qabs[:, None] - self.window) if ls: # sinks: appended, always attendable kw = jnp.concatenate([kw, ks], axis=1) vw = jnp.concatenate([vw, vs], axis=1) allowed = jnp.concatenate( [allowed, jnp.ones((blk, ls), bool)], axis=1) ob = jax.nn.dot_product_attention(qb, kw, vw, mask=allowed[None, None]) return _, ob else: kpos = jnp.arange(Tk) def body(_, xb): i, qb = xb # i:[] qb:[B,blk,hq,hd] qabs = i * blk + jnp.arange(blk) # absolute query positions allowed = (kpos[None, :] <= qabs[:, None]) & (kpos[None, :] < real_k) # causal, real keys if self.window is not None: allowed = allowed & (kpos[None, :] > qabs[:, None] - self.window) if self.n_sink: allowed = allowed | ((kpos[None, :] < self.n_sink) & (kpos[None, :] <= qabs[:, None])) if self.learned_sink: allowed = allowed | (kpos[None, :] >= real_k) # learned sinks: always attendable ob = jax.nn.dot_product_attention(qb, k, v, mask=allowed[None, None]) return _, ob _, outs = jax.lax.scan(jax.checkpoint(body), None, (jnp.arange(nb), qblocks)) return outs.swapaxes(0, 1).reshape(B, Tq, hq, v.shape[-1]) # [B,Tq,hq,vd] def _ffn_block(self, ff_in, mlp_sink, gate_sink, route): """SwiGLU FFN, optionally Mixture-of-Depths gated. Returns the residual contribution. route="train"/"prefill": fixed-capacity top-k over the sequence -> static shapes (TPU) and real FFN savings -- GATHER the k = keep*T kept tokens, run the FFN matmul on just those, SCATTER back; the rest skip to the residual. Kept tokens are scaled by the router weight sigmoid(s) so the main CE/MTP loss trains the router via the gradient path. In "train" we also train a causal predictor (BCE, stop-grad) to reproduce the top-k membership from each token alone. route="decode": the predictor decides PER TOKEN on its own representation -- so a single decode token needs no other positions (and a full sequence works as a batched proxy). For one token (B=T=1) we actually skip the FFN via lax.cond.""" if not self.use_ffn_gate: return self.mlp(ff_in, sink=mlp_sink) B, Tloc, _ = ff_in.shape s = self.ffn_router(ff_in)[..., 0] # [B, Tloc] routing score if route == "decode": keep_logit = self.ffn_predictor(ff_in)[..., 0] # causal, per-token decision if B == 1 and Tloc == 1: # single-token decode: truly skip gw = jax.nn.sigmoid(s)[..., None].astype(ff_in.dtype) return jax.lax.cond(keep_logit[0, 0] > 0.0, lambda: gw * self.mlp(ff_in), lambda: jnp.zeros_like(ff_in)) gate = jnp.where(keep_logit > 0.0, jax.nn.sigmoid(s), 0.0) return gate[..., None].astype(ff_in.dtype) * self.mlp(ff_in, sink=mlp_sink) k = max(1, int(round(self.ffn_keep * Tloc))) # capacity (static; Tloc known) top_s, idx = jax.lax.top_k(s, k) # [B, k] ff_in_k = jnp.take_along_axis(ff_in, idx[..., None], axis=1) # gather kept -> [B,k,D] ff_k = jax.nn.sigmoid(top_s)[..., None].astype(ff_in.dtype) * self.mlp(ff_in_k, sink=mlp_sink) out = jnp.zeros_like(ff_in).at[jnp.arange(B)[:, None], idx].set(ff_k) # scatter back if route == "train" and gate_sink is not None: # predictor BCE (mimic top-k; stop-grad) tgt = jnp.zeros_like(s).at[jnp.arange(B)[:, None], idx].set(1.0) logit = self.ffn_predictor(jax.lax.stop_gradient(ff_in))[..., 0] bce = jnp.maximum(logit, 0.0) - logit * tgt + jnp.log1p(jnp.exp(-jnp.abs(logit))) gate_sink.append(bce.mean()) return out def _read_mem(self, q_read, prior_mem): mk, mv = prior_mem # [B,T,G,hkv,hd], [B,T,G,hkv,vd] rep = self.hq // mk.shape[3] mk = jnp.repeat(mk, rep, axis=3) mv = jnp.repeat(mv, rep, axis=3) ms = jnp.einsum("bthd,btghd->bthg", q_read, mk).astype(F32) * self.scale ma = jax.nn.softmax(ms, axis=-1).astype(q_read.dtype) return jnp.einsum("bthg,btghe->bthe", ma, mv) def __call__(self, h, rope, kv, prior_mem, mlp_sink=None, gate_sink=None, route="train"): cfg = self.cfg B, T, _ = h.shape cos, sin = rope hd = cfg.head_dim xq = self.norm1(h) if self.diff: q = self.wq(xq).reshape(B, T, self.hq, 2, hd) q1, q2 = q[:, :, :, 0], q[:, :, :, 1] if cfg.use_qk_norm: q1, q2 = self.q_norm(q1), self.q_norm(q2) q1r, q2r = apply_rope(q1, cos, sin), apply_rope(q2, cos, sin) k1, k2, v = kv va, vb = v[..., :hd], v[..., hd:] lam = self._lam(h.dtype) # (A1 - lam*A2) @ [va|vb] == [A1@va - lam*A2@va | A1@vb - lam*A2@vb] oa = self._flash(q1r, k1, va, T) - lam * self._flash(q2r, k2, va, T) ob = self._flash(q1r, k1, vb, T) - lam * self._flash(q2r, k2, vb, T) out = jnp.concatenate([oa, ob], axis=-1) # [B,T,hq,2hd] out = self.head_norm(out) * (1.0 - self.lambda_init) if self.use_mem and prior_mem is not None: out = out + self.mem_gate.value.astype(h.dtype) * self._read_mem(q1, prior_mem) out = out.reshape(B, T, self.hq * 2 * hd) else: q = self.wq(xq).reshape(B, T, self.hq, hd) if cfg.use_qk_norm: q = self.q_norm(q) qr = apply_rope(q, cos, sin) k, v = kv out = self._flash(qr, k, v, T) # [B,T,hq,hd] if self.use_mem and prior_mem is not None: out = out + self.mem_gate.value.astype(h.dtype) * self._read_mem(q, prior_mem) out = out.reshape(B, T, self.hq * hd) h = h + self.wo(out) h = h + self._ffn_block(self.norm2(h), mlp_sink, gate_sink, route) return h def decode(self, h, rope_p, kv, prior_mem, valid): """Single-token decode. `h`:[B,1,D]; `kv` is the full preallocated cache (out_kv format); `valid`:[max_len] bool marks cached positions (<= pos). The one query attends the whole valid cache directly (it is the latest position, so no causal realignment).""" cfg = self.cfg; B = h.shape[0]; hd = cfg.head_dim cos, sin = rope_p xq = self.norm1(h) m = valid[None, None, None, :] # [1,1,1,max_len]; True == attend def attn(q, k, v): return jax.nn.dot_product_attention(q, k, v, mask=m) if self.diff: q = self.wq(xq).reshape(B, 1, self.hq, 2, hd) q1, q2 = q[:, :, :, 0], q[:, :, :, 1] if cfg.use_qk_norm: q1, q2 = self.q_norm(q1), self.q_norm(q2) q1r, q2r = apply_rope(q1, cos, sin), apply_rope(q2, cos, sin) k1, k2, v = kv; va, vb = v[..., :hd], v[..., hd:] lam = self._lam(h.dtype) oa = attn(q1r, k1, va) - lam * attn(q2r, k2, va) ob = attn(q1r, k1, vb) - lam * attn(q2r, k2, vb) out = jnp.concatenate([oa, ob], axis=-1) out = self.head_norm(out) * (1.0 - self.lambda_init) if self.use_mem and prior_mem is not None: out = out + self.mem_gate.value.astype(h.dtype) * self._read_mem(q1, prior_mem) out = out.reshape(B, 1, self.hq * 2 * hd) else: q = self.wq(xq).reshape(B, 1, self.hq, hd) if cfg.use_qk_norm: q = self.q_norm(q) qr = apply_rope(q, cos, sin); k, v = kv out = attn(qr, k, v) if self.use_mem and prior_mem is not None: out = out + self.mem_gate.value.astype(h.dtype) * self._read_mem(q, prior_mem) out = out.reshape(B, 1, self.hq * hd) h = h + self.wo(out) h = h + self._ffn_block(self.norm2(h), None, None, route="decode") return h # ---------------------------------------------------------------------------- # Block Attention Residuals: learned softmax aggregation over prior block states # ---------------------------------------------------------------------------- class AttnResAggregator(nnx.Module): def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs): self.norm = _rmsnorm(cfg.d_model, cfg, rngs) self.query = nnx.Param(jnp.zeros((cfg.d_model,), cfg.param_dtype)) def __call__(self, states): # states: [n, B, T, D] keys = self.norm(states) logits = jnp.einsum("d,nbtd->nbt", self.query.value.astype(states.dtype), keys) alpha = jax.nn.softmax(logits.astype(F32), axis=0).astype(states.dtype) return jnp.einsum("nbt,nbtd->btd", alpha, states) # ---------------------------------------------------------------------------- # Multi-Token Prediction: one shared, attention-free block reused across depths # ---------------------------------------------------------------------------- class MTPHead(nnx.Module): def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs): self.norm_h = _rmsnorm(cfg.d_model, cfg, rngs) self.norm_e = _rmsnorm(cfg.d_model, cfg, rngs) self.fuse = _linear(2 * cfg.d_model, cfg.d_model, cfg, rngs) self.block_norm = _rmsnorm(cfg.d_model, cfg, rngs) self.mlp = MLP(cfg, rngs) def step(self, h_prev, tok_emb): z = self.fuse(jnp.concatenate([self.norm_h(h_prev), self.norm_e(tok_emb)], axis=-1)) return z + self.mlp(self.block_norm(z)) # alias so the block can be driven through jax.checkpoint / nnx.split helpers # (which call the module, not a named method) when reused recurrently as a draft head. def __call__(self, h_prev, tok_emb): return self.step(h_prev, tok_emb) # ---------------------------------------------------------------------------- # Per-layer gradient checkpointing. "full" recomputes every sublayer internal in the # backward (attention scores, SwiGLU activations) and keeps only the residual stream -- # the memory-minimal policy. "selective" additionally keeps the big no-batch-dim GEMM # outputs (FFN + projections), recomputing only the cheaper attention/elementwise; this # recovers most of the recompute tax (higher MFU) but costs HBM -- only viable once the # sharded optimizer has freed the param-scale memory. # ---------------------------------------------------------------------------- _REMAT_POLICIES = { "full": None, # default jax.checkpoint policy = save nothing "selective": jax.checkpoint_policies.dots_with_no_batch_dims_saveable, } def _remat_sub(sub, h, rope, kv, prior, route, policy=None, want_gate=False): """Per-layer checkpoint. The MoD predictor BCE is returned as a VALUE (not via the gate_sink side-channel list) so it crosses jax.checkpoint cleanly -- remat + MoD compose.""" gdef, state = nnx.split(sub) def pure(state, h, rope, kv, prior): sink = [] if want_gate else None out = nnx.merge(gdef, state)(h, rope, kv, prior, gate_sink=sink, route=route) bce = sink[0] if sink else jnp.zeros((), F32) return out, bce return jax.checkpoint(pure, policy=policy)(state, h, rope, kv, prior) # ---------------------------------------------------------------------------- # Full model # ---------------------------------------------------------------------------- class CortexLM(nnx.Module): def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs): assert cfg.layers_per_block % cfg.kv_share_group == 0, ( f"layers_per_block ({cfg.layers_per_block}) must be a multiple of " f"kv_share_group ({cfg.kv_share_group})") assert not (cfg.global_every and cfg.kv_share_group > 1), ( "global_every is unsupported with CLA; use n_sink for global reach") self.cfg = cfg self.use_mem = cfg.use_growing_memory self.n_groups = cfg.n_layers // cfg.kv_share_group self.embed = nnx.Embed( cfg.vocab_size, cfg.d_model, embedding_init=nnx.initializers.normal(stddev=0.02), dtype=cfg.compute_dtype, param_dtype=cfg.param_dtype, rngs=rngs) self.block_kv = nnx.List([_BlockKV(cfg, rngs, self.use_mem) for _ in range(self.n_groups)]) self.subs = nnx.List([_Sublayer(cfg, rngs, self.use_mem, i) for i in range(cfg.n_layers)]) if cfg.use_attn_res: self.aggregators = nnx.List([AttnResAggregator(cfg, rngs) for _ in range(cfg.n_blocks + 1)]) self.final_norm = _rmsnorm(cfg.d_model, cfg, rngs) if cfg.use_factorized_head: self.head_transform = nnx.Linear( cfg.d_model, cfg.d_model, use_bias=False, kernel_init=_identity_init, dtype=cfg.compute_dtype, param_dtype=cfg.param_dtype, rngs=rngs) if not cfg.tie_embeddings: self.lm_head = _linear(cfg.d_model, cfg.vocab_size, cfg, rngs) if cfg.use_mtp: self.mtp = MTPHead(cfg, rngs) def _trunk(self, tokens, collect: bool = False, mlp_sink=None, gate_sink=None, route="train"): cfg = self.cfg T = tokens.shape[1] x = self.embed(tokens) rope = compute_rope(T, cfg.head_dim, cfg.rope_base, cfg, x.dtype) g = cfg.kv_share_group pool = [x] mem_list = [] kv = prior = new_mem = None gi = 0 for bi in range(cfg.n_blocks): h = self.aggregators[bi](jnp.stack(pool, 0)) if cfg.use_attn_res else pool[-1] for li in range(cfg.layers_per_block): if gi % g == 0: kv, new_mem = self.block_kv[gi // g](h, rope) prior = None if self.use_mem and mem_list: mk = jnp.stack([m[0] for m in mem_list], axis=2) mv = jnp.stack([m[1] for m in mem_list], axis=2) prior = (mk, mv) if cfg.use_remat and mlp_sink is None: want_gate = cfg.use_ffn_gate and gate_sink is not None h, bce = _remat_sub(self.subs[gi], h, rope, kv, prior, route, # per-layer checkpoint _REMAT_POLICIES[cfg.remat_policy], want_gate=want_gate) if want_gate: gate_sink.append(bce) else: h = self.subs[gi](h, rope, kv, prior, mlp_sink=mlp_sink, gate_sink=gate_sink, route=route) if gi % g == g - 1 and new_mem is not None: mem_list.append(new_mem) gi += 1 pool.append(h) if cfg.use_attn_res and cfg.use_attn_res_readout: x = self.aggregators[cfg.n_blocks](jnp.stack(pool, 0)) states = pool + [x] else: x = pool[-1] states = pool return (x, states) if collect else x def collect_mlp_acts(self, tokens): """Per-neuron mean |SwiGLU activation| for every trunk MLP -> [n_layers, d_ff]. Drives the dormant-neuron probe (cortex.neuron_probe); jit-safe.""" sink: list = [] self._trunk(tokens, mlp_sink=sink) return jnp.stack(sink, axis=0) def _head(self, hidden): """Tied/factorized LM head (final_norm -> d x d transform -> E^T). Safe to call on a sequence CHUNK [B, c, D] -> [B, c, V] for chunked loss.""" x = self.final_norm(hidden) if self.cfg.use_factorized_head: x = self.head_transform(x) if self.cfg.tie_embeddings: emb = self.embed.embedding.value.astype(x.dtype) logits = jnp.einsum("btd,vd->btv", x, emb) else: logits = self.lm_head(x) return logits.astype(F32) def forward_train_hidden(self, tokens): """Pre-head hidden states for chunked-CE training: (main_hidden, [mtp_hidden_k], gate_aux). gate_aux is the mean MoD causal-predictor BCE (0.0 if the gate is off); the train loop adds ``ffn_gate_aux_weight * gate_aux`` to the loss. Avoids ever building the [B,T,vocab] logit array in the trunk.""" gate_sink: list = [] hidden = self._trunk(tokens, gate_sink=gate_sink) mtp_hiddens = [] if self.cfg.use_mtp: emb = self.embed(tokens) B, T, D = emb.shape h = hidden for k in range(1, self.cfg.mtp_depth + 1): tok_emb = jnp.concatenate([emb[:, k:, :], jnp.zeros((B, k, D), emb.dtype)], axis=1) h = self.mtp.step(h, tok_emb) mtp_hiddens.append(h) gate_aux = jnp.mean(jnp.stack(gate_sink)) if gate_sink else jnp.asarray(0.0, F32) return hidden, mtp_hiddens, gate_aux def __call__(self, tokens): """Inference forward (causal): main next-token logits. MoD routing dispatches by length -- a prompt (T>1) uses the top-k capacity gather (static, real FFN savings, identical to training); a single decode token (T==1) uses the causal predictor and skips the FFN when it says so. Both are correct for any length, incl. one token.""" route = "decode" if tokens.shape[1] == 1 else "prefill" return self._head(self._trunk(tokens, route=route)) def forward_train(self, tokens): """(main_logits, [mtp_logits], gate_aux) -- materializes full logits; fine for small vocab (experiments/eval). Production training uses forward_train_hidden + losses.chunked_lm_loss to avoid the big logit array.""" hidden, mtp_hiddens, gate_aux = self.forward_train_hidden(tokens) return self._head(hidden), [self._head(h) for h in mtp_hiddens], gate_aux # ---- KV-cache decode: O(1)-tokens-through-the-trunk autoregressive generation ------- def init_decode_cache(self, B, max_len): """Preallocated KV (+ growing-memory running-sum) cache, one slot per KV-share group. Attention masks out the unwritten tail each step, so shapes stay static for jit.""" cfg = self.cfg; hd = cfg.head_dim; dt = cfg.compute_dtype kv, msum = [], [] for grp in range(self.n_groups): bk = self.block_kv[grp]; hkv, vd = bk.hkv, bk.vd if bk.diff: kv.append((jnp.zeros((B, max_len, hkv, hd), dt), jnp.zeros((B, max_len, hkv, hd), dt), jnp.zeros((B, max_len, hkv, vd), dt))) else: kv.append((jnp.zeros((B, max_len, hkv, hd), dt), jnp.zeros((B, max_len, hkv, vd), dt))) msum.append(jnp.zeros((B, hkv, vd), F32) if self.use_mem else None) return {"kv": kv, "msum": msum} def decode_step(self, token, pos, cache, cos, sin): """One autoregressive step against the cache. `token`:[B,1]; `pos`: scalar position; `cos`/`sin`: full RoPE tables [max_len, head_dim]. Returns (logits[B,V], new_cache). Only the new token flows through the FFN/projections -- the speedup vs full recompute.""" cfg = self.cfg; g = cfg.kv_share_group x = self.embed(token) # [B,1,D] rope_p = (jax.lax.dynamic_slice_in_dim(cos, pos, 1, axis=0), jax.lax.dynamic_slice_in_dim(sin, pos, 1, axis=0)) # [1, head_dim] valid = jnp.arange(cos.shape[0]) <= pos new_kv = list(cache["kv"]); new_msum = list(cache["msum"]) pool = [x]; mem_cur = []; gi = 0 kv_full = prior = mem_p = None for bi in range(cfg.n_blocks): h = self.aggregators[bi](jnp.stack(pool, 0)) if cfg.use_attn_res else pool[-1] for li in range(cfg.layers_per_block): if gi % g == 0: grp = gi // g kv_p, mem_p, msum_new = self.block_kv[grp].kv_decode( h, rope_p, cache["msum"][grp], pos) cur = cache["kv"][grp] kv_full = tuple(jax.lax.dynamic_update_slice_in_dim(cur[j], kv_p[j], pos, axis=1) for j in range(len(cur))) new_kv[grp] = kv_full; new_msum[grp] = msum_new prior = None if self.use_mem and mem_cur: mk = jnp.stack([mm[0] for mm in mem_cur], axis=2) # [B,1,Gprior,hkv,hd] mv = jnp.stack([mm[1] for mm in mem_cur], axis=2) prior = (mk, mv) h = self.subs[gi].decode(h, rope_p, kv_full, prior, valid) if gi % g == g - 1 and mem_p is not None: mem_cur.append(mem_p) gi += 1 pool.append(h) if cfg.use_attn_res and cfg.use_attn_res_readout: x = self.aggregators[cfg.n_blocks](jnp.stack(pool, 0)) else: x = pool[-1] return self._head(x)[:, 0], {"kv": new_kv, "msum": new_msum} # logits [B,V]