| """Per-forward, differentiable prompt memory for a single writable position.""" |
|
|
| from dataclasses import dataclass, field |
|
|
| import torch |
| from torch import Tensor |
|
|
|
|
| @dataclass |
| class PromptMemory: |
| """Capture layer inputs on pass one and reuse their non-working positions. |
| |
| write_mask is [B,S,1]; captured tensors are [B,S,H]. Never detach: later |
| losses must reach adapters that formed the prompt memory on the first pass. |
| Objects live for one model forward only, not across examples or updates. |
| This is an activation-memory reference path, not an attention KV cache. |
| """ |
|
|
| write_mask: Tensor |
| layers: list[Tensor] = field(default_factory=list) |
| output: Tensor | None = None |
| keys_values: list[tuple[Tensor, Tensor]] = field(default_factory=list) |
| layer_outputs: list[Tensor] = field(default_factory=list) |
|
|
| def layer_input(self, index: int, hidden: Tensor) -> Tensor: |
| if self.output is None: |
| self.layers.append(hidden) |
| return hidden |
| return torch.where(self.write_mask, hidden, self.layers[index]) |
|
|
| def block_output(self, hidden: Tensor) -> Tensor: |
| if self.output is None: |
| self.output = hidden |
| return hidden |
| return torch.where(self.write_mask, hidden, self.output) |
|
|
