cortex/model.py

"""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]