File size: 1,287 Bytes
e29f88c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | """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)
|