# SPDX-License-Identifier: Apache-2.0 # Original project extension; not part of the upstream LOOM snapshot. # Licensed under the Apache License, Version 2.0. # See source/LICENSE in this distribution for the license terms. # This distribution adds this license notice only; executable Python # syntax is unchanged from the original trained-model artifact. """Differentiable per-unit MELT memory and causal history selection.""" from typing import NamedTuple import torch import torch.nn.functional as F from torch import Tensor, nn class MeltGate(nn.Module): def __init__(self, width: int, *, device=None, dtype=None): super().__init__() self.input_weight = nn.Parameter(torch.zeros(width, width, device=device, dtype=dtype)) self.memory_weight = nn.Parameter(torch.zeros(width, width, device=device, dtype=dtype)) self.bias = nn.Parameter(torch.full((width,), -4.0, device=device, dtype=dtype)) def forward(self, current: Tensor, previous: Tensor) -> Tensor: retention = torch.sigmoid( F.linear(current, self.input_weight, self.bias) + F.linear(previous, self.memory_weight) ) return retention * previous + (1.0 - retention) * current class KV(NamedTuple): keys: Tensor values: Tensor class MeltHistory: """One dense equal-length LM unit; tensors stay connected across chunks. Container assignments commit immutable tensors, never in-place tensor writes. Final-loop history survives for the whole unit. SHARED's earlier-loop Python references retain only the recent window; autograd may retain older tensors. """ def __init__(self, layers: int, loops: int, policy: str, window_size: int, *, gate_enabled: bool = True): if policy not in ("full", "shared"): raise ValueError("Policy must be full or shared") if type(window_size) is not int or window_size < 0: raise ValueError("Window size must be a nonnegative integer") if layers < 1 or loops < 1: raise ValueError("Layers and loops must be positive") self.layers, self.loops = layers, loops self.policy, self.window_size = policy, window_size self.gate_enabled = gate_enabled self.history: list[list[KV | None]] = [[None] * loops for _ in range(layers)] self.length = 0 self.failed = False self._active = False self._batch_size: int | None = None def begin_chunk(self, batch_size: int, size: int) -> None: if self.failed or self._active: raise RuntimeError("History is failed or has an unfinished chunk") if size < 1 or size > self.window_size + 1: raise ValueError("Exact chunks require 1 <= C <= W+1") if self._batch_size is not None and batch_size != self._batch_size: raise ValueError("Unit batch size cannot change") self._batch_size = batch_size self._size = size self._next_visit = 0 self._memory: list[Tensor | None] = [None] * self.layers self._chunk: list[list[KV | None]] = [[None] * self.loops for _ in range(self.layers)] self._active = True def _check_visit(self, layer: int, loop: int) -> None: if self.failed or not self._active: raise RuntimeError("Begin a healthy chunk before layer visits") if self._next_visit != loop * self.layers + layer: raise RuntimeError("MELT layer/loop visits must follow the native loop order") def memory(self, layer: int, loop: int, current: Tensor, gate: MeltGate) -> Tensor: self._check_visit(layer, loop) if current.ndim != 3 or current.shape[:2] != (self._batch_size, self._size): raise ValueError("MELT requires a dense synchronous B x C x width input") previous = self._memory[layer] if loop == 0 or not self.gate_enabled: result = current else: if previous is None: raise RuntimeError("Previous-loop latent memory is missing") result = gate(current, previous) self._memory[layer] = result return result def attend(self, layer: int, loop: int, query: Tensor, key: Tensor, value: Tensor) -> Tensor: self._check_visit(layer, loop) current = KV(key, value) prefix = self.history[layer][loop] query_pos = torch.arange(self.length, self.length + self._size, device=query.device) current_pos = query_pos parts = [] masks = [] if self.policy == "full" or loop == self.loops - 1: if prefix is not None: parts.append(prefix) masks.append(torch.ones(self._size, prefix.keys.shape[1], device=query.device, dtype=torch.bool)) else: boundary = (query_pos - self.window_size).clamp_min(0) final = self.history[layer][-1] if final is not None: parts.append(final) positions = torch.arange(self.length, device=query.device) masks.append(positions[None, :] < boundary[:, None]) if prefix is not None: parts.append(prefix) positions = torch.arange(self.length - prefix.keys.shape[1], self.length, device=query.device) masks.append(positions[None, :] >= boundary[:, None]) parts.append(current) masks.append(current_pos[None, :] <= query_pos[:, None]) keys = torch.cat([part.keys for part in parts], dim=1).transpose(1, 2) values = torch.cat([part.values for part in parts], dim=1).transpose(1, 2) heads = query.shape[2] if heads % keys.shape[1]: raise ValueError("Query heads must be divisible by KV heads") repeats = heads // keys.shape[1] if repeats != 1: keys = keys.repeat_interleave(repeats, dim=1) values = values.repeat_interleave(repeats, dim=1) mask = torch.cat(masks, dim=-1)[None, None] output = F.scaled_dot_product_attention( query.transpose(1, 2), keys, values, attn_mask=mask, dropout_p=0.0, is_causal=False, ).transpose(1, 2) self._chunk[layer][loop] = current self._next_visit += 1 return output def commit_chunk(self) -> None: if self.failed or not self._active or self._next_visit != self.layers * self.loops: raise RuntimeError("Complete every native layer/loop visit before commit") for layer in range(self.layers): for loop in range(self.loops): previous, current = self.history[layer][loop], self._chunk[layer][loop] if current is None: raise RuntimeError("Missing current K/V") if previous is None: combined = current else: combined = KV(torch.cat((previous.keys, current.keys), dim=1), torch.cat((previous.values, current.values), dim=1)) if self.policy == "shared" and loop != self.loops - 1: start = max(0, combined.keys.shape[1] - self.window_size) combined = KV(combined.keys[:, start:].clone(), combined.values[:, start:].clone()) self.history[layer][loop] = combined self.length += self._size self._memory = [] self._chunk = [] self._active = False @property def live_kv_bytes(self) -> int: return sum(tensor.numel() * tensor.element_size() for layer in self.history for entry in layer if entry is not None for tensor in entry) @property def referenced_kv_storage_bytes(self) -> int: storages = {} for layer in self.history: for entry in layer: if entry is not None: for tensor in entry: storage = tensor.untyped_storage() storages[(str(tensor.device), storage.data_ptr())] = storage.nbytes() return sum(storages.values())