Skip to content

KV Cache

KeyValueCache and KeyValueCacheEntry provide small containers for autoregressive key/value activations. They are intentionally simple so model adapters can expose cache state without binding SafeLens to a specific transformers implementation.

Supported operations:

  • Lazy per-layer entry creation through cache[layer].
  • Sequence-axis append through append(layer, keys, values).
  • Best-effort sequence_length inference for tensor-like and nested-list data.
  • Serialization-friendly to_dict().

Example:

from SafeLens.core.kv_cache import KeyValueCache

cache = KeyValueCache()
cache.append(0, keys=[[[1]]], values=[[[2]]])
cache.append(0, keys=[[[3]]], values=[[[4]]])

assert cache[0].keys == [[[1], [3]]]
assert cache[0].sequence_length == 2

TransformerLens-compatible key/value cache containers.

KeyValueCache dataclass

Dictionary-like key/value cache keyed by layer index.

Source code in src/SafeLens/core/kv_cache.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
@dataclass
class KeyValueCache:
    """Dictionary-like key/value cache keyed by layer index."""

    entries: dict[int, KeyValueCacheEntry] | list[KeyValueCacheEntry] = field(default_factory=dict)
    previous_attention_mask: Any | None = None
    frozen: bool = False

    @classmethod
    def init_cache(
        cls,
        cfg: Any,
        device: Any,
        batch_size: int = 1,
    ) -> KeyValueCache:
        """Create an empty cache following TransformerLens' ``init_cache`` API."""
        import torch

        n_layers = getattr(cfg, "n_layers", None)
        if n_layers is None:
            n_layers = getattr(cfg, "n_layer", None)
        if n_layers is None:
            num_hidden_layers = getattr(cfg, "num_hidden_layers", None)
            n_layers = 0 if num_hidden_layers is None else num_hidden_layers
        device_for_mask = device if device is not None else getattr(cfg, "device", None)
        if device_for_mask is None:
            device_for_mask = torch.device("cpu")

        entries = [
            KeyValueCacheEntry.init_cache_entry(
                cfg,
                _device_for_cache_layer(cfg, device, layer_index),
                batch_size=batch_size,
            )
            for layer_index in range(int(n_layers))
        ]
        return cls(
            entries=entries,
            previous_attention_mask=torch.empty(
                (int(batch_size), 0),
                device=device_for_mask,
                dtype=torch.int,
            ),
        )

    def __getitem__(self, layer: int) -> KeyValueCacheEntry:
        if isinstance(self.entries, Mapping):
            if layer not in self.entries:
                self.entries[layer] = KeyValueCacheEntry(frozen=self.frozen)
            return self.entries[layer]
        return self.entries[layer]

    def append(self, layer: int, keys: Any, values: Any, *, dim: int = 1) -> tuple[Any, Any]:
        """Append key/value activations for one layer."""
        return self[layer].append(keys, values, dim=dim)

    def freeze(self) -> None:
        """Prevent future appends from mutating this cache."""
        self.frozen = True
        for entry in _iter_entries(self.entries):
            entry.frozen = True

    def unfreeze(self) -> None:
        """Allow future appends to mutate this cache."""
        self.frozen = False
        for entry in _iter_entries(self.entries):
            entry.frozen = False

    def append_attention_mask(self, attention_mask: Any) -> Any:
        """Append a batch attention mask and return the full mask."""
        previous_attention_mask = self.previous_attention_mask
        if previous_attention_mask is None:
            previous_attention_mask = _empty_attention_mask_like(attention_mask)
        updated_attention_mask = concat_values(
            previous_attention_mask,
            attention_mask,
            dim=-1,
        )
        if not self.frozen:
            self.previous_attention_mask = updated_attention_mask
        return updated_attention_mask

    def to_dict(self) -> dict[int, dict[str, Any]]:
        """Return a serializable view."""
        if isinstance(self.entries, Mapping):
            return {layer: entry.to_dict() for layer, entry in self.entries.items()}
        return {layer: entry.to_dict() for layer, entry in enumerate(self.entries)}

append(layer, keys, values, *, dim=1)

Append key/value activations for one layer.

Source code in src/SafeLens/core/kv_cache.py 
170
171
172
def append(self, layer: int, keys: Any, values: Any, *, dim: int = 1) -> tuple[Any, Any]:
    """Append key/value activations for one layer."""
    return self[layer].append(keys, values, dim=dim)

append_attention_mask(attention_mask)

Append a batch attention mask and return the full mask.

Source code in src/SafeLens/core/kv_cache.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
def append_attention_mask(self, attention_mask: Any) -> Any:
    """Append a batch attention mask and return the full mask."""
    previous_attention_mask = self.previous_attention_mask
    if previous_attention_mask is None:
        previous_attention_mask = _empty_attention_mask_like(attention_mask)
    updated_attention_mask = concat_values(
        previous_attention_mask,
        attention_mask,
        dim=-1,
    )
    if not self.frozen:
        self.previous_attention_mask = updated_attention_mask
    return updated_attention_mask

freeze()

Prevent future appends from mutating this cache.

Source code in src/SafeLens/core/kv_cache.py 
174
175
176
177
178
def freeze(self) -> None:
    """Prevent future appends from mutating this cache."""
    self.frozen = True
    for entry in _iter_entries(self.entries):
        entry.frozen = True

init_cache(cfg, device, batch_size=1) classmethod

Create an empty cache following TransformerLens' init_cache API.

Source code in src/SafeLens/core/kv_cache.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
@classmethod
def init_cache(
    cls,
    cfg: Any,
    device: Any,
    batch_size: int = 1,
) -> KeyValueCache:
    """Create an empty cache following TransformerLens' ``init_cache`` API."""
    import torch

    n_layers = getattr(cfg, "n_layers", None)
    if n_layers is None:
        n_layers = getattr(cfg, "n_layer", None)
    if n_layers is None:
        num_hidden_layers = getattr(cfg, "num_hidden_layers", None)
        n_layers = 0 if num_hidden_layers is None else num_hidden_layers
    device_for_mask = device if device is not None else getattr(cfg, "device", None)
    if device_for_mask is None:
        device_for_mask = torch.device("cpu")

    entries = [
        KeyValueCacheEntry.init_cache_entry(
            cfg,
            _device_for_cache_layer(cfg, device, layer_index),
            batch_size=batch_size,
        )
        for layer_index in range(int(n_layers))
    ]
    return cls(
        entries=entries,
        previous_attention_mask=torch.empty(
            (int(batch_size), 0),
            device=device_for_mask,
            dtype=torch.int,
        ),
    )

to_dict()

Return a serializable view.

Source code in src/SafeLens/core/kv_cache.py 
200
201
202
203
204
def to_dict(self) -> dict[int, dict[str, Any]]:
    """Return a serializable view."""
    if isinstance(self.entries, Mapping):
        return {layer: entry.to_dict() for layer, entry in self.entries.items()}
    return {layer: entry.to_dict() for layer, entry in enumerate(self.entries)}

unfreeze()

Allow future appends to mutate this cache.

Source code in src/SafeLens/core/kv_cache.py 
180
181
182
183
184
def unfreeze(self) -> None:
    """Allow future appends to mutate this cache."""
    self.frozen = False
    for entry in _iter_entries(self.entries):
        entry.frozen = False

KeyValueCacheEntry dataclass

Cache entry for one layer's key and value activations.

Source code in src/SafeLens/core/kv_cache.py 
 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
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
@dataclass(init=False)
class KeyValueCacheEntry:
    """Cache entry for one layer's key and value activations."""

    keys: Any | None = None
    values: Any | None = None
    frozen: bool = False

    def __init__(
        self,
        keys: Any | None = None,
        values: Any | None = None,
        *,
        past_keys: Any = _MISSING,
        past_values: Any = _MISSING,
        frozen: bool = False,
    ) -> None:
        if past_keys is not _MISSING:
            keys = past_keys
        if past_values is not _MISSING:
            values = past_values
        self.keys = keys
        self.values = values
        self.frozen = bool(frozen)

    @classmethod
    def init_cache_entry(
        cls,
        cfg: Any,
        device: Any,
        batch_size: int = 1,
    ) -> KeyValueCacheEntry:
        """Create an empty TL-layout entry shaped ``[batch, 0, heads, d_head]``."""
        import torch

        n_heads = getattr(cfg, "n_key_value_heads", None)
        if n_heads is None:
            n_heads = getattr(cfg, "num_key_value_heads", None)
        if n_heads is None:
            n_heads = getattr(cfg, "n_heads", None)
        if n_heads is None:
            n_heads = getattr(cfg, "n_head", None)
        if n_heads is None:
            raise AttributeError("KV cache initialization requires cfg.n_heads.")
        d_head = getattr(cfg, "d_head", None)
        if d_head is None:
            hidden_size = getattr(cfg, "d_model", None)
            if hidden_size is None:
                hidden_size = getattr(cfg, "n_embd", None)
            if hidden_size is None:
                hidden_size = getattr(cfg, "hidden_size", None)
            if hidden_size is None:
                raise AttributeError("KV cache initialization requires cfg.d_head.")
            d_head = int(hidden_size) // int(n_heads)
        dtype = _torch_dtype_from_config(cfg, torch)
        return cls(
            past_keys=torch.empty(
                (int(batch_size), 0, int(n_heads), int(d_head)),
                device=device,
                dtype=dtype,
            ),
            past_values=torch.empty(
                (int(batch_size), 0, int(n_heads), int(d_head)),
                device=device,
                dtype=dtype,
            ),
        )

    @property
    def past_keys(self) -> Any | None:
        """TransformerLens name for cached keys."""
        return self.keys

    @past_keys.setter
    def past_keys(self, value: Any | None) -> None:
        self.keys = value

    @property
    def past_values(self) -> Any | None:
        """TransformerLens name for cached values."""
        return self.values

    @past_values.setter
    def past_values(self, value: Any | None) -> None:
        self.values = value

    def append(self, keys: Any, values: Any, *, dim: int = 1) -> tuple[Any, Any]:
        """Append new key/value tensors along the sequence dimension."""
        updated_keys = concat_values(self.keys, keys, dim=dim)
        updated_values = concat_values(self.values, values, dim=dim)
        if not self.frozen:
            self.keys = updated_keys
            self.values = updated_values
        return updated_keys, updated_values

    @property
    def sequence_length(self) -> int:
        """Return cached sequence length when shape is available."""
        shape = shape_of(self.keys)
        return int(shape[1]) if len(shape) >= 2 else 0

    def to_dict(self) -> dict[str, Any]:
        """Return a serializable view."""
        return {"keys": self.keys, "values": self.values, "sequence_length": self.sequence_length}

past_keys property writable

TransformerLens name for cached keys.

past_values property writable

TransformerLens name for cached values.

sequence_length property

Return cached sequence length when shape is available.

append(keys, values, *, dim=1)

Append new key/value tensors along the sequence dimension.

Source code in src/SafeLens/core/kv_cache.py 
 98
 99
100
101
102
103
104
105
def append(self, keys: Any, values: Any, *, dim: int = 1) -> tuple[Any, Any]:
    """Append new key/value tensors along the sequence dimension."""
    updated_keys = concat_values(self.keys, keys, dim=dim)
    updated_values = concat_values(self.values, values, dim=dim)
    if not self.frozen:
        self.keys = updated_keys
        self.values = updated_values
    return updated_keys, updated_values

init_cache_entry(cfg, device, batch_size=1) classmethod

Create an empty TL-layout entry shaped [batch, 0, heads, d_head].

Source code in src/SafeLens/core/kv_cache.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
@classmethod
def init_cache_entry(
    cls,
    cfg: Any,
    device: Any,
    batch_size: int = 1,
) -> KeyValueCacheEntry:
    """Create an empty TL-layout entry shaped ``[batch, 0, heads, d_head]``."""
    import torch

    n_heads = getattr(cfg, "n_key_value_heads", None)
    if n_heads is None:
        n_heads = getattr(cfg, "num_key_value_heads", None)
    if n_heads is None:
        n_heads = getattr(cfg, "n_heads", None)
    if n_heads is None:
        n_heads = getattr(cfg, "n_head", None)
    if n_heads is None:
        raise AttributeError("KV cache initialization requires cfg.n_heads.")
    d_head = getattr(cfg, "d_head", None)
    if d_head is None:
        hidden_size = getattr(cfg, "d_model", None)
        if hidden_size is None:
            hidden_size = getattr(cfg, "n_embd", None)
        if hidden_size is None:
            hidden_size = getattr(cfg, "hidden_size", None)
        if hidden_size is None:
            raise AttributeError("KV cache initialization requires cfg.d_head.")
        d_head = int(hidden_size) // int(n_heads)
    dtype = _torch_dtype_from_config(cfg, torch)
    return cls(
        past_keys=torch.empty(
            (int(batch_size), 0, int(n_heads), int(d_head)),
            device=device,
            dtype=dtype,
        ),
        past_values=torch.empty(
            (int(batch_size), 0, int(n_heads), int(d_head)),
            device=device,
            dtype=dtype,
        ),
    )

to_dict()

Return a serializable view.

Source code in src/SafeLens/core/kv_cache.py 
113
114
115
def to_dict(self) -> dict[str, Any]:
    """Return a serializable view."""
    return {"keys": self.keys, "values": self.values, "sequence_length": self.sequence_length}

concat_values(old, new, *, dim=1)

Concatenate tensor-like or nested-list values.

Source code in src/SafeLens/core/kv_cache.py 
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
def concat_values(old: Any | None, new: Any, *, dim: int = 1) -> Any:
    """Concatenate tensor-like or nested-list values."""
    if old is None:
        return new
    try:
        import torch

        if isinstance(old, torch.Tensor) or isinstance(new, torch.Tensor):
            if not isinstance(old, torch.Tensor):
                old = torch.as_tensor(
                    old,
                    dtype=getattr(new, "dtype", None),
                    device=getattr(new, "device", None),
                )
            if not isinstance(new, torch.Tensor):
                new = torch.as_tensor(new, dtype=old.dtype, device=old.device)
            return torch.cat([old, new.to(device=old.device, dtype=old.dtype)], dim=dim)
    except Exception:
        pass
    try:
        import numpy as np

        if isinstance(old, np.ndarray) or isinstance(new, np.ndarray):
            return np.concatenate([np.asarray(old), np.asarray(new)], axis=dim)
    except Exception:
        pass
    if dim < 0:
        shape = shape_of(old)
        if shape:
            dim += len(shape)
    if dim == 0:
        return list(old) + list(new)
    return [
        concat_values(old_item, new_item, dim=dim - 1)
        for old_item, new_item in zip(old, new, strict=True)
    ]

shape_of(value)

Return best-effort shape.

Source code in src/SafeLens/core/kv_cache.py 
249
250
251
252
253
254
255
256
257
258
def shape_of(value: Any) -> tuple[int, ...]:
    """Return best-effort shape."""
    shape = getattr(value, "shape", None)
    if shape is not None:
        return tuple(int(dim) for dim in shape)
    if _is_sequence(value):
        if not value:
            return (0,)
        return (len(value), *shape_of(value[0]))
    return ()