1"""Shared sparse editor: two shared-weight blocks acting on the 1-2 active slots of each position.23Each block: (1) active queries attend to the causally available source bank (K_t, V_t);4(2) active queries attend to all M slots of a coherent pre-write snapshot; (3) SwiGLU FFN.5Pre-RMSNorm + residual around each sublayer. After the blocks, the proposed vectors are turned6into clipped, gated writes; non-selected slots are returned bitwise unchanged.7"""8from __future__ import annotations910import math11from typing import Optional1213import torch14import torch.nn as nn15import torch.nn.functional as F1617from .attention import RMSNorm, SwiGLU18from .config import PrismeConfig19from .workspace import action_table, action_group_size, rms202122class EditorBlock(nn.Module):23 def __init__(self, cfg: PrismeConfig):24 super().__init__()25 w, h = cfg.w, cfg.n_editor_heads26 self.h, self.hd = h, w // h27 self.src_norm = RMSNorm(w, cfg.norm_eps)28 self.src_q = nn.Linear(w, w, bias=False)29 self.src_o = nn.Linear(w, w, bias=False)30 self.slot_norm = RMSNorm(w, cfg.norm_eps)31 self.slot_q = nn.Linear(w, w, bias=False)32 self.slot_k = nn.Linear(w, w, bias=False)33 self.slot_v = nn.Linear(w, w, bias=False)34 self.slot_o = nn.Linear(w, w, bias=False)35 self.ffn_norm = RMSNorm(w, cfg.norm_eps)36 self.ffn = SwiGLU(w, cfg.ffn_w)3738 def forward(self, x: torch.Tensor, snapshot: torch.Tensor, bank_k: torch.Tensor, bank_v: torch.Tensor,39 bank_mask: torch.Tensor) -> torch.Tensor:40 """x: [B, Tq, k, w] active queries; snapshot: [B, Tq, M, w]; bank_k/v: [B, Tb, w];41 bank_mask: [B, Tq, Tb] bool (True = key readable by that query position)."""42 B, Tq, k, w = x.shape43 Tb = bank_k.shape[1]44 # --- (1) source-bank cross-attention (causal) --------------------------------------45 q = self.src_q(self.src_norm(x)).view(B, Tq * k, self.h, self.hd).transpose(1, 2) # [B,h,Tq*k,hd]46 kk = bank_k.view(B, Tb, self.h, self.hd).transpose(1, 2) # [B,h,Tb,hd]47 vv = bank_v.view(B, Tb, self.h, self.hd).transpose(1, 2)48 m = bank_mask.unsqueeze(2).expand(B, Tq, k, Tb).reshape(B, 1, Tq * k, Tb)49 a = F.scaled_dot_product_attention(q, kk, vv, attn_mask=m)50 x = x + self.src_o(a.transpose(1, 2).reshape(B, Tq, k, w))51 # --- (2) attention over the full slot snapshot ---------------------------------------52 M = snapshot.shape[2]53 sn = self.slot_norm(snapshot)54 q = self.slot_q(self.slot_norm(x)).view(B * Tq, k, self.h, self.hd).transpose(1, 2) # [BT,h,k,hd]55 kk = self.slot_k(sn).view(B * Tq, M, self.h, self.hd).transpose(1, 2)56 vv = self.slot_v(sn).view(B * Tq, M, self.h, self.hd).transpose(1, 2)57 a = F.scaled_dot_product_attention(q, kk, vv)58 x = x + self.slot_o(a.transpose(1, 2).reshape(B, Tq, k, w))59 # --- (3) FFN ---------------------------------------------------------------------------60 x = x + self.ffn(self.ffn_norm(x))61 return x626364class SparseEditor(nn.Module):65 """Applies one action (per position) to the workspace. Weights of the block are shared across the66 cfg.n_editor_blocks sequential applications."""6768 N_CTX = 3 # step/R, remaining/R, k/26970 def __init__(self, cfg: PrismeConfig):71 super().__init__()72 self.cfg = cfg73 self.block = EditorBlock(cfg)74 self.ctx_proj = nn.Linear(self.N_CTX, cfg.w, bias=True)75 self.gate = nn.Linear(2 * cfg.w, 1, bias=True)76 self.register_buffer("actions", action_table(cfg), persistent=False)77 self.register_buffer("group_size", action_group_size(cfg), persistent=False)78 nn.init.zeros_(self.ctx_proj.weight)79 nn.init.zeros_(self.ctx_proj.bias)80 nn.init.zeros_(self.gate.weight)81 nn.init.zeros_(self.gate.bias) # eta starts at eta_max/28283 def active_slots(self, action: torch.Tensor):84 """action: [B, Tq] long -> (gidx [B,Tq,k] long >= 0, amask [B,Tq,k] bool)."""85 if self.cfg.dense_writes:86 M = self.cfg.n_slots87 gidx = torch.arange(M, device=action.device).view(1, 1, M).expand(*action.shape, M)88 amask = (action != self.cfg.stop_action).unsqueeze(-1).expand_as(gidx)89 return gidx, amask90 slots = self.actions[action] # [B,Tq,2]91 amask = slots >= 092 return slots.clamp(min=0), amask9394 def forward(self, S: torch.Tensor, bank_k: torch.Tensor, bank_v: torch.Tensor, bank_mask: torch.Tensor,95 action: torch.Tensor, step: torch.Tensor, remaining: torch.Tensor):96 """S: [B,Tq,M,w] current workspace; action: [B,Tq] in [0, n_actions]; step, remaining: [B,Tq] long.97 Returns (S_new [B,Tq,M,w], info dict). Positions with STOP are returned bitwise unchanged."""98 cfg = self.cfg99 B, Tq, M, w = S.shape100 gidx, amask = self.active_slots(action) # [B,Tq,k]101 k = gidx.shape[-1]102 idx = gidx.unsqueeze(-1).expand(B, Tq, k, w)103 cur = torch.gather(S, 2, idx) # [B,Tq,k,w] current values of active slots104 ksz = self.group_size[action].to(S.dtype) # [B,Tq]105 ctx = torch.stack([step.to(S.dtype) / cfg.r_max, remaining.to(S.dtype) / cfg.r_max, ksz / 2.0], -1) # [B,Tq,3]106 x = cur + self.ctx_proj(ctx).unsqueeze(2)107 snapshot = S # coherent pre-write snapshot for all blocks and both lanes of a pair108 for _ in range(cfg.n_editor_blocks):109 x = self.block(x, snapshot, bank_k, bank_v, bank_mask)110 proposed = x111 delta = proposed - cur112 r = rms(delta).unsqueeze(-1) # [B,Tq,k,1] (float32)113 scale = torch.clamp(cfg.clip_c / torch.clamp(r, min=cfg.clip_eps), max=1.0).to(delta.dtype)114 delta_c = delta * scale115 eta = cfg.eta_max * torch.sigmoid(self.gate(torch.cat([cur, proposed], -1))) # [B,Tq,k,1]116 new = cur + eta * delta_c117 # --- order-independent sparse write: one-hot scatter from the same snapshot ---------------118 onehot = F.one_hot(gidx, M).to(S.dtype) * amask.unsqueeze(-1).to(S.dtype) # [B,Tq,k,M]119 cand = torch.einsum("btkm,btkw->btmw", onehot, new) # [B,Tq,M,w]120 write_mask = (onehot.sum(2) > 0).unsqueeze(-1) # [B,Tq,M,1]121 S_new = torch.where(write_mask, cand, S)122 info = {"eta": eta.squeeze(-1), "delta_rms": r.squeeze(-1), "scale": scale.squeeze(-1),123 "active_mask": amask, "active_idx": gidx, "write_mask": write_mask.squeeze(-1)}124 return S_new, info125