1"""PrismeModel: backbone + projected source bank + ephemeral workspace + sparse editor + controller + readout.23Contracts implemented here (see docs/architecture_spec.md):4 * explicit bypass: max_budget == 0 (or actions None) returns the backbone head without touching the workspace;5 * causality: the bank visible to position t is 1..t; nothing from t+1 onward enters S, the editor or the controller;6 * locality: only the slots of the chosen action are written (bitwise identity elsewhere);7 * the editor never modifies H, K, V or z (they are inputs only);8 * the deployed controller scores actions from target-free features and never executes candidate edits.9"""10from __future__ import annotations1112from dataclasses import dataclass13from typing import Optional1415import torch16import torch.nn as nn17import torch.nn.functional as F1819from .attention import RMSNorm, LayerKVCache20from .backbone import Backbone21from .config import PrismeConfig22from .controller import ValueController23from .editor import SparseEditor24from .readout import DifferentialReadout25from .workspace import action_table, action_group_size, init_workspace262728@dataclass29class EncodedContext:30 """Everything the workspace needs for a batch of query positions. All tensors are read-only for the editor."""31 H: torch.Tensor # [B,Tq,d] backbone residual stream at the query positions32 K: torch.Tensor # [B,Tb,w] source bank keys (causally readable)33 V: torch.Tensor # [B,Tb,w]34 z: torch.Tensor # [B,Tq,w]35 bank_mask: torch.Tensor # [B,Tq,Tb] bool36 S0: torch.Tensor # [B,Tq,M,w]373839class PrismeModel(nn.Module):40 def __init__(self, cfg: PrismeConfig):41 super().__init__()42 cfg.validate()43 self.cfg = cfg44 self.backbone = Backbone(cfg)45 if cfg.use_workspace:46 self.bank_norm = RMSNorm(cfg.d_model, cfg.norm_eps)47 self.W_K = nn.Linear(cfg.d_model, cfg.w, bias=False)48 self.W_V = nn.Linear(cfg.d_model, cfg.w, bias=False)49 self.P = nn.Linear(cfg.d_model, cfg.w, bias=False)50 self.slot_embed = nn.Parameter(torch.randn(cfg.n_slots, cfg.w) * cfg.init_std)51 self.editor = SparseEditor(cfg)52 self.readout = DifferentialReadout(cfg)53 self.controller = ValueController(cfg)54 for lin in (self.W_K, self.W_V, self.P):55 nn.init.normal_(lin.weight, std=cfg.init_std)56 self.register_buffer("action_slots", action_table(cfg), persistent=False)57 self.register_buffer("action_k", action_group_size(cfg), persistent=False)5859 # ------------------------------------------------------------------ parameter groups60 def backbone_parameters(self):61 return list(self.backbone.parameters())6263 def workspace_parameters(self):64 names = ("bank_norm", "W_K", "W_V", "P", "slot_embed", "editor", "readout")65 return [p for n, p in self.named_parameters() if n.split(".")[0] in names]6667 def controller_parameters(self):68 return list(self.controller.parameters())6970 # ------------------------------------------------------------------ encoding71 def encode(self, tokens: torch.Tensor) -> EncodedContext:72 """Teacher-forced parallel encoding of all positions; bank mask is strictly causal (key j <= query t)."""73 H = self.backbone(tokens)74 return self.context_from_H(H, H)7576 def context_from_H(self, H_bank: torch.Tensor, H_query: torch.Tensor, query_pos: Optional[torch.Tensor] = None,77 bank_len: Optional[torch.Tensor] = None) -> EncodedContext:78 """H_bank: [B,Tb,d] history (incl. the query positions); H_query: [B,Tq,d].79 query_pos: [B,Tq] absolute positions of the queries inside the bank (default: arange).80 bank_len: [B] valid bank length per request (for padded batches)."""81 B, Tb, _ = H_bank.shape82 Tq = H_query.shape[1]83 hb = self.bank_norm(H_bank)84 K, V = self.W_K(hb), self.W_V(hb)85 z = self.P(self.bank_norm(H_query))86 if query_pos is None:87 query_pos = torch.arange(Tq, device=H_bank.device).view(1, Tq).expand(B, Tq)88 key_pos = torch.arange(Tb, device=H_bank.device)89 bank_mask = key_pos.view(1, 1, Tb) <= query_pos.unsqueeze(-1)90 if bank_len is not None:91 bank_mask = bank_mask & (key_pos.view(1, 1, Tb) < bank_len.view(B, 1, 1))92 S0 = init_workspace(self.slot_embed, z)93 return EncodedContext(H=H_query, K=K, V=V, z=z, bank_mask=bank_mask, S0=S0)9495 # ------------------------------------------------------------------ one action96 def apply_action(self, S: torch.Tensor, ctx: EncodedContext, action: torch.Tensor, step, remaining):97 B, Tq = action.shape98 if not torch.is_tensor(step):99 step = torch.full((B, Tq), int(step), device=S.device, dtype=torch.long)100 if not torch.is_tensor(remaining):101 remaining = torch.full((B, Tq), int(remaining), device=S.device, dtype=torch.long)102 return self.editor(S, ctx.K, ctx.V, ctx.bank_mask, action, step, remaining)103104 # ------------------------------------------------------------------ readout105 def logits_from_state(self, S: torch.Tensor, ctx: EncodedContext, edited: Optional[torch.Tensor] = None):106 """h_corrected = H + U(r(S) - r(S0)). If `edited` ([B,Tq] bool) is given, positions that were never107 edited use the backbone head directly (exact bypass)."""108 delta = self.readout(S, ctx.S0, ctx.z)109 if edited is not None:110 delta = torch.where(edited.unsqueeze(-1), delta, torch.zeros_like(delta))111 return self.backbone.head(ctx.H + delta), delta112113 # ------------------------------------------------------------------ forced-action forward (training / tests)114 def forward(self, tokens: torch.Tensor, actions: Optional[torch.Tensor] = None, return_state: bool = False):115 """tokens: [B,T]; actions: [B,T,R] long (STOP for unused steps) or None (-> exact bypass).116 Returns dict(logits=[B,T,V], ...)."""117 cfg = self.cfg118 if actions is None or not cfg.use_workspace or actions.shape[-1] == 0:119 H = self.backbone(tokens)120 return {"logits": self.backbone.head(H), "H": H, "n_edits": None}121 R = actions.shape[-1]122 assert R <= cfg.r_max, f"hard action cap exceeded: {R} > r_max={cfg.r_max}"123 ctx = self.encode(tokens)124 S = ctx.S0125 n_edits = torch.zeros(actions.shape[:2], dtype=torch.long, device=tokens.device)126 stopped = torch.zeros(actions.shape[:2], dtype=torch.bool, device=tokens.device)127 infos = []128 for r in range(R):129 a = actions[:, :, r]130 a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a) # STOP is absorbing131 nonstop = a != cfg.stop_action132 stopped = stopped | ~nonstop133 remaining = (R - r) - 1134 S, info = self.apply_action(S, ctx, a, r, remaining)135 n_edits = n_edits + nonstop.long()136 infos.append(info)137 edited = n_edits > 0138 logits, delta = self.logits_from_state(S, ctx, edited)139 out = {"logits": logits, "H": ctx.H, "n_edits": n_edits, "delta_out": delta}140 if return_state:141 out.update({"ctx": ctx, "S_final": S, "infos": infos})142 return out143144 # ------------------------------------------------------------------ controller-driven forward (teacher forced)145 @torch.no_grad()146 def forward_policy(self, tokens: torch.Tensor, max_budget: int, cost: torch.Tensor,147 lambda_cost: float | None = None, kappa_risk: float | None = None, return_state: bool = False):148 """Parallel teacher-forced rollout of the deployed controller at every position (evaluation only).149 cost: [A] analytic cost proxy per nonterminal action. Returns logits and the action trace [B,T,max_budget]."""150 cfg = self.cfg151 if max_budget == 0:152 H = self.backbone(tokens)153 return {"logits": self.backbone.head(H), "actions": None, "n_edits": torch.zeros(tokens.shape, dtype=torch.long, device=tokens.device)}154 max_budget = min(max_budget, cfg.r_max)155 ctx = self.encode(tokens)156 S = ctx.S0157 B, T = tokens.shape158 stopped = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)159 trace = torch.full((B, T, max_budget), cfg.stop_action, dtype=torch.long, device=tokens.device)160 n_edits = torch.zeros(B, T, dtype=torch.long, device=tokens.device)161 for r in range(max_budget):162 step = torch.full((B, T), r, device=tokens.device, dtype=torch.long)163 remaining = torch.full((B, T), max_budget - r - 1, device=tokens.device, dtype=torch.long)164 preds = self.controller(ctx.z, S, step, remaining)165 a = self.controller.select(preds, cost, lambda_cost=lambda_cost, kappa_risk=kappa_risk)166 a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)167 stopped = stopped | (a == cfg.stop_action)168 trace[:, :, r] = a169 S, _ = self.apply_action(S, ctx, a, step, remaining)170 n_edits = n_edits + (a != cfg.stop_action).long()171 if bool(stopped.all()):172 break173 logits, delta = self.logits_from_state(S, ctx, n_edits > 0)174 out = {"logits": logits, "actions": trace, "n_edits": n_edits}175 if return_state:176 out.update({"ctx": ctx, "S_final": S})177 return out178