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 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))43 S = ctx.S044 R = min(max_budget, cfg.r_max)45 trace = torch.full((B, n, R), cfg.stop_action, dtype=torch.long, device=H.device)46 n_edits = torch.zeros(B, n, dtype=torch.long, device=H.device)47 stopped = torch.zeros(B, n, dtype=torch.bool, device=H.device)48 for r in range(R):49 step = torch.full((B, n), r, device=H.device, dtype=torch.long)50 remaining = torch.full((B, n), R - r - 1, device=H.device, dtype=torch.long)51 if forced_actions is not None:52 a = forced_actions[:, :, r]53 else:54 preds = model.controller(z, S, step, remaining)55 a = model.controller.select(preds, cost, lambda_cost=lambda_cost, kappa_risk=kappa_risk)56 a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)57 stopped = stopped | (a == cfg.stop_action)58 trace[:, :, r] = a59 S, _ = model.apply_action(S, ctx, a, step, remaining)60 n_edits += (a != cfg.stop_action).long()61 if bool(stopped.all()):62 break63 logits, _ = model.logits_from_state(S, ctx, n_edits > 0)64 return logits, trace, n_edits656667@torch.no_grad()68def generate(model: PrismeModel, prompt: torch.Tensor, max_new_tokens: int, max_budget: int = 0,69 cost: Optional[torch.Tensor] = None, temperature: float = 0.0, eos_id: Optional[int] = None,70 forced_budget_actions: Optional[Callable[[int, int], torch.Tensor]] = None, generator=None,71 lambda_cost=None, kappa_risk=None):72 """Greedy (temperature 0) or sampled decoding. prompt: [B, T0]. Returns dict(tokens [B,T0+new], n_edits [B,new],73 actions list). Prompt positions are processed with the same per-token routing (prefill)."""74 cfg = model.cfg75 cache = RequestBatchCache(layers=model.backbone.new_cache())76 B, T0 = prompt.shape77 fa = forced_budget_actions(B, T0) if forced_budget_actions is not None else None78 logits, _, _ = step_logits(model, prompt, cache, max_budget, cost, fa, lambda_cost, kappa_risk)79 out = [prompt]80 edits, traces = [], []81 done = torch.zeros(B, dtype=torch.bool, device=prompt.device)82 nxt = logits[:, -1]83 for i in range(max_new_tokens):84 if temperature <= 0:85 tok = nxt.argmax(-1, keepdim=True)86 else:87 probs = torch.softmax(nxt / temperature, -1)88 tok = torch.multinomial(probs, 1, generator=generator)89 if eos_id is not None:90 tok = torch.where(done.unsqueeze(-1), torch.full_like(tok, eos_id), tok)91 done = done | (tok.squeeze(-1) == eos_id)92 out.append(tok)93 if i == max_new_tokens - 1 or bool(done.all()):94 break95 fa = forced_budget_actions(B, 1) if forced_budget_actions is not None else None96 logits, tr, ne = step_logits(model, tok, cache, max_budget, cost, fa, lambda_cost, kappa_risk)97 edits.append(ne[:, 0]); traces.append(tr)98 nxt = logits[:, -1]99 return {"tokens": torch.cat(out, 1), "n_edits": torch.stack(edits, 1) if edits else None, "cache_bytes": cache.nbytes()}100