1"""Cache-aware incremental decoding with per-token workspace routing (forced budget or learned controller).23For every new position t: backbone step with KV cache -> H_t; bank append (K_t, V_t); z_t; fresh S0; actions4(forced or controller-chosen, hard cap r_max); differential readout; logits; sample. The workspace is discarded5after the token (no persistent memory). The explicit budget-0 branch never allocates slots.6"""7from __future__ import annotations89from typing import Callable, Optional1011import torch1213from .cache import RequestBatchCache14from .model import PrismeModel151617@torch.no_grad()18def step_logits(model: PrismeModel, tokens_new: torch.Tensor, cache: RequestBatchCache, max_budget: int,19 cost: Optional[torch.Tensor] = None, forced_actions: Optional[torch.Tensor] = None,20 lambda_cost=None, kappa_risk=None):21 """tokens_new: [B, n] new tokens (prefill n>=1 or decode n=1). Returns logits [B,n,V] and the action trace."""22 cfg = model.cfg23 B, n = tokens_new.shape24 start = cache.length25 H = model.backbone(tokens_new, cache=cache.layers) # [B,n,d]26 if not cfg.use_workspace:27 return model.backbone.head(H), None, torch.zeros(B, n, dtype=torch.long, device=H.device)28 if max_budget == 0:29 # explicit bypass: no workspace allocation, no bank read; still extend the bank for later tokens30 hb = model.bank_norm(H)31 cache.bank.append(model.W_K(hb), model.W_V(hb))32 return model.backbone.head(H), None, torch.zeros(B, n, dtype=torch.long, device=H.device)33 hb = model.bank_norm(H)34 cache.bank.append(model.W_K(hb), model.W_V(hb))35 Tb = cache.bank.length36 z = model.P(hb)37 qpos = torch.arange(start, start + n, device=H.device).view(1, n).expand(B, n)38 key_pos = torch.arange(Tb, device=H.device)39 bank_mask = key_pos.view(1, 1, Tb) <= qpos.unsqueeze(-1)40 from .model import EncodedContext41 from .workspace import init_workspace42 entropy = None43 if cfg.ctrl_use_entropy:44 p = torch.softmax(model.backbone.head(H).float(), -1)45 entropy = -(p * torch.log(p.clamp_min(1e-12))).sum(-1)46 ctx = EncodedContext(H=H, K=cache.bank.K, V=cache.bank.V, z=z, bank_mask=bank_mask, S0=init_workspace(model.slot_embed, z),47 ctrl_in=(hb if cfg.ctrl_input == "h" else z), entropy=entropy)48 S = ctx.S049 R = min(max_budget, cfg.r_max)50 trace = torch.full((B, n, R), cfg.stop_action, dtype=torch.long, device=H.device)51 n_edits = torch.zeros(B, n, dtype=torch.long, device=H.device)52 stopped = torch.zeros(B, n, dtype=torch.bool, device=H.device)53 for r in range(R):54 step = torch.full((B, n), r, device=H.device, dtype=torch.long)55 remaining = torch.full((B, n), R - r - 1, device=H.device, dtype=torch.long)56 if forced_actions is not None:57 a = forced_actions[:, :, r]58 else:59 preds = model.controller(ctx.ctrl_in, S, step, remaining, ctx.entropy)60 a = model.controller.select(preds, cost, lambda_cost=lambda_cost, kappa_risk=kappa_risk)61 a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)62 stopped = stopped | (a == cfg.stop_action)63 trace[:, :, r] = a64 S, _ = model.apply_action(S, ctx, a, step, remaining)65 n_edits += (a != cfg.stop_action).long()66 if bool(stopped.all()):67 break68 logits, _ = model.logits_from_state(S, ctx, n_edits > 0)69 return logits, trace, n_edits707172@torch.no_grad()73def generate_uncached(model, prompt: torch.Tensor, max_new_tokens: int, max_budget: int = 0, temperature: float = 0.0,74 eos_id=None, forced_budget_actions=None, generator=None, **_):75 """Fallback for models without the PRISME caches (e.g. the global-recurrence baseline): recompute the full76 sequence at every step with forced actions for all positions. O(T^2) per token; used only for evaluation."""77 cfg = model.cfg78 B, T0 = prompt.shape79 seq = prompt80 done = torch.zeros(B, dtype=torch.bool, device=prompt.device)81 edits = []82 for i in range(max_new_tokens):83 T = seq.shape[1]84 acts = forced_budget_actions(B, T) if (forced_budget_actions is not None and max_budget > 0) else None85 out = model(seq, acts)86 nxt = out["logits"][:, -1]87 if out.get("n_edits") is not None:88 edits.append(out["n_edits"][:, -1])89 if temperature <= 0:90 tok = nxt.argmax(-1, keepdim=True)91 else:92 tok = torch.multinomial(torch.softmax(nxt / temperature, -1), 1, generator=generator)93 if eos_id is not None:94 tok = torch.where(done.unsqueeze(-1), torch.full_like(tok, eos_id), tok)95 done = done | (tok.squeeze(-1) == eos_id)96 seq = torch.cat([seq, tok], 1)97 if bool(done.all()):98 break99 return {"tokens": seq, "n_edits": torch.stack(edits, 1) if edits else None, "cache_bytes": 0}100101102@torch.no_grad()103def generate(model: PrismeModel, prompt: torch.Tensor, max_new_tokens: int, max_budget: int = 0,104 cost: Optional[torch.Tensor] = None, temperature: float = 0.0, eos_id: Optional[int] = None,105 forced_budget_actions: Optional[Callable[[int, int], torch.Tensor]] = None, generator=None,106 lambda_cost=None, kappa_risk=None):107 """Greedy (temperature 0) or sampled decoding. prompt: [B, T0]. Returns dict(tokens [B,T0+new], n_edits [B,new],108 actions list). Prompt positions are processed with the same per-token routing (prefill)."""109 cfg = model.cfg110 if not hasattr(model, "W_K") and getattr(cfg, "use_workspace", False):111 return generate_uncached(model, prompt, max_new_tokens, max_budget, temperature, eos_id, forced_budget_actions, generator)112 cache = RequestBatchCache(layers=model.backbone.new_cache())113 B, T0 = prompt.shape114 fa = forced_budget_actions(B, T0) if forced_budget_actions is not None else None115 logits, _, _ = step_logits(model, prompt, cache, max_budget, cost, fa, lambda_cost, kappa_risk)116 out = [prompt]117 edits, traces = [], []118 done = torch.zeros(B, dtype=torch.bool, device=prompt.device)119 nxt = logits[:, -1]120 for i in range(max_new_tokens):121 if temperature <= 0:122 tok = nxt.argmax(-1, keepdim=True)123 else:124 probs = torch.softmax(nxt / temperature, -1)125 tok = torch.multinomial(probs, 1, generator=generator)126 if eos_id is not None:127 tok = torch.where(done.unsqueeze(-1), torch.full_like(tok, eos_id), tok)128 done = done | (tok.squeeze(-1) == eos_id)129 out.append(tok)130 if i == max_new_tokens - 1 or bool(done.all()):131 break132 fa = forced_budget_actions(B, 1) if forced_budget_actions is not None else None133 logits, tr, ne = step_logits(model, tok, cache, max_budget, cost, fa, lambda_cost, kappa_risk)134 edits.append(ne[:, 0]); traces.append(tr)135 nxt = logits[:, -1]136 return {"tokens": torch.cat(out, 1), "n_edits": torch.stack(edits, 1) if edits else None, "cache_bytes": cache.nbytes()}137