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