cortex/latent_diffusion.py

"""Cortex-A 0.6 — latent-AR planner + autoregressive writer (v30).

Two-level causal LM: ~5x cheaper deep compute and a smaller KV cache than a full AR
transformer, while keeping AR's reliable, coherent decoding.

  tokens --embed--> [B,T,d]
        --chunk-pool(P)--> chunk means C [B,N,d]   (N = T/P)
  PLANNER  (deep latent-AR over CHUNKS): predicts the next chunk embedding chat[i] ~
        C[i+1].  Runs over T/P positions -> the MFU/KV win.  Cosine aux loss.  Kept
        blocks: RoPE + QK-norm + SwiGLU + Block-Attention-Residuals + SWA + learned
        sinks + growing causal KV-memory (cortex.model._Sublayer/_BlockKV).
  WRITER   (shallow causal AR over TOKENS): a standard next-token transformer, each
        position additionally conditioned on the planner's hint for the chunk it is
        writing -- cond[t] = chat[(t+1)//P - 1], which is strictly causal (chat[c-1]
        depends only on tokens < c*P <= t).  Next-token CE loss; the CE gradient flows
        back through `cond` into the planner, so the hint learns to encode whatever
        helps the writer pick the token, NOT just the lossy chunk mean.

No diffusion (no noise / timesteps / denoise / self-conditioning).  The writer is the
same certified causal block as the backbone, so flash-causal training AND KV-cache
decode are reused verbatim.  The readout CE is now standard next-token perplexity --
directly comparable to a vanilla AR transformer.
"""

from __future__ import annotations

import dataclasses
import math

import jax
import jax.numpy as jnp
from flax import nnx

from .config import ModelConfig
from .losses import _chunked_ce
from .model import (
    F32, MLP, MTPHead, AttnResAggregator, _BlockKV, _REMAT_POLICIES, _Sublayer,
    _identity_init, _linear, _rmsnorm, compute_rope,
)


# ----------------------------------------------------------------------------
# Config clones.  Both planner and writer run the kept blocks (RoPE/QK-norm/SwiGLU/
# Block-Attn-Residuals) with diff-attn/GQA/CLA/MoD removed; _Sublayer then degenerates
# to a plain causal MHA+SwiGLU block and _BlockKV to per-layer K/V.
# ----------------------------------------------------------------------------
def backbone_cfg(cfg: ModelConfig) -> ModelConfig:
    # PLANNER over chunks.  Keeps SWA + learned sinks + growing memory (rescaling a
    # token-defined window into chunk units).
    win = cfg.sliding_window
    if win is not None:
        win = max(1, -(-win // cfg.patch_size))
    return dataclasses.replace(
        cfg, use_diff_attn=False, n_kv_heads=cfg.n_q_heads, kv_share_group=1,
        use_ffn_gate=False, ffn_keep=1.0, use_mtp=False,   # inherit use_pallas_attn from parent
        sliding_window=win)


def writer_cfg(cfg: ModelConfig) -> ModelConfig:
    # WRITER over tokens: dec_layers deep, organized into dec_blocks Block-Attn-Residual
    # blocks, causal sliding-window (dec_window, in TOKEN units), no sinks/memory.
    return dataclasses.replace(
        cfg, use_diff_attn=False, n_kv_heads=cfg.n_q_heads, kv_share_group=1,
        use_ffn_gate=False, ffn_keep=1.0, use_mtp=False,   # inherit use_pallas_attn from parent
        use_growing_memory=False, learned_sink=0, n_sink=0,
        sliding_window=cfg.dec_window,
        n_blocks=cfg.dec_blocks, layers_per_block=cfg.dec_layers // cfg.dec_blocks)


def _ckpt_call(mod, *args, policy=None):
    """Per-layer gradient checkpoint (split -> pure call -> jax.checkpoint)."""
    gdef, state = nnx.split(mod)

    def pure(state, *a):
        return nnx.merge(gdef, state)(*a)

    return jax.checkpoint(pure, policy=policy)(state, *args)


def _cond_fuse_init(key, shape, dtype=jnp.float32):
    """Init for the writer's full-conditioning fuse Linear (2d -> d).  Kernel is [in=2d, out=d];
    the input is concat([norm(token), norm(cond)]).  Top d rows = I (pass the token through),
    bottom d rows = 0.1*I (warm but small plan contribution) -> writer_in starts at
    norm(token) + 0.1*norm(cond): the proven token pathway survives a --reinit resume while the
    planner plan is used from step 1 and grows as the CE gradient demands."""
    d = shape[1]
    return jnp.concatenate([jnp.eye(d, d, dtype=dtype), 0.1 * jnp.eye(d, d, dtype=dtype)], axis=0)


# ----------------------------------------------------------------------------
# Parallel multi-token writer head (Medusa / MTP / blockwise-parallel decoding).
# ----------------------------------------------------------------------------
class _SlotAttn(nnx.Module):
    """Bidirectional multi-head self-attention over the P-1 draft slots of ONE chunk,
    batched over [B, N].  No RoPE and no causal mask: the slots are order-tagged by the
    head's learned slot embeddings and are drafted together, so each slot may attend to
    the others' (plan/position-derived) QUERY context -- never to their predicted tokens."""

    def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs):
        self.h, self.hd = cfg.n_q_heads, cfg.head_dim
        self.use_qk_norm = cfg.use_qk_norm
        self.qkv = _linear(cfg.d_model, 3 * cfg.d_model, cfg, rngs)
        self.out = _linear(cfg.d_model, cfg.d_model, cfg, rngs,
                           scale=1.0 / math.sqrt(2 * max(cfg.n_layers, 1)))
        if cfg.use_qk_norm:
            self.q_norm = _rmsnorm(self.hd, cfg, rngs)
            self.k_norm = _rmsnorm(self.hd, cfg, rngs)

    def __call__(self, x):                                      # x:[B,N,S,d] -> [B,N,S,d]
        B, N, S, d = x.shape
        qkv = self.qkv(x).reshape(B, N, S, 3, self.h, self.hd)
        q, k, v = qkv[..., 0, :, :], qkv[..., 1, :, :], qkv[..., 2, :, :]   # [B,N,S,h,hd]
        if self.use_qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)
        attn = jnp.einsum("bnshd,bnthd->bnhst", q, k) * (1.0 / math.sqrt(self.hd))
        a = jax.nn.softmax(attn.astype(F32), axis=-1).astype(x.dtype)
        o = jnp.einsum("bnhst,bnthd->bnshd", a, v).reshape(B, N, S, d)
        return self.out(o)


class ParallelWriterHead(nnx.Module):
    """Drafts tokens 2..P of every chunk IN PARALLEL (Medusa/MTP, adapted to our chunks).

    Per chunk c the head sees only three things, all available BEFORE tokens c*P+1.. exist:
      g  = the planner's plan for chunk c (chat[c-1]; depends on tokens < c*P),
      e1 = the embedding of the chunk's FIRST token (the one the AR writer emits normally),
      and a learned per-slot position (slot j -> token c*P+j).
    It NEVER consumes tokens c*P+1.., so the identical compute runs at inference; the AR
    writer then VERIFIES the drafts (speculative decoding) -> output == pure AR writer.

    Lightweight: one transformer block (bidirectional slot self-attention + SwiGLU FFN),
    reusing the certified MLP/RMSNorm/Linear primitives.  Logits go through the model's
    shared tied/factorized head (drafts live in the writer's own logit space)."""

    def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs):
        self.n_slot = cfg.patch_size - 1
        self.norm_g = _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.slot_emb = nnx.Param(nnx.initializers.normal(stddev=0.02)(
            rngs.params(), (self.n_slot, cfg.d_model), cfg.param_dtype))
        self.attn_norm = _rmsnorm(cfg.d_model, cfg, rngs)
        self.attn = _SlotAttn(cfg, rngs)
        self.mlp_norm = _rmsnorm(cfg.d_model, cfg, rngs)
        self.mlp = MLP(cfg, rngs)

    def __call__(self, g, e1):                                  # g,e1:[B,N,d] -> [B,N,P-1,d]
        ctx = self.fuse(jnp.concatenate([self.norm_g(g), self.norm_e(e1)], axis=-1))   # [B,N,d]
        x = ctx[:, :, None, :] + self.slot_emb.value.astype(ctx.dtype)[None, None]     # [B,N,S,d]
        x = x + self.attn(self.attn_norm(x))                    # bidirectional over the S slots
        x = x + self.mlp(self.mlp_norm(x))                      # SwiGLU FFN
        return x


class CortexLatentDiffusion(nnx.Module):
    # (name kept for the resume.json arch tag + checkpoint compatibility; it is now a
    #  latent-AR planner + AR writer, no diffusion.)
    def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs):
        self.cfg = cfg
        self.P = cfg.patch_size
        bcfg = backbone_cfg(cfg)
        wcfg = writer_cfg(cfg)
        self.bcfg, self.wcfg = bcfg, wcfg

        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)

        # ---- PLANNER (unchanged from the working backbone; resumes from checkpoint) ----
        self.use_mem = bcfg.use_growing_memory
        self.n_groups = bcfg.n_layers
        lpb = bcfg.layers_per_block
        self.block_kv = nnx.List([
            _BlockKV(bcfg, rngs, self.use_mem and ((i + 1) % lpb == 0))
            for i in range(self.n_groups)])
        self.subs = nnx.List([_Sublayer(bcfg, rngs, self.use_mem, i) for i in range(bcfg.n_layers)])
        if bcfg.use_attn_res:
            self.aggregators = nnx.List(
                [AttnResAggregator(bcfg, rngs) for _ in range(bcfg.n_blocks + 1)])
        self.latent_norm = _rmsnorm(cfg.d_model, cfg, rngs)
        self.latent_head = _linear(cfg.d_model, cfg.d_model, cfg, rngs)

        # ---- tied / factorized output head (kept) ----
        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)

        # ---- WRITER conditioning: full plan fusion (use_full_cond) or v30 additive hint ----
        self.use_full_cond = cfg.use_full_cond
        if cfg.use_full_cond:
            self.cond_norm_tok = _rmsnorm(cfg.d_model, cfg, rngs)
            self.cond_norm_cond = _rmsnorm(cfg.d_model, cfg, rngs)
            self.cond_fuse = nnx.Linear(                      # [norm(tok) | norm(plan)] -> writer input
                2 * cfg.d_model, cfg.d_model, use_bias=False, kernel_init=_cond_fuse_init,
                dtype=cfg.compute_dtype, param_dtype=cfg.param_dtype, rngs=rngs)
        else:
            self.cond_proj = _linear(cfg.d_model, cfg.d_model, cfg, rngs)   # v30 additive hint
        self.w_block_kv = nnx.List([_BlockKV(wcfg, rngs, False) for _ in range(wcfg.n_layers)])
        self.w_subs = nnx.List([_Sublayer(wcfg, rngs, False, i) for i in range(wcfg.n_layers)])
        if wcfg.use_attn_res:
            self.w_agg = nnx.List(
                [AttnResAggregator(wcfg, rngs) for _ in range(wcfg.n_blocks + 1)])

        # ---- PARALLEL WRITER HEAD (optional; Medusa/MTP speculative drafting, trains fresh) ----
        self.use_pwriter = cfg.use_parallel_writer
        if self.use_pwriter:
            self.pwriter = ParallelWriterHead(cfg, rngs)

        # ---- RECURRENT EAGLE/MTP DRAFT HEAD (optional): ONE shared block reused across K depths
        # to draft tokens t+2..t+1+K from the writer hidden feature.  Trains fresh (--reinit). ----
        self.use_writer_mtp = cfg.use_writer_mtp
        if self.use_writer_mtp:
            self.writer_mtp = MTPHead(cfg, rngs)

    # ---- shared kept-block trunk (planner over chunks, writer over tokens) ----
    def _run_trunk(self, x, rope, cfg_, block_kv, subs, aggs, use_mem):
        pool = [x]
        mem_list = []
        gi = 0
        for bi in range(cfg_.n_blocks):
            h = aggs[bi](jnp.stack(pool, 0)) if cfg_.use_attn_res else pool[-1]
            for _ in range(cfg_.layers_per_block):
                kv, new_mem = block_kv[gi](h, rope)
                prior = None
                if 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:
                    h = _ckpt_call(subs[gi], h, rope, kv, prior,
                                   policy=_REMAT_POLICIES[cfg_.remat_policy])
                else:
                    h = subs[gi](h, rope, kv, prior)
                if 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:
            return aggs[cfg_.n_blocks](jnp.stack(pool, 0))
        return pool[-1]

    def _backbone(self, x):                                    # x:[B,N,d] -> chunk hidden
        cfg = self.bcfg
        rope = compute_rope(x.shape[1], cfg.head_dim, cfg.rope_base, cfg, x.dtype)
        return self._run_trunk(x, rope, cfg, self.block_kv, self.subs, self.aggregators, self.use_mem)

    def _plan(self, emb):                                      # emb:[B,T,d] -> chat:[B,N,d]
        Cm = jax.lax.stop_gradient(self._chunk_means(emb))     # cosine target (E shaped only by CE)
        chat = self.latent_head(self.latent_norm(self._backbone(Cm)))
        return Cm, chat

    def _writer_input(self, emb, cond):                       # emb,cond:[B,*,d] -> writer input [B,*,d]
        if self.use_full_cond:                                # full-bandwidth plan fusion (concat 2d->d)
            return self.cond_fuse(jnp.concatenate(
                [self.cond_norm_tok(emb), self.cond_norm_cond(cond)], axis=-1))
        return emb + self.cond_proj(cond)                     # v30 additive hint

    def _writer(self, emb, cond):                              # emb,cond:[B,T,d] -> hidden:[B,T,d]
        cfg = self.wcfg
        rope = compute_rope(emb.shape[1], cfg.head_dim, cfg.rope_base, cfg, emb.dtype)
        x = self._writer_input(emb, cond)
        return self._run_trunk(x, rope, cfg, self.w_block_kv, self.w_subs, self.w_agg, False)

    def _head(self, hidden):
        x = self.final_norm(hidden)
        if self.cfg.use_factorized_head:
            x = self.head_transform(x)
        emb = self.embed.embedding.value.astype(x.dtype)
        return jnp.einsum("btd,vd->btv", x, emb).astype(F32)

    def _chunk_means(self, emb):
        B, T, d = emb.shape
        return emb.reshape(B, T // self.P, self.P, d).mean(axis=2)

    def _cond_from_chat(self, chat, T):
        """Per-token planner hint: position t predicts token t+1 (chunk c=(t+1)//P),
        hinted by chat[c-1].  chat[c-1] depends only on tokens < c*P <= t -> strictly
        causal.  Positions predicting chunk-0 tokens (no prior chunk) get a zero hint."""
        N = chat.shape[1]
        src = (jnp.arange(T) + 1) // self.P - 1                # [T]
        cond = chat[:, jnp.clip(src, 0, N - 1)]                # [B,T,d]
        return jnp.where((src >= 0)[None, :, None], cond, 0.0).astype(chat.dtype)

    # ---- parallel multi-token draft head (Medusa/MTP).  Trains the planner + embeddings to
    # pack enough into each chunk plan that the next P-1 tokens decode from it in ONE shot ----
    def _parallel_draft_inputs(self, emb, chat):
        """Per-chunk (plan, first-token) for the parallel draft head.  The plan for chunk c is
        chat[c-1] (chunk 0 has none -> zero); the first token is emb[:, c*P].  NO stop-gradient
        anywhere -> the draft CE flows into the planner (through chat) and the embedding table
        (through e1 + the tied head), co-training the WHOLE model, not just the head."""
        B, T, d = emb.shape
        N, P = T // self.P, self.P
        e1 = emb.reshape(B, N, P, d)[:, :, 0]                            # [B,N,d] chunk first token
        g = jnp.concatenate([jnp.zeros((B, 1, d), chat.dtype), chat[:, :-1]], axis=1)  # [B,N,d] plan
        return g, e1

    def _parallel_draft_ce(self, emb, chat, tokens, ce_chunk, z_loss_coef):
        B, T = tokens.shape
        N, P, d = T // self.P, self.P, emb.shape[-1]
        g, e1 = self._parallel_draft_inputs(emb, chat)
        if self.cfg.use_remat:                                          # recompute the head in backward
            hid = _ckpt_call(self.pwriter, g, e1, policy=_REMAT_POLICIES[self.cfg.remat_policy])
        else:
            hid = self.pwriter(g, e1)
        hid = hid.reshape(B, N * (P - 1), d)                            # [B,N*(P-1),d]
        tgt = tokens.reshape(B, N, P)[:, :, 1:].reshape(B, N * (P - 1))  # tokens 2..P of each chunk
        ce_tot, _ = _chunked_ce(self._head, hid, tgt, ce_chunk, z_loss_coef)
        return ce_tot

    # ---- recurrent EAGLE/MTP draft: reuse ONE shared block over K depths to predict the next K
    # tokens from the writer hidden.  "As many tokens as required" == cfg.writer_mtp_depth. ----
    def _writer_mtp_ce(self, hidden, emb, tokens, ce_chunk, z_loss_coef, key=None):
        """Step k advances the feature f (f0 = writer hidden, which already carries the planner
        plan via `cond`) by fusing it with emb(token_{t+k}) -- the real teacher-forced token, so
        f_k[:,t] sees only the writer hidden at t and tokens t+1..t+k (all < t+1+k): strictly
        causal / leak-free, identical compute at inference.  Its tied-head logits predict token
        t+1+k.  Depth-decayed CE; the grad co-trains the writer, the planner (through `hidden`)
        and the embeddings (fed token + tied head).  ONE block (self.writer_mtp) is reused for
        every depth -> FLOPs-light + a runtime-variable draft length.

        writer_mtp_subsample<1 scores only a strided random subset of positions.  The head is
        position-wise, so one fixed subset can carry the whole recurrence and this is an EXACT
        per-position estimate of the depth-averaged CE -- and because the full-vocab head is the
        cost, it cuts the MTP MFU tax ~proportionally (0.25 -> ~4x cheaper)."""
        cfg = self.cfg
        B, T, d = emb.shape
        K = cfg.writer_mtp_depth
        maxvalid = T - (1 + K)              # pos p needs token p+1+k for all k<=K -> p <= T-2-K
        if maxvalid <= 0:
            return jnp.asarray(0.0, F32)

        if cfg.writer_mtp_subsample < 1.0:                          # ---- subsampled (strided + random offset)
            m = min(maxvalid, max(64, int(round(cfg.writer_mtp_subsample * maxvalid))))
            stride = max(1, maxvalid // m)
            off = jax.random.randint(key, (), 0, stride) if key is not None else jnp.int32(0)
            idx = (off + stride * jnp.arange(m)) % maxvalid         # [m] positions in [0, maxvalid)
            f = jnp.take(hidden, idx, axis=1)                       # [B,m,d]
            terms, wsum = [], 0.0
            for k in range(1, K + 1):
                tok_emb = jnp.take(emb, idx + k, axis=1)            # emb(token p+k) at the subset
                if cfg.use_remat:
                    f = _ckpt_call(self.writer_mtp, f, tok_emb, policy=_REMAT_POLICIES[cfg.remat_policy])
                else:
                    f = self.writer_mtp(f, tok_emb)
                tgt = jnp.take(tokens, idx + 1 + k, axis=1)         # [B,m] target token p+1+k
                ce_k, _ = _chunked_ce(self._head, f, tgt, ce_chunk, z_loss_coef)
                w = cfg.writer_mtp_decay ** (k - 1)
                terms.append(w * ce_k); wsum += w
            return sum(terms) / wsum

        f = hidden                                                  # ---- full (all positions); f[:,t]->t+1
        terms, wsum = [], 0.0
        for k in range(1, K + 1):
            valid = T - (1 + k)
            tok_emb = jnp.concatenate(                              # emb(token t+k) aligned to pos t
                [emb[:, k:], jnp.zeros((B, k, d), emb.dtype)], axis=1)
            if cfg.use_remat:
                f = _ckpt_call(self.writer_mtp, f, tok_emb, policy=_REMAT_POLICIES[cfg.remat_policy])
            else:
                f = self.writer_mtp(f, tok_emb)
            ce_k, _ = _chunked_ce(self._head, f[:, :valid], tokens[:, 1 + k:], ce_chunk, z_loss_coef)
            w = cfg.writer_mtp_decay ** (k - 1)
            terms.append(w * ce_k); wsum += w
        return sum(terms) / wsum

    # ------------------------------------------------------------------ training
    def compute_loss(self, tokens, key=None, *, ce_chunk: int = 512, z_loss_coef: float = 1e-4):
        """Planner cosine aux + writer next-token CE.  `key` (optional) drives per-chunk
        conditioning dropout so the writer also learns from context alone."""
        cfg = self.cfg
        P, (B, T) = self.P, tokens.shape
        N = T // P
        emb = self.embed(tokens)
        Cm, chat = self._plan(emb)
        pf = chat[:, :N - 1].astype(F32)
        tf = Cm[:, 1:].astype(F32)
        # SAFE norms: eps INSIDE the sqrt. jnp.linalg.norm(x) has a 0/0 (NaN) gradient when x
        # is a zero vector -- and a predicted chunk embedding pf CAN collapse to ~0, which then
        # poisons the whole grad. sqrt(sum(x^2)+eps) keeps the backward finite at x==0.
        pf_n = jnp.sqrt(jnp.sum(pf * pf, axis=-1) + 1e-12)
        tf_n = jnp.sqrt(jnp.sum(tf * tf, axis=-1) + 1e-12)
        cos = (1.0 - (pf * tf).sum(-1) / (pf_n * tf_n + 1e-6)).mean()

        mtp_key = None
        if key is not None:
            key, mtp_key = jax.random.split(key)               # separate stream for MTP subsampling
        cond = self._cond_from_chat(chat, T)
        if key is not None and cfg.cond_dropout > 0:           # per-chunk hint dropout
            keep = jax.random.bernoulli(key, 1.0 - cfg.cond_dropout, (B, N, 1)).astype(cond.dtype)
            cond = cond * keep[:, jnp.clip((jnp.arange(T) + 1) // P - 1, 0, N - 1)]
        hidden = self._writer(emb, cond)
        ce_tot, ce = _chunked_ce(self._head, hidden[:, :-1], tokens[:, 1:], ce_chunk, z_loss_coef)
        total = cfg.w_cos * cos + cfg.w_ce * ce_tot
        pw_ce = jnp.asarray(0.0, F32)
        if self.use_pwriter:                                    # parallel-draft MTP aux loss; its grad
            pw_ce = self._parallel_draft_ce(emb, chat, tokens, ce_chunk, z_loss_coef)  # co-trains head +
            total = total + cfg.pw_loss_weight * pw_ce          # planner (via chat) + embeddings (via e1)
        mtp_ce = jnp.asarray(0.0, F32)
        if self.use_writer_mtp:                                 # recurrent EAGLE/MTP draft (next K tokens
            mtp_ce = self._writer_mtp_ce(hidden, emb, tokens, ce_chunk, z_loss_coef, mtp_key)  # writer hidden
            total = total + cfg.writer_mtp_weight * mtp_ce
        return total, {"loss": total, "ce": ce, "cos": cos,
                       "denoise": jnp.asarray(0.0, F32), "mtp_ce": mtp_ce,
                       "pw_ce": pw_ce}

    def eval_ce(self, tokens, *, ce_chunk: int = 512, ablate_plan: bool = False):
        """Standard teacher-forced next-token CE -> honest perplexity (exp(CE)).  ablate_plan
        ZEROES the planner conditioning (writer runs WITHOUT the latent backbone); the gap vs
        the normal CE measures how much the writer actually relies on the planner."""
        emb = self.embed(tokens)
        _, chat = self._plan(emb)
        cond = self._cond_from_chat(chat, tokens.shape[1])
        if ablate_plan:
            cond = jnp.zeros_like(cond)
        hidden = self._writer(emb, cond)
        _, ce = _chunked_ce(self._head, hidden[:, :-1], tokens[:, 1:], ce_chunk, 0.0)
        return ce

    # ------------------------------------------------------------------ inference
    # PLANNER cache (chunks): KV + learned sinks + growing-memory running sums.
    def init_chunk_cache(self, B: int, max_chunks: int):
        bcfg = self.bcfg
        hd, dt, ls = bcfg.head_dim, bcfg.compute_dtype, bcfg.learned_sink
        k_c, v_c, msum = [], [], []
        for bk in self.block_kv:
            k = jnp.zeros((B, max_chunks + ls, bk.hkv, hd), dt)
            v = jnp.zeros((B, max_chunks + ls, bk.hkv, bk.vd), dt)
            if ls:
                bc = lambda p: jnp.broadcast_to(p.value.astype(dt)[None], (B,) + p.value.shape)
                k = k.at[:, max_chunks:].set(bc(bk.sink_k))
                v = v.at[:, max_chunks:].set(bc(bk.sink_v))
            k_c.append(k); v_c.append(v)
            msum.append(jnp.zeros((B, bk.hkv, bk.vd), F32) if bk.use_mem else None)
        return {"k": k_c, "v": v_c, "msum": msum}

    def predict_chunk(self, cm, pos, cache, cos, sin):
        """One chunk mean [B,1,d] at chunk `pos` through the cached planner -> next-chunk
        embedding [B,1,d]."""
        bcfg = self.bcfg
        mc = cache["k"][0].shape[1] - bcfg.learned_sink
        cos, sin = cos.astype(cm.dtype), sin.astype(cm.dtype)
        rope_p = (jax.lax.dynamic_slice_in_dim(cos, pos, 1, axis=0),
                  jax.lax.dynamic_slice_in_dim(sin, pos, 1, axis=0))
        idx = jnp.arange(mc + bcfg.learned_sink)
        w = bcfg.sliding_window                                 # windowed-causal (match training)
        causal = (idx <= pos) & ((idx > pos - w) if w is not None else True)
        valid = causal | (idx >= mc)                           # + learned sinks always visible
        new_k, new_v, new_ms = list(cache["k"]), list(cache["v"]), list(cache["msum"])
        pool = [cm]; mem_cur = []; gi = 0
        for bi in range(bcfg.n_blocks):
            h = self.aggregators[bi](jnp.stack(pool, 0)) if bcfg.use_attn_res else pool[-1]
            for _ in range(bcfg.layers_per_block):
                bk = self.block_kv[gi]
                kv_p, mem_p, ms_new = bk.kv_decode(h, rope_p, cache["msum"][gi], pos)
                new_k[gi] = jax.lax.dynamic_update_slice_in_dim(new_k[gi], kv_p[0], pos, axis=1)
                new_v[gi] = jax.lax.dynamic_update_slice_in_dim(new_v[gi], kv_p[1], pos, axis=1)
                new_ms[gi] = ms_new
                prior = None
                if self.use_mem and mem_cur:
                    mk = jnp.stack([m[0] for m in mem_cur], axis=2)
                    mv = jnp.stack([m[1] for m in mem_cur], axis=2)
                    prior = (mk, mv)
                h = self.subs[gi].decode(h, rope_p, (new_k[gi], new_v[gi]), prior, valid)
                if mem_p is not None:
                    mem_cur.append(mem_p)
                gi += 1
            pool.append(h)
        if bcfg.use_attn_res and bcfg.use_attn_res_readout:
            h = self.aggregators[bcfg.n_blocks](jnp.stack(pool, 0))
        else:
            h = pool[-1]
        pred = self.latent_head(self.latent_norm(h))
        return pred, {"k": new_k, "v": new_v, "msum": new_ms}

    # WRITER cache (tokens): plain per-layer KV (no sinks/memory).
    def init_writer_cache(self, B: int, max_len: int):
        wcfg = self.wcfg
        hd, dt = wcfg.head_dim, wcfg.compute_dtype
        k_c = [jnp.zeros((B, max_len, bk.hkv, hd), dt) for bk in self.w_block_kv]
        v_c = [jnp.zeros((B, max_len, bk.hkv, bk.vd), dt) for bk in self.w_block_kv]
        return {"k": k_c, "v": v_c}

    def writer_step(self, tok, cond_vec, pos, wcache, cos, sin):
        """One token: embed(tok)+cond_proj(hint) through the cached causal writer ->
        next-token logits [B,V]."""
        wcfg = self.wcfg
        x = self._writer_input(self.embed(tok), cond_vec)      # [B,1,d]
        cos, sin = cos.astype(x.dtype), sin.astype(x.dtype)
        rope_p = (jax.lax.dynamic_slice_in_dim(cos, pos, 1, axis=0),
                  jax.lax.dynamic_slice_in_dim(sin, pos, 1, axis=0))
        idx = jnp.arange(wcache["k"][0].shape[1])
        w = wcfg.sliding_window                                 # windowed-causal (match training)
        valid = (idx <= pos) & ((idx > pos - w) if w is not None else True)
        new_k, new_v = list(wcache["k"]), list(wcache["v"])
        pool = [x]; gi = 0
        for bi in range(wcfg.n_blocks):
            h = self.w_agg[bi](jnp.stack(pool, 0)) if wcfg.use_attn_res else pool[-1]
            for _ in range(wcfg.layers_per_block):
                kv_p, _, _ = self.w_block_kv[gi].kv_decode(h, rope_p, None, pos)
                new_k[gi] = jax.lax.dynamic_update_slice_in_dim(new_k[gi], kv_p[0], pos, axis=1)
                new_v[gi] = jax.lax.dynamic_update_slice_in_dim(new_v[gi], kv_p[1], pos, axis=1)
                h = self.w_subs[gi].decode(h, rope_p, (new_k[gi], new_v[gi]), None, valid)
                gi += 1
            pool.append(h)
        if wcfg.use_attn_res and wcfg.use_attn_res_readout:
            h = self.w_agg[wcfg.n_blocks](jnp.stack(pool, 0))
        else:
            h = pool[-1]
        return self._head(h)[:, 0], {"k": new_k, "v": new_v}   # logits [B,V]

    def draft_chunk(self, plan, first_tok):
        """Speculative proposal: from the planner plan [B,1,d] for a chunk and that chunk's
        already-emitted first token id [B,1], draft the next P-1 token ids [B,P-1] in ONE
        parallel pass.  The AR writer then VERIFIES them (writer_step), so accepted output is
        bit-identical to pure autoregressive writer decoding -- the head only saves steps."""
        e1 = self.embed(first_tok)                             # [B,1,d]
        hid = self.pwriter(plan, e1)                           # [B,1,P-1,d]
        return jnp.argmax(self._head(hid[:, 0]), axis=-1)      # [B,P-1] token ids

    def draft_mtp(self, hidden_t, k_draft):
        """Recurrent EAGLE/MTP speculative draft: from the writer hidden at the current position
        [B,1,d], reuse the shared block autoregressively to draft the next `k_draft` tokens (greedy
        feature feedback).  k_draft is a RUNTIME argument -- draft as many tokens as you want, the
        one trained block is reused for every step.  The AR writer then VERIFIES the drafts
        (writer_step), so accepted output is bit-identical to pure AR decoding."""
        f = hidden_t                                           # [B,1,d]
        ids = []
        for _ in range(int(k_draft)):
            nxt = jnp.argmax(self._head(f)[:, -1], axis=-1)    # [B] next-token prediction
            ids.append(nxt)
            f = self.writer_mtp(f, self.embed(nxt[:, None]))   # advance feature with the drafted token
        return jnp.stack(ids, axis=1)                          # [B, k_draft] token ids