"""Unified deep KV and runtime-state cache for recurrent Dendro execution.""" from __future__ import annotations from dataclasses import dataclass from typing import Any, Iterable, Iterator import torch try: # Optional: make Transformers generation recognize the cache class. from transformers.cache_utils import Cache as HFCacheBase except Exception: # pragma: no cover - local fallback. class HFCacheBase: # type: ignore[no-redef] pass @dataclass(slots=True) class QuantizedCacheTensor: """Per-vector symmetric int8 cache tensor.""" data: torch.Tensor scale: torch.Tensor @classmethod def from_tensor(cls, tensor: torch.Tensor) -> "QuantizedCacheTensor": scale = tensor.detach().abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / 127.0 data = torch.round(tensor.detach() / scale).clamp(-127, 127).to(torch.int8) return cls(data=data, scale=scale.to(dtype=torch.float16 if tensor.dtype != torch.float64 else torch.float32)) def dequantize( self, *, device: torch.device | str | None = None, dtype: torch.dtype | None = None, ) -> torch.Tensor: data = self.data.to(device=device) scale = self.scale.to(device=data.device) return (data.float() * scale.float()).to(dtype=dtype or torch.float32) def index_select(self, dim: int, index: torch.Tensor) -> "QuantizedCacheTensor": return QuantizedCacheTensor( data=self.data.index_select(dim, index.to(self.data.device)), scale=self.scale.index_select(dim, index.to(self.scale.device)), ) def to(self, device: torch.device | str) -> "QuantizedCacheTensor": return QuantizedCacheTensor(self.data.to(device), self.scale.to(device)) @property def shape(self) -> torch.Size: return self.data.shape CacheValue = torch.Tensor | QuantizedCacheTensor | None RUNTIME_STATE_NAMES = ( "workspace", "memory", "memory_scores", "plasticity_trace", "associative_keys", "associative_values", "associative_scores", "head_communication", "route_history", ) def _move_value(value: Any, device: torch.device | str) -> Any: if value is None: return None if isinstance(value, QuantizedCacheTensor): return value.to(device) if torch.is_tensor(value): return value.to(device) return value def _detach_value(value: Any) -> Any: if value is None: return None if isinstance(value, QuantizedCacheTensor): return QuantizedCacheTensor(value.data.detach(), value.scale.detach()) if torch.is_tensor(value): return value.detach() return value def _reorder_value(value: Any, beam_idx: torch.Tensor) -> Any: if value is None: return None if isinstance(value, QuantizedCacheTensor): return value.index_select(0, beam_idx) if torch.is_tensor(value): return value.index_select(0, beam_idx.to(value.device)) return value class DendroKVCache(HFCacheBase): """Per-recurrent-depth KV storage plus all non-parameter Dendro runtime state. The model has one physical recurrent cell, but each virtual depth sees a distinct hidden state. Correct autoregressive decoding therefore requires one K/V stream per recurrent depth. This cache stores those streams without introducing model parameters and also carries workspace, memory, plasticity and associative state. """ def __init__( self, *, max_depth: int, implementation: str = "dynamic", max_cache_length: int | None = None, sliding_window: int | None = None, offload_device: str = "cpu", detach_runtime_state: bool = True, ) -> None: # Deliberately do not call a version-specific HF Cache constructor. self.max_depth = int(max_depth) self.implementation = str(implementation) self.max_cache_length = None if max_cache_length is None else int(max_cache_length) self.sliding_window = None if sliding_window is None else int(sliding_window) self.offload_device = str(offload_device) self.detach_runtime_state = bool(detach_runtime_state) if self.max_depth <= 0: raise ValueError("max_depth must be positive") if self.implementation not in {"dynamic", "static", "sliding", "int8", "offloaded"}: raise ValueError(f"Unsupported cache implementation: {self.implementation}") if self.implementation == "static" and not self.max_cache_length: raise ValueError("static cache requires max_cache_length") if self.implementation == "sliding" and not (self.sliding_window or self.max_cache_length): raise ValueError("sliding cache requires sliding_window or max_cache_length") self.key_cache: list[CacheValue] = [None] * self.max_depth self.value_cache: list[CacheValue] = [None] * self.max_depth self.lengths: list[int] = [0] * self.max_depth self.recurrent_depth: int | None = None self.total_tokens_seen = 0 self.key_is_prefix: torch.Tensor | None = None self.key_positions: torch.Tensor | None = None self.key_attention_mask: torch.Tensor | None = None self.modality_ids: torch.Tensor | None = None self._last_sliding_indices: torch.Tensor | None = None # Sliding selection is computed once at depth zero and then reused by # every recurrent depth. Recording the pre-selection length lets later # depths validate alignment without reading a CUDA scalar via # ``indices.max().item()`` on every generated token. self._last_sliding_source_length: int | None = None # Runtime state is isolated by *virtual recurrent depth*. The model still # has one physical cell and one parameter source; these are activation/cache # tensors only. Per-depth state prevents a summary produced by depth N from # leaking into depth N+1 while processing the same training sequence. self._runtime_state_by_depth: dict[str, list[torch.Tensor | None]] = { name: [None] * self.max_depth for name in RUNTIME_STATE_NAMES } self.reasoning_state: dict[str, Any] = {} def get_runtime_state(self, name: str, depth_idx: int = 0) -> torch.Tensor | None: if name not in self._runtime_state_by_depth: raise KeyError(f"Unknown runtime state {name!r}") if not 0 <= int(depth_idx) < self.max_depth: raise IndexError(f"depth_idx={depth_idx} outside [0, {self.max_depth})") return self._runtime_state_by_depth[name][int(depth_idx)] # Read-only depth-zero compatibility properties. New architecture code should # use get_runtime_state(name, depth_idx), but these keep the public cache surface # convenient for diagnostics and older callers. @property def workspace(self) -> torch.Tensor | None: return self.get_runtime_state("workspace") @property def memory(self) -> torch.Tensor | None: return self.get_runtime_state("memory") @property def memory_scores(self) -> torch.Tensor | None: return self.get_runtime_state("memory_scores") @property def plasticity_trace(self) -> torch.Tensor | None: return self.get_runtime_state("plasticity_trace") @property def associative_keys(self) -> torch.Tensor | None: return self.get_runtime_state("associative_keys") @property def associative_values(self) -> torch.Tensor | None: return self.get_runtime_state("associative_values") @property def associative_scores(self) -> torch.Tensor | None: return self.get_runtime_state("associative_scores") @property def head_communication(self) -> torch.Tensor | None: return self.get_runtime_state("head_communication") @property def route_history(self) -> torch.Tensor | None: return self.get_runtime_state("route_history") @property def is_compileable(self) -> bool: # Mirrors the current Hugging Face Cache contract. The static backing # tensors have stable addresses after first allocation; all other modes # may resize or materialize and are therefore intentionally non-compileable. return self.implementation == "static" @property def is_initialized(self) -> bool: return any(value is not None for value in self.key_cache) @property def batch_size(self) -> int: """Return the cached batch size, or ``-1`` before first initialization.""" for metadata in ( self.key_attention_mask, self.key_positions, self.key_is_prefix, self.modality_ids, ): if metadata is not None: return int(metadata.shape[0]) for value in self.key_cache: if value is not None: return int(value.shape[0]) return -1 @property def max_batch_size(self) -> int: """Backward-compatible alias used by older Transformers releases.""" return self.batch_size @property def seen_tokens(self) -> int: """Backward-compatible absolute token counter. ``get_seq_length`` reports physically retained K/V length, while this counter continues increasing for bounded/sliding caches. """ return int(self.total_tokens_seen) @property def is_sliding(self) -> list[bool]: """Per-virtual-depth sliding markers expected by current HF cache code.""" return [self.implementation == "sliding"] * self.max_depth @property def is_linear(self) -> list[bool]: """Dendro uses attention K/V at every recurrent depth, not linear-attention layers.""" return [False] * self.max_depth def __len__(self) -> int: return self.max_depth def __iter__(self) -> Iterator[tuple[torch.Tensor, torch.Tensor] | None]: for idx in range(self.max_depth): if self.key_cache[idx] is None: yield None else: yield self[idx] def __getitem__(self, depth_idx: int) -> tuple[torch.Tensor, torch.Tensor]: key = self.key_cache[depth_idx] value = self.value_cache[depth_idx] if key is None or value is None: raise IndexError(f"No cache exists for recurrent depth {depth_idx}") return self._materialize(key), self._materialize(value) def _storage_device(self, incoming: torch.Tensor) -> torch.device: if self.implementation == "offloaded": return torch.device(self.offload_device) return incoming.device @staticmethod def _materialize( value: torch.Tensor | QuantizedCacheTensor, *, device: torch.device | str | None = None, dtype: torch.dtype | None = None, ) -> torch.Tensor: if isinstance(value, QuantizedCacheTensor): return value.dequantize(device=device, dtype=dtype) return value.to(device=device, dtype=dtype) if device is not None or dtype is not None else value def _store(self, tensor: torch.Tensor) -> torch.Tensor | QuantizedCacheTensor: if self.detach_runtime_state: tensor = tensor.detach() storage_device = self._storage_device(tensor) tensor = tensor.to(storage_device) if self.implementation == "int8": return QuantizedCacheTensor.from_tensor(tensor) return tensor def _append(self, previous: CacheValue, current: torch.Tensor, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor: if previous is None: return current previous_tensor = self._materialize(previous, device=device, dtype=dtype) return torch.cat([previous_tensor, current], dim=-2) def _update_metadata( self, *, is_prefix: torch.Tensor | None, positions: torch.Tensor | None, attention_mask: torch.Tensor | None, modality_ids: torch.Tensor | None, ) -> None: def append(old: torch.Tensor | None, new: torch.Tensor | None, fill: int | bool) -> torch.Tensor | None: if new is None: if old is None: return None batch = old.shape[0] new = torch.full((batch, 1), fill, device=old.device, dtype=old.dtype) if old is None: return new.detach() if self.detach_runtime_state else new return torch.cat([old.to(new.device), new], dim=-1) self.key_is_prefix = append(self.key_is_prefix, is_prefix, False) self.key_positions = append(self.key_positions, positions, 0) self.key_attention_mask = append(self.key_attention_mask, attention_mask, True) self.modality_ids = append(self.modality_ids, modality_ids, 0) def _sliding_keep_indices(self, device: torch.device) -> torch.Tensor: total = int(self.key_is_prefix.shape[-1]) if self.key_is_prefix is not None else self.lengths[0] window = int(self.sliding_window or self.max_cache_length or total) if total <= window: return torch.arange(total, device=device) if self.key_is_prefix is None: return torch.arange(total - window, total, device=device) # Prefix metadata should be identical across batch for packed generation. If # not, preserve every position marked prefix by at least one sample. prefix_any = self.key_is_prefix.any(dim=0) prefix_idx = torch.nonzero(prefix_any, as_tuple=False).flatten().to(device) text_idx = torch.nonzero(~prefix_any, as_tuple=False).flatten().to(device) text_keep = text_idx[-window:] return torch.unique(torch.cat([prefix_idx, text_keep]), sorted=True) def _apply_indices_to_metadata(self, indices: torch.Tensor) -> None: for name in ("key_is_prefix", "key_positions", "key_attention_mask", "modality_ids"): value = getattr(self, name) if value is not None: setattr(self, name, value.index_select(-1, indices.to(value.device))) def update( self, key_states: torch.Tensor, value_states: torch.Tensor, depth_idx: int, cache_kwargs: dict[str, Any] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Append or place K/V for one recurrent depth and return usable full K/V.""" if not 0 <= depth_idx < self.max_depth: raise IndexError(f"depth_idx={depth_idx} outside [0, {self.max_depth})") if key_states.shape != value_states.shape: raise ValueError("key_states and value_states must have identical shapes") kwargs = cache_kwargs or {} original_device, original_dtype = key_states.device, key_states.dtype query_len = int(key_states.shape[-2]) # Validate fixed-position writes before mutating either cache metadata or # the absolute token counter. A failed static write must leave the cache # transactionally unchanged so callers can recover or allocate a larger # cache without stale prefix/position rows. static_cache_position: torch.Tensor | None = None if self.implementation == "static": capacity = int(self.max_cache_length or 0) raw_cache_position = kwargs.get("cache_position") if raw_cache_position is None: start = self.lengths[depth_idx] raw_cache_position = torch.arange( start, start + query_len, device=key_states.device, ) static_cache_position = raw_cache_position.to( key_states.device, dtype=torch.long, ).reshape(-1) if static_cache_position.numel() != query_len: raise ValueError( "cache_position length must match the K/V query length: " f"{static_cache_position.numel()} != {query_len}" ) if static_cache_position.numel() and ( int(static_cache_position.min().item()) < 0 or int(static_cache_position.max().item()) >= capacity ): raise RuntimeError(f"Static cache capacity {capacity} exceeded") if depth_idx == 0: self._update_metadata( is_prefix=kwargs.get("is_prefix"), positions=kwargs.get("positions"), attention_mask=kwargs.get("attention_mask"), modality_ids=kwargs.get("modality_ids"), ) self.total_tokens_seen += query_len if self.implementation == "static": capacity = int(self.max_cache_length or 0) previous_key = self.key_cache[depth_idx] if previous_key is None: shape = list(key_states.shape) shape[-2] = capacity previous_key = torch.zeros(shape, device=key_states.device, dtype=key_states.dtype) previous_value = torch.zeros_like(previous_key) self.key_cache[depth_idx] = previous_key self.value_cache[depth_idx] = previous_value else: previous_value = self.value_cache[depth_idx] assert torch.is_tensor(previous_key) and torch.is_tensor(previous_value) assert static_cache_position is not None cache_position = static_cache_position previous_key.index_copy_(-2, cache_position, key_states) previous_value.index_copy_(-2, cache_position, value_states) self.lengths[depth_idx] = max(self.lengths[depth_idx], int(cache_position.max().item()) + 1) return ( previous_key[..., : self.lengths[depth_idx], :], previous_value[..., : self.lengths[depth_idx], :], ) combined_key = self._append( self.key_cache[depth_idx], key_states, device=original_device, dtype=original_dtype ) combined_value = self._append( self.value_cache[depth_idx], value_states, device=original_device, dtype=original_dtype ) if self.implementation == "sliding": if depth_idx == 0: self._last_sliding_source_length = int(combined_key.shape[-2]) window = int( self.sliding_window or self.max_cache_length or self._last_sliding_source_length ) if self._last_sliding_source_length <= window: # The common text-only decode case needs neither an index # tensor nor any K/V selection until its window is full. self._last_sliding_indices = None else: keep_indices = self._sliding_keep_indices(combined_key.device) # Immutable prefix tokens can make total length exceed the # text window even though every position is still retained. # In that case the sorted keep indices are exactly identity. if int(keep_indices.numel()) == self._last_sliding_source_length: self._last_sliding_indices = None else: self._last_sliding_indices = keep_indices self._apply_indices_to_metadata(keep_indices) indices = self._last_sliding_indices source_length = self._last_sliding_source_length if source_length is None: raise RuntimeError("Sliding cache metadata was not initialized at depth zero") # A newly introduced or misaligned depth cannot reconstruct already # discarded history. Shape comparison is sufficient because all # recurrent depths consume the same token layout, and unlike # ``indices.max().item()`` it introduces no GPU-to-host sync. if int(combined_key.shape[-2]) != source_length: raise RuntimeError("Changing recurrent depth during cached sliding generation is unsupported") if indices is not None: combined_key = combined_key.index_select(-2, indices.to(combined_key.device)) combined_value = combined_value.index_select(-2, indices.to(combined_value.device)) if self.max_cache_length and self.implementation not in {"sliding", "static"}: combined_key = combined_key[..., -self.max_cache_length :, :] combined_value = combined_value[..., -self.max_cache_length :, :] if depth_idx == 0: for name in ("key_is_prefix", "key_positions", "key_attention_mask", "modality_ids"): metadata = getattr(self, name) if metadata is not None: setattr(self, name, metadata[..., -self.max_cache_length :]) self.lengths[depth_idx] = int(combined_key.shape[-2]) self.key_cache[depth_idx] = self._store(combined_key) self.value_cache[depth_idx] = self._store(combined_value) return ( self._materialize(self.key_cache[depth_idx], device=original_device, dtype=original_dtype), self._materialize(self.value_cache[depth_idx], device=original_device, dtype=original_dtype), ) def get_seq_length(self, layer_idx: int = 0) -> int: layer_idx = int(layer_idx) if not 0 <= layer_idx < self.max_depth: return 0 return self.lengths[layer_idx] if self.lengths else 0 def get_max_length(self, layer_idx: int | None = None) -> int: """Return the current Hugging Face maximum-cache contract. ``-1`` means unbounded/undefined. A sliding cache reports its text window; Dendro may retain an immutable multimodal prefix in addition to that window, which is represented explicitly by cache metadata. """ if layer_idx is not None and not 0 <= int(layer_idx) < self.max_depth: return -1 if self.implementation == "sliding": return int(self.sliding_window or self.max_cache_length or -1) if self.implementation == "static": return int(self.max_cache_length or -1) return int(self.max_cache_length) if self.max_cache_length is not None else -1 def get_max_cache_shape(self, layer_idx: int = 0) -> int: """Compatibility alias retained for Transformers versions before 5.16.""" return self.get_max_length(layer_idx) def get_mask_sizes(self, query_length: int, layer_idx: int = 0) -> tuple[int, int]: """Return usable key length and physical offset for HF mask builders. Dendro constructs its authoritative multimodal/prefix mask internally; this method makes the custom cache conform to the public HF Cache API. """ query_length = int(query_length) previous = self.get_seq_length(layer_idx) if self.implementation == "static": return int(self.max_cache_length or previous + query_length), 0 retained_after_update = previous + query_length maximum = self.get_max_length(layer_idx) if maximum > 0 and self.implementation != "sliding": retained_after_update = min(retained_after_update, maximum) offset = max(0, self.total_tokens_seen - previous) return retained_after_update, offset def get_query_offset(self, layer_idx: int = 0) -> int: """Absolute next-query offset used by current Transformers generation.""" del layer_idx return int(self.total_tokens_seen) def get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int: previous = self.get_seq_length(layer_idx) maximum = self.get_max_length(layer_idx) if maximum < 0: return previous return max(0, min(previous, maximum - int(new_seq_length))) def ensure_recurrent_depth(self, recurrent_depth: int) -> None: recurrent_depth = int(recurrent_depth) if recurrent_depth > self.max_depth: raise ValueError(f"Cache max_depth={self.max_depth} cannot hold recurrent depth {recurrent_depth}") if self.recurrent_depth is None: self.recurrent_depth = recurrent_depth elif self.recurrent_depth != recurrent_depth and self.get_seq_length() > 0: raise RuntimeError( "reasoning effort/recurrent depth changed while reusing a populated cache; " "start a fresh cache or keep the same effort" ) def set_runtime_state( self, name: str, value: torch.Tensor | None, depth_idx: int = 0, ) -> None: if name not in self._runtime_state_by_depth: raise KeyError(f"Unknown runtime state {name!r}") if not 0 <= int(depth_idx) < self.max_depth: raise IndexError(f"depth_idx={depth_idx} outside [0, {self.max_depth})") if value is not None and self.detach_runtime_state: value = value.detach() self._runtime_state_by_depth[name][int(depth_idx)] = value def reorder_cache(self, beam_idx: torch.Tensor) -> "DendroKVCache": self.key_cache = [_reorder_value(value, beam_idx) for value in self.key_cache] self.value_cache = [_reorder_value(value, beam_idx) for value in self.value_cache] for name in ("key_is_prefix", "key_positions", "key_attention_mask", "modality_ids"): setattr(self, name, _reorder_value(getattr(self, name), beam_idx)) for name, values in self._runtime_state_by_depth.items(): self._runtime_state_by_depth[name] = [_reorder_value(value, beam_idx) for value in values] return self def batch_repeat_interleave(self, repeats: int) -> "DendroKVCache": if repeats <= 0: raise ValueError("repeats must be positive") if self.key_cache[0] is None: return self device = self._materialize(self.key_cache[0]).device batch = self._materialize(self.key_cache[0]).shape[0] index = torch.arange(batch, device=device).repeat_interleave(repeats) return self.reorder_cache(index) def batch_select_indices(self, indices: torch.Tensor) -> "DendroKVCache": return self.reorder_cache(indices) def activate_past_recording(self) -> "DendroKVCache": """Advertise rollback-capable history for assisted generation. Current Transformers releases call this hook on custom cache objects when speculative/assisted decoding may crop rejected draft tokens. Dendro's dynamic, static, int8, and offloaded caches already retain the required history and :meth:`crop` rolls every recurrent-depth K/V stream back in lockstep. The marker is diagnostic only and introduces no tensor storage. A bounded sliding cache cannot resurrect tokens that it has deliberately evicted, but it can still crop retained draft positions. That limitation is recorded in ``reasoning_state`` rather than silently changing the cache implementation. """ self.reasoning_state["past_recording_active"] = True self.reasoning_state["past_recording_complete_history"] = self.implementation != "sliding" return self def crop(self, max_length: int) -> None: if max_length < 0: max_length = max(0, self.get_seq_length() + max_length) for depth_idx in range(self.max_depth): key, value = self.key_cache[depth_idx], self.value_cache[depth_idx] if key is None or value is None: continue if isinstance(key, QuantizedCacheTensor): self.key_cache[depth_idx] = QuantizedCacheTensor( key.data[..., :max_length, :], key.scale[..., :max_length, :] ) assert isinstance(value, QuantizedCacheTensor) self.value_cache[depth_idx] = QuantizedCacheTensor( value.data[..., :max_length, :], value.scale[..., :max_length, :] ) else: assert torch.is_tensor(value) self.key_cache[depth_idx] = key[..., :max_length, :] self.value_cache[depth_idx] = value[..., :max_length, :] self.lengths[depth_idx] = min(self.lengths[depth_idx], max_length) for name in ("key_is_prefix", "key_positions", "key_attention_mask", "modality_ids"): metadata = getattr(self, name) if metadata is not None: setattr(self, name, metadata[..., :max_length]) self._last_sliding_indices = None self._last_sliding_source_length = None def detach(self) -> "DendroKVCache": self.key_cache = [_detach_value(value) for value in self.key_cache] self.value_cache = [_detach_value(value) for value in self.value_cache] for name in ("key_is_prefix", "key_positions", "key_attention_mask", "modality_ids"): setattr(self, name, _detach_value(getattr(self, name))) for name, values in self._runtime_state_by_depth.items(): self._runtime_state_by_depth[name] = [_detach_value(value) for value in values] return self def to(self, device: torch.device | str) -> "DendroKVCache": self.key_cache = [_move_value(value, device) for value in self.key_cache] self.value_cache = [_move_value(value, device) for value in self.value_cache] for name in ("key_is_prefix", "key_positions", "key_attention_mask", "modality_ids"): setattr(self, name, _move_value(getattr(self, name), device)) for name, values in self._runtime_state_by_depth.items(): self._runtime_state_by_depth[name] = [_move_value(value, device) for value in values] return self def reset(self) -> None: self.key_cache = [None] * self.max_depth self.value_cache = [None] * self.max_depth self.lengths = [0] * self.max_depth self.recurrent_depth = None self.total_tokens_seen = 0 for name in ("key_is_prefix", "key_positions", "key_attention_mask", "modality_ids"): setattr(self, name, None) self._runtime_state_by_depth = { name: [None] * self.max_depth for name in RUNTIME_STATE_NAMES } self.reasoning_state = {} self._last_sliding_indices = None self._last_sliding_source_length = None def to_legacy_cache(self) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: result: list[tuple[torch.Tensor, torch.Tensor]] = [] for depth in range(self.recurrent_depth or self.max_depth): if self.key_cache[depth] is None: break result.append(self[depth]) return tuple(result) @classmethod def from_legacy_cache( cls, legacy: Iterable[tuple[torch.Tensor, torch.Tensor]], *, max_depth: int | None = None, implementation: str = "dynamic", ) -> "DendroKVCache": rows = list(legacy) cache = cls(max_depth=max_depth or max(1, len(rows)), implementation=implementation) for depth, (key, value) in enumerate(rows): cache.key_cache[depth] = cache._store(key) cache.value_cache[depth] = cache._store(value) cache.lengths[depth] = int(key.shape[-2]) cache.recurrent_depth = len(rows) return cache def memory_bytes(self) -> int: total = 0 values: list[Any] = [*self.key_cache, *self.value_cache] values.extend(getattr(self, name) for name in ("key_is_prefix", "key_positions", "key_attention_mask", "modality_ids")) for runtime_values in self._runtime_state_by_depth.values(): values.extend(runtime_values) seen_storage: set[tuple[int, int]] = set() for value in values: if isinstance(value, QuantizedCacheTensor): for tensor in (value.data, value.scale): key = (tensor.untyped_storage().data_ptr(), tensor.untyped_storage().nbytes()) if key not in seen_storage: total += tensor.untyped_storage().nbytes() seen_storage.add(key) elif torch.is_tensor(value): key = (value.untyped_storage().data_ptr(), value.untyped_storage().nbytes()) if key not in seen_storage: total += value.untyped_storage().nbytes() seen_storage.add(key) return total def summary(self) -> dict[str, Any]: return { "implementation": self.implementation, "max_depth": self.max_depth, "recurrent_depth": self.recurrent_depth, "sequence_length": self.get_seq_length(), "total_tokens_seen": self.total_tokens_seen, "memory_bytes": self.memory_bytes(), "runtime_depths_initialized": sum( any(value is not None for value in self._runtime_state_by_depth[name]) for name in RUNTIME_STATE_NAMES ), "has_workspace": any(value is not None for value in self._runtime_state_by_depth["workspace"]), "has_memory": any(value is not None for value in self._runtime_state_by_depth["memory"]), "has_plasticity": any( value is not None for value in self._runtime_state_by_depth["plasticity_trace"] ), "has_associative_memory": any( value is not None for value in self._runtime_state_by_depth["associative_keys"] ), }