"""Shared sparse editor: two shared-weight blocks acting on the 1-2 active slots of each position. Each block: (1) active queries attend to the causally available source bank (K_t, V_t); (2) active queries attend to all M slots of a coherent pre-write snapshot; (3) SwiGLU FFN. Pre-RMSNorm + residual around each sublayer. After the blocks, the proposed vectors are turned into clipped, gated writes; non-selected slots are returned bitwise unchanged. """ from __future__ import annotations import math from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from .attention import RMSNorm, SwiGLU from .config import PrismeConfig from .workspace import action_table, action_group_size, rms class EditorBlock(nn.Module): def __init__(self, cfg: PrismeConfig): super().__init__() w, h = cfg.w, cfg.n_editor_heads self.h, self.hd = h, w // h self.src_norm = RMSNorm(w, cfg.norm_eps) self.src_q = nn.Linear(w, w, bias=False) self.src_o = nn.Linear(w, w, bias=False) self.slot_norm = RMSNorm(w, cfg.norm_eps) self.slot_q = nn.Linear(w, w, bias=False) self.slot_k = nn.Linear(w, w, bias=False) self.slot_v = nn.Linear(w, w, bias=False) self.slot_o = nn.Linear(w, w, bias=False) self.ffn_norm = RMSNorm(w, cfg.norm_eps) self.ffn = SwiGLU(w, cfg.ffn_w) def forward(self, x: torch.Tensor, snapshot: torch.Tensor, bank_k: torch.Tensor, bank_v: torch.Tensor, bank_mask: torch.Tensor) -> torch.Tensor: """x: [B, Tq, k, w] active queries; snapshot: [B, Tq, M, w]; bank_k/v: [B, Tb, w]; bank_mask: [B, Tq, Tb] bool (True = key readable by that query position).""" B, Tq, k, w = x.shape Tb = bank_k.shape[1] # --- (1) source-bank cross-attention (causal) -------------------------------------- q = self.src_q(self.src_norm(x)).view(B, Tq * k, self.h, self.hd).transpose(1, 2) # [B,h,Tq*k,hd] kk = bank_k.view(B, Tb, self.h, self.hd).transpose(1, 2) # [B,h,Tb,hd] vv = bank_v.view(B, Tb, self.h, self.hd).transpose(1, 2) m = bank_mask.unsqueeze(2).expand(B, Tq, k, Tb).reshape(B, 1, Tq * k, Tb) a = F.scaled_dot_product_attention(q, kk, vv, attn_mask=m) x = x + self.src_o(a.transpose(1, 2).reshape(B, Tq, k, w)) # --- (2) attention over the full slot snapshot --------------------------------------- M = snapshot.shape[2] sn = self.slot_norm(snapshot) q = self.slot_q(self.slot_norm(x)).view(B * Tq, k, self.h, self.hd).transpose(1, 2) # [BT,h,k,hd] kk = self.slot_k(sn).view(B * Tq, M, self.h, self.hd).transpose(1, 2) vv = self.slot_v(sn).view(B * Tq, M, self.h, self.hd).transpose(1, 2) a = F.scaled_dot_product_attention(q, kk, vv) x = x + self.slot_o(a.transpose(1, 2).reshape(B, Tq, k, w)) # --- (3) FFN --------------------------------------------------------------------------- x = x + self.ffn(self.ffn_norm(x)) return x class SparseEditor(nn.Module): """Applies one action (per position) to the workspace. Weights of the block are shared across the cfg.n_editor_blocks sequential applications.""" N_CTX = 3 # step/R, remaining/R, k/2 def __init__(self, cfg: PrismeConfig): super().__init__() self.cfg = cfg self.block = EditorBlock(cfg) self.ctx_proj = nn.Linear(self.N_CTX, cfg.w, bias=True) self.gate = nn.Linear(2 * cfg.w, 1, bias=True) self.register_buffer("actions", action_table(cfg), persistent=False) self.register_buffer("group_size", action_group_size(cfg), persistent=False) nn.init.zeros_(self.ctx_proj.weight) nn.init.zeros_(self.ctx_proj.bias) nn.init.zeros_(self.gate.weight) nn.init.zeros_(self.gate.bias) # eta starts at eta_max/2 def active_slots(self, action: torch.Tensor): """action: [B, Tq] long -> (gidx [B,Tq,k] long >= 0, amask [B,Tq,k] bool).""" if self.cfg.dense_writes: M = self.cfg.n_slots gidx = torch.arange(M, device=action.device).view(1, 1, M).expand(*action.shape, M) amask = (action != self.cfg.stop_action).unsqueeze(-1).expand_as(gidx) return gidx, amask slots = self.actions[action] # [B,Tq,2] amask = slots >= 0 return slots.clamp(min=0), amask def forward(self, S: torch.Tensor, bank_k: torch.Tensor, bank_v: torch.Tensor, bank_mask: torch.Tensor, action: torch.Tensor, step: torch.Tensor, remaining: torch.Tensor): """S: [B,Tq,M,w] current workspace; action: [B,Tq] in [0, n_actions]; step, remaining: [B,Tq] long. Returns (S_new [B,Tq,M,w], info dict). Positions with STOP are returned bitwise unchanged.""" cfg = self.cfg B, Tq, M, w = S.shape gidx, amask = self.active_slots(action) # [B,Tq,k] k = gidx.shape[-1] idx = gidx.unsqueeze(-1).expand(B, Tq, k, w) cur = torch.gather(S, 2, idx) # [B,Tq,k,w] current values of active slots ksz = self.group_size[action].to(S.dtype) # [B,Tq] ctx = torch.stack([step.to(S.dtype) / cfg.r_max, remaining.to(S.dtype) / cfg.r_max, ksz / 2.0], -1) # [B,Tq,3] x = cur + self.ctx_proj(ctx).unsqueeze(2) snapshot = S # coherent pre-write snapshot for all blocks and both lanes of a pair for _ in range(cfg.n_editor_blocks): x = self.block(x, snapshot, bank_k, bank_v, bank_mask) proposed = x delta = proposed - cur r = rms(delta).unsqueeze(-1) # [B,Tq,k,1] (float32) scale = torch.clamp(cfg.clip_c / torch.clamp(r, min=cfg.clip_eps), max=1.0).to(delta.dtype) delta_c = delta * scale eta = cfg.eta_max * torch.sigmoid(self.gate(torch.cat([cur, proposed], -1))) # [B,Tq,k,1] new = cur + eta * delta_c # --- order-independent sparse write: one-hot scatter from the same snapshot --------------- onehot = F.one_hot(gidx, M).to(S.dtype) * amask.unsqueeze(-1).to(S.dtype) # [B,Tq,k,M] cand = torch.einsum("btkm,btkw->btmw", onehot, new) # [B,Tq,M,w] write_mask = (onehot.sum(2) > 0).unsqueeze(-1) # [B,Tq,M,1] S_new = torch.where(write_mask, cand, S) info = {"eta": eta.squeeze(-1), "delta_rms": r.squeeze(-1), "scale": scale.squeeze(-1), "active_mask": amask, "active_idx": gidx, "write_mask": write_mask.squeeze(-1)} return S_new, info