| |
| |
| |
| |
| |
| |
|
|
| """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()) |
|
|
