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 shared = 2 * w * w + 2 * 2 * w + 6 * w * hidden68 per_action = 2 * w * w + 2 * 4 * w69 return shared + A * per_action707172def readout_ops(c: PrismeConfig) -> float:73 w, d, M = c.w, c.d_model, c.n_slots74 return 2 * w * w + 2 * (2 * M * w * w + 4 * M * w) + 2 * w * d757677def ops_per_token(c: PrismeConfig, T: int, n_actions: int, k: int, use_controller: bool = True,78 controller_steps: int | None = None) -> OpsBreakdown:79 """Ops for one decode step with `n_actions` executed edits of group size k (k = M for dense writes).80 controller_steps defaults to n_actions + 1 (one evaluation per executed action plus the one that STOPs)."""81 d, V = c.d_model, c.vocab_size82 bb_lin = 2 * backbone_nonembedding_params(c)83 bb_att = 4 * c.n_layers * T * d84 head = 2 * d * V85 if n_actions == 0:86 return OpsBreakdown(bb_lin, bb_att, head, 0, 0, 0, 0, 0, 0, 0)87 bank = 2 * 3 * d * c.w88 blk = editor_block_ops(c, k, T)89 nb = c.n_editor_blocks * n_actions90 gate = n_actions * 2 * k * 2 * c.w91 steps = (n_actions + 1 if controller_steps is None else controller_steps) if use_controller else 092 steps = min(steps, c.r_max)93 ctrl = steps * controller_step_ops(c)94 return OpsBreakdown(bb_lin, bb_att, head, bank, nb * blk["source_attn"], nb * blk["slot_attn"], nb * blk["ffn"],95 gate, ctrl, readout_ops(c))969798def action_cost_proxy(c: PrismeConfig, T: int) -> list[float]:99 """Known analytic cost of each nonterminal action (ops of one edit, incl. the next controller evaluation),100 normalised by the backbone-only ops/token at context T. Used as c(a) by the controller."""101 base = ops_per_token(c, T, 0, 1).backbone102 costs = []103 for k in [1] * c.n_slots + [2] * c.n_pairs:104 blk = editor_block_ops(c, k, T)105 e = c.n_editor_blocks * sum(blk.values()) + 2 * k * 2 * c.w + controller_step_ops(c)106 costs.append(e / base)107 return costs108