1"""Analytic operation counts (2 x multiply-adds = "ops") per generated token at context length T.23This is an ARITHMETIC model of the reference implementation, not a hardware timing model. Conventions:4 * one multiply-add = 2 ops; softmax/normalisation/activation costs are ignored (they are O(dim), not O(dim^2));5 * decode-step accounting: the backbone attends over the whole context T (KV cache), the LM head is included;6 * source-bank projections are incremental (only the new position is projected);7 * the editor re-reads the full bank (T keys) for every active query of every block, and refreshes the slot K/V8 of all M slots in every block (reference implementation; an optimised kernel could cache unchanged slots);9 * controller: one shared state MLP per step + one w x w projection per candidate action (A = M + M/2);10 * readout: two attention reads (final and initial state) + the U projection.11"""12from __future__ import annotations1314from dataclasses import dataclass1516from .config import PrismeConfig171819@dataclass20class OpsBreakdown:21 backbone_linear: float22 backbone_attention: float23 lm_head: float24 bank_projection: float25 editor_source_attn: float26 editor_slot_attn: float27 editor_ffn: float28 editor_gate: float29 controller: float30 readout: float3132 @property33 def backbone(self) -> float:34 return self.backbone_linear + self.backbone_attention + self.lm_head3536 @property37 def workspace(self) -> float:38 return (self.bank_projection + self.editor_source_attn + self.editor_slot_attn + self.editor_ffn39 + self.editor_gate + self.controller + self.readout)4041 @property42 def total(self) -> float:43 return self.backbone + self.workspace4445 def as_dict(self) -> dict:46 d = self.__dict__.copy()47 d.update(backbone=self.backbone, workspace=self.workspace, total=self.total)48 return d495051def backbone_nonembedding_params(c: PrismeConfig) -> int:52 d, f, L, kv = c.d_model, c.ffn_width, c.n_layers, c.kv_dim53 return L * (2 * d * d + 2 * d * kv + 3 * d * f + 2 * d) + d545556def editor_block_ops(c: PrismeConfig, k: int, T: int) -> dict[str, float]:57 w, fw, M = c.w, c.ffn_w, c.n_slots58 return {59 "source_attn": 4 * k * w * w + 4 * k * T * w, # q/o projections + QK^T + AV over T bank entries60 "slot_attn": 4 * k * w * w + 4 * M * w * w + 4 * k * M * w, # q/o + K,V refresh of M slots + scores/AV61 "ffn": 6 * k * w * fw,62 }636465def controller_step_ops(c: PrismeConfig) -> float:66 w, hidden, A = c.w, (c.controller_hidden or c.ffn_w), c.n_actions67 din = c.d_model if c.ctrl_input == "h" else w68 shared = 2 * din * w + 2 * 2 * w + 6 * w * hidden + (2 * 2 * w if c.ctrl_use_entropy else 0)69 per_action = 2 * w * w + 2 * 4 * w70 return shared + A * per_action717273def controller_token_ops(c: PrismeConfig) -> float:74 """Once-per-token controller features: the backbone entropy needs one head evaluation (budget 0 would need it75 anyway only if the token stops immediately; we charge it fully, conservatively)."""76 return 2 * c.d_model * c.vocab_size if c.ctrl_use_entropy else 0.0777879def readout_ops(c: PrismeConfig) -> float:80 w, d, M = c.w, c.d_model, c.n_slots81 return 2 * w * w + 2 * (2 * M * w * w + 4 * M * w) + 2 * w * d828384def ops_per_token(c: PrismeConfig, T: int, n_actions: int, k: int, use_controller: bool = True,85 controller_steps: int | None = None) -> OpsBreakdown:86 """Ops for one decode step with `n_actions` executed edits of group size k (k = M for dense writes).87 controller_steps defaults to n_actions + 1 (one evaluation per executed action plus the one that STOPs)."""88 d, V = c.d_model, c.vocab_size89 bb_lin = 2 * backbone_nonembedding_params(c)90 bb_att = 4 * c.n_layers * T * d91 head = 2 * d * V92 if n_actions == 0:93 return OpsBreakdown(bb_lin, bb_att, head, 0, 0, 0, 0, 0, 0, 0)94 bank = 2 * 3 * d * c.w95 blk = editor_block_ops(c, k, T)96 nb = c.n_editor_blocks * n_actions97 gate = n_actions * 2 * k * 2 * c.w98 steps = (n_actions + 1 if controller_steps is None else controller_steps) if use_controller else 099 steps = min(steps, c.r_max)100 ctrl = steps * controller_step_ops(c) + controller_token_ops(c)101 return OpsBreakdown(bb_lin, bb_att, head, bank, nb * blk["source_attn"], nb * blk["slot_attn"], nb * blk["ffn"],102 gate, ctrl, readout_ops(c))103104105def action_cost_proxy(c: PrismeConfig, T: int) -> list[float]:106 """Known analytic cost of each nonterminal action (ops of one edit, incl. the next controller evaluation),107 normalised by the backbone-only ops/token at context T. Used as c(a) by the controller."""108 base = ops_per_token(c, T, 0, 1).backbone109 costs = []110 for k in [1] * c.n_slots + [2] * c.n_pairs:111 blk = editor_block_ops(c, k, T)112 e = c.n_editor_blocks * sum(blk.values()) + 2 * k * 2 * c.w + controller_step_ops(c)113 costs.append(e / base)114 return costs115116117def global_recurrence_step_ops(c: PrismeConfig, T: int, ffn_w: int | None = None) -> float:118 """One refinement step of the global-recurrence adapter (baselines/global_recurrence.py) for one position at119 context T: q/k/v/o projections (8w^2), causal attention over T positions (4Tw), SwiGLU (6 w f), plus the120 down/up projections charged once per token (4dw)."""121 w, f = c.w, (ffn_w or c.ffn_w)122 return 8 * w * w + 4 * T * w + 6 * w * f123124125def global_recurrence_token_ops(c: PrismeConfig) -> float:126 return 2 * 2 * c.d_model * c.w127