cache.py

"""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"]
            ),
        }