From 6f4b12b5031032e4845527cec7f60eed7d1004d1 Mon Sep 17 00:00:00 2001
From: Ziyang Guo <121015044+RerankerGuo@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:56:09 +0800
Subject: [PATCH] fix(activation): keep the stored KV cache read-only during
 generation (#2301)

The activation memory stored in `KVCacheMemory` was handed to the model as
the very same object, and `transformers` appends the newly generated
key/value tensors to that object **in place** (`DynamicLayer.update`
rebinds `keys`/`values` on the same layer object).

So the stored activation memory grew with the prompt and the generated
tokens of every turn, and `ActivationMemoryManager` wrote the growth back
to disk on the next dump. The stored cache was supposed to represent a
fixed set of memories, not the conversation.

Three hops were involved:

1. `MemOS.chat` (`mem_os/core.py`), `mem_os/main.py`, `mem_chat/simple.py`
   and the scheduler analyzer pass the stored `kv_cache.memory` straight
   into `HFLLM.generate(past_key_values=...)`;
2. `HFLLM._prefill` forwards it to `model(past_key_values=...)`, and
   transformers appends the new tokens' K/V in place;
3. `KVCacheMemory._concat_caches` returns `caches[0]` unchanged for a
   single id, so the `get_cache()` merge path hands out the stored object
   too.

This commit treats stored caches as read-only by construction, at the two
boundaries where a cache leaves the store or enters the model:

- `clone_dynamic_cache()` (new, `memories/activation/kv.py`): an
  independent copy with cloned K/V tensors. The modern layout
  (`layers`, transformers >= 4.55) is preferred whenever the attribute
  exists, so that the copy never shares its layer list with the
  original, even for an empty cache that the model fills with its first
  layer. The legacy layout (`key_cache`/`value_cache` lists, transformers
  < 4.55) is copied as well. On a layer object, `keys`/`values` and
  `key_cache`/`value_cache` are copied independently of each other, so a
  layer that populates only one side does not lose data. The layer
  objects are shallow-copied, so all non-tensor state (for example the
  initialization flag that `DynamicLayer.get_seq_length()` relies on) is
  preserved.
- `KVCacheMemory._concat_caches`: the single-cache case now returns a
  clone instead of the stored object.
- `HFLLM.generate` / `HFLLM.generate_stream`: clone the incoming
  `past_key_values` once at the boundary. This covers all four call
  sites without touching them; the cost is one cache copy per generate
  call, the same order as the prefill work itself.

No production dependency, configuration, schema or API change.

Tests:

- `tests/llms/test_hf.py`: `test_generate_does_not_mutate_caller_cache`
  and `test_generate_stream_does_not_mutate_caller_cache` use a stand-in
  model that appends K/V in place, and assert that the model receives a
  copy and that the caller's cache still holds only the original tokens.
  Both fail before the fix (the model receives the caller's cache
  itself) and pass after it.
- `tests/memories/activation/test_kv.py`:
  `test_clone_dynamic_cache_does_not_share_tensors`,
  `test_get_cache_single_item_returns_independent_copy`,
  `test_get_cache_merge_does_not_alias_inputs`,
  `test_clone_dynamic_cache_handles_empty_cache`,
  `test_clone_dynamic_cache_prefers_the_layers_layout`,
  `test_clone_dynamic_cache_handles_legacy_layout` and
  `test_clone_dynamic_cache_copies_layers_with_legacy_attribute_names`
  cover the store boundary and both cache layouts.
- `tests/utils.py`: shared helpers to build a cache and to read it
  independently of the cache layout.

The layout-independent helpers also replace the `key_cache`-only helper in
`tests/memories/activation/test_kv.py`, which cannot build a populated
cache at all with transformers >= 4.55.

The two touched test modules now import torch through
`pytest.importorskip("torch")`, so that a test environment without torch
skips them instead of failing during collection.

Fixes #2301.
---
 src/memos/llms/hf.py                 |  33 ++++-
 src/memos/memories/activation/kv.py  |  78 +++++++++++-
 tests/llms/test_hf.py                |  84 ++++++++++++-
 tests/memories/activation/test_kv.py | 174 +++++++++++++++++++++++++--
 tests/utils.py                       |  82 +++++++++++++
 5 files changed, 438 insertions(+), 13 deletions(-)

diff --git a/src/memos/llms/hf.py b/src/memos/llms/hf.py
index 0dd841c1a..ad0363187 100644
--- a/src/memos/llms/hf.py
+++ b/src/memos/llms/hf.py
@@ -82,7 +82,7 @@ def generate(
         if past_key_values is None:
             return self._generate_full(prompt, **kwargs)
         else:
-            return self._generate_with_cache(prompt, past_key_values, **kwargs)
+            return self._generate_with_cache(prompt, self._clone_cache(past_key_values), **kwargs)
 
     def generate_stream(
         self, messages: MessageList, past_key_values: DynamicCache | None = None, **kwargs
@@ -102,7 +102,36 @@ def generate_stream(
         if past_key_values is None:
             yield from self._generate_full_stream(prompt)
         else:
-            yield from self._generate_with_cache_stream(prompt, past_key_values)
+            yield from self._generate_with_cache_stream(prompt, self._clone_cache(past_key_values))
+
+    @staticmethod
+    def _clone_cache(cache: DynamicCache) -> DynamicCache:
+        """
+        Copy a KV cache before it is handed to the model.
+
+        ``transformers`` appends the generated key/value tensors to the cache
+        object in place, so the cache that is stored in the activation memory
+        would otherwise grow with every turn and the growth would be written
+        back to disk on the next dump.
+
+        The copy is the price of keeping the stored cache untouched: one extra
+        copy of the activation cache is alive while the model generates. That
+        is the same order as the prefill work itself, and it is preferred over
+        cropping the stored cache back to its original length afterwards,
+        which would mutate the stored cache during the generation and leave
+        it grown whenever the generation is interrupted.
+
+        Args:
+            cache (DynamicCache): The cache provided by the caller.
+
+        Returns:
+            DynamicCache: A copy that does not share any tensor with ``cache``.
+        """
+        # Imported lazily: memos.llms.factory imports this module, so a
+        # module-level import of memos.memories.activation.kv would be circular.
+        from memos.memories.activation.kv import clone_dynamic_cache
+
+        return clone_dynamic_cache(cache)
 
     def _generate_full(self, prompt: str, **kwargs) -> str:
         """
diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py
index 1981b958f..166f61c90 100644
--- a/src/memos/memories/activation/kv.py
+++ b/src/memos/memories/activation/kv.py
@@ -1,7 +1,9 @@
+import copy
 import os
 import pickle
 
 from datetime import datetime
+from typing import Any
 
 from transformers import DynamicCache
 
@@ -206,7 +208,9 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
 
         assert caches, "Need at least one cache"
         if len(caches) == 1:
-            return caches[0]
+            # Hand out an independent copy: the stored cache must stay read-only,
+            # otherwise the model appends the generated tokens to it in place.
+            return clone_dynamic_cache(caches[0])
 
         merged = DynamicCache()
 
@@ -255,6 +259,78 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
         return merged
 
 
+def _clone_tensor(tensor: Any) -> Any:
+    """Clone a cache tensor, keeping ``None`` as ``None``."""
+    if tensor is None:
+        return None
+    return tensor.clone()
+
+
+def _clone_cache_layer(layer: Any) -> Any:
+    """
+    Copy a single cache layer without sharing any tensor with the original.
+
+    The layer object itself is shallow-copied so that all non-tensor state
+    (e.g. the cumulative sequence length tracked by newer ``transformers``
+    versions) is preserved, while every key/value tensor is replaced by an
+    independent clone.
+    """
+    cloned = copy.copy(layer)
+    for attribute in ("keys", "key_cache", "values", "value_cache"):
+        tensor = getattr(layer, attribute, None)
+        if tensor is not None:
+            setattr(cloned, attribute, _clone_tensor(tensor))
+    return cloned
+
+
+def clone_dynamic_cache(cache: DynamicCache) -> DynamicCache:
+    """
+    Return an independent copy of a ``DynamicCache``.
+
+    ``transformers`` appends the newly generated key/value tensors **in place**
+    to the cache object it is given (``DynamicLayer.update`` rebinds
+    ``keys``/``values`` on the very same layer object). When the cache handed
+    to the model is the object stored in the activation memory, every chat turn
+    permanently grows the stored activation memory, and the growth is written
+    back to disk on the next dump.
+
+    The returned cache shares no tensor with the original, so it can safely be
+    passed to the model. Both cache layouts are supported:
+
+    * ``transformers < 4.55``: ``key_cache`` / ``value_cache`` lists;
+    * ``transformers >= 4.55``: a ``layers`` list whose layers carry either
+      ``keys`` / ``values`` or ``key_cache`` / ``value_cache``.
+
+    Args:
+        cache: The cache to copy.
+
+    Returns:
+        DynamicCache: A copy that does not share any tensor with ``cache``.
+    """
+    layers = getattr(cache, "layers", None)
+    if layers is not None:
+        # The modern layout. The check is on the attribute itself and not on
+        # its content: an empty cache still has to be copied into an
+        # independent one, because the model appends its first layer to the
+        # cache it receives.
+        cloned = copy.copy(cache)
+        cloned.layers = [_clone_cache_layer(layer) for layer in layers]
+        return cloned
+
+    key_cache = getattr(cache, "key_cache", None)
+    value_cache = getattr(cache, "value_cache", None)
+    if key_cache is not None or value_cache is not None:
+        # The legacy layout, where the tensors live in two flat lists.
+        cloned = copy.copy(cache)
+        if key_cache is not None:
+            cloned.key_cache = [_clone_tensor(tensor) for tensor in key_cache]
+        if value_cache is not None:
+            cloned.value_cache = [_clone_tensor(tensor) for tensor in value_cache]
+        return cloned
+
+    raise AttributeError("DynamicCache object has neither 'layers' nor 'key_cache' attributes")
+
+
 def move_dynamic_cache_htod(dynamic_cache: DynamicCache, device: str) -> DynamicCache:
     """
     Move DynamicCache from CPU to GPU device.
diff --git a/tests/llms/test_hf.py b/tests/llms/test_hf.py
index 375bf2247..d5265e2f5 100644
--- a/tests/llms/test_hf.py
+++ b/tests/llms/test_hf.py
@@ -2,13 +2,46 @@
 
 from unittest.mock import MagicMock, patch
 
-import torch
+import pytest
 
 from transformers import DynamicCache
 
 from memos.configs.llm import HFLLMConfig, LLMConfigFactory
 from memos.llms.factory import LLMFactory
 from memos.llms.hf import HFLLM
+from tests.utils import dynamic_cache_seq_len, make_dynamic_cache
+
+
+# The KV cache is a torch structure; skip the whole module where torch is
+# not installed instead of failing the collection of the test suite.
+torch = pytest.importorskip("torch")
+
+
+class CacheGrowingModel:
+    """A stand-in for a transformers model that grows the cache in place.
+
+    Real transformers models append the newly generated key/value tensors to
+    the ``past_key_values`` object they are given (``DynamicLayer.update``
+    rebinds ``keys``/``values`` on the very same layer object). Handing the
+    stored activation cache to the model therefore grows the stored cache with
+    every turn.
+    """
+
+    def __init__(self, device: str = "cpu") -> None:
+        self.device = device
+        self.received_caches = []
+
+    def __call__(
+        self, input_ids=None, use_cache=False, past_key_values=None, return_dict=True, **kwargs
+    ):
+        self.received_caches.append(past_key_values)
+        if past_key_values is not None:
+            past_key_values.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+        output = MagicMock()
+        output.logits = torch.full((1, 1, 100), -1e9)
+        output.logits[0, 0, 7] = 1.0
+        output.past_key_values = past_key_values
+        return output
 
 
 @patch("transformers.AutoModelForCausalLM", MagicMock())
@@ -182,3 +215,52 @@ def test_kv_cache_generation_with_sampling(self):
         kv_cache = DynamicCache()
         resp = llm.generate([{"role": "user", "content": "Sampling"}], past_key_values=kv_cache)
         self.assertEqual(resp, self.standard_response)
+
+    def test_generate_does_not_mutate_caller_cache(self):
+        # Regression for #2301: the cache stored in the activation memory must
+        # not grow with the tokens generated for the current turn.
+        config = HFLLMConfig(
+            model_name_or_path="qwen3:0.6b",
+            max_tokens=5,
+            do_sample=False,
+            add_generation_prompt=True,
+        )
+        llm = self._create_llm(config)
+        model = CacheGrowingModel()
+        llm.model = model
+
+        cache = make_dynamic_cache(seq_len=3)
+        resp = llm.generate([{"role": "user", "content": "Hello"}], past_key_values=cache)
+
+        self.assertEqual(resp, self.standard_response)
+        # The model received a copy, not the caller's cache.
+        self.assertIsNot(model.received_caches[0], cache)
+        # The copy did grow, i.e. the model really appends in place.
+        self.assertGreater(dynamic_cache_seq_len(model.received_caches[0]), 3)
+        # The caller's cache still holds only the original tokens.
+        self.assertEqual(dynamic_cache_seq_len(cache), 3)
+
+    def test_generate_stream_does_not_mutate_caller_cache(self):
+        # Same regression for the streaming path.
+        config = HFLLMConfig(
+            model_name_or_path="qwen3:0.6b",
+            max_tokens=5,
+            do_sample=False,
+            add_generation_prompt=True,
+        )
+        llm = self._create_llm(config)
+        model = CacheGrowingModel()
+        llm.model = model
+
+        cache = make_dynamic_cache(seq_len=3)
+        chunks = list(
+            llm.generate_stream([{"role": "user", "content": "Hello"}], past_key_values=cache)
+        )
+
+        self.assertTrue(chunks)
+        # The model received a copy, not the caller's cache.
+        self.assertIsNot(model.received_caches[0], cache)
+        # The copy did grow, i.e. the model really appends in place.
+        self.assertGreater(dynamic_cache_seq_len(model.received_caches[0]), 3)
+        # The caller's cache still holds only the original tokens.
+        self.assertEqual(dynamic_cache_seq_len(cache), 3)
diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py
index 6490d687f..2a3b5f645 100644
--- a/tests/memories/activation/test_kv.py
+++ b/tests/memories/activation/test_kv.py
@@ -1,13 +1,23 @@
 from unittest.mock import MagicMock
 
 import pytest
-import torch
 
 from transformers import DynamicCache
 
 from memos.configs.memory import KVCacheMemoryConfig
 from memos.memories.activation.item import KVCacheItem
-from memos.memories.activation.kv import KVCacheMemory
+from memos.memories.activation.kv import KVCacheMemory, clone_dynamic_cache
+from tests.utils import (
+    dynamic_cache_key_tensors,
+    dynamic_cache_layer_count,
+    dynamic_cache_seq_len,
+    make_dynamic_cache,
+)
+
+
+# The KV cache is a torch structure; skip the whole module where torch is
+# not installed instead of failing the collection of the test suite.
+torch = pytest.importorskip("torch")
 
 
 @pytest.fixture
@@ -35,10 +45,7 @@ def kv_memory(dummy_config):
 
 def make_filled_cache():
     # Create a DynamicCache with at least one dummy tensor layer
-    cache = DynamicCache()
-    cache.key_cache.append(torch.zeros(1, 2, 3))
-    cache.value_cache.append(torch.zeros(1, 2, 3))
-    return cache
+    return make_dynamic_cache(seq_len=3)
 
 
 def test_extract_and_add_and_get(kv_memory):
@@ -58,9 +65,8 @@ def test_get_cache_merge(kv_memory):
     kv_memory.add([item1, item2])
     merged = kv_memory.get_cache([item1.id, item2.id])
     assert isinstance(merged, DynamicCache)
-    # Check the number of layers in merged key/value cache
-    assert len(merged.key_cache) == 1
-    assert len(merged.value_cache) == 1
+    # Check the number of layers in merged cache
+    assert dynamic_cache_layer_count(merged) == 1
 
 
 def test_delete_and_get_all(kv_memory):
@@ -84,3 +90,153 @@ class DummyTextualMemory:
     item = kv_memory.from_textual_memory(DummyTextualMemory())
     assert isinstance(item, KVCacheItem)
     assert item.metadata["bar"] == 1
+
+
+def test_clone_dynamic_cache_does_not_share_tensors():
+    # A clone must not share any tensor with the original: transformers appends
+    # the generated tokens to the cache object in place.
+    cache = make_filled_cache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert dynamic_cache_layer_count(cloned) == dynamic_cache_layer_count(cache)
+
+    original_keys = dynamic_cache_key_tensors(cache)[0]
+    cloned_keys = dynamic_cache_key_tensors(cloned)[0]
+    assert cloned_keys is not original_keys
+    assert cloned_keys.data_ptr() != original_keys.data_ptr()
+    assert torch.equal(cloned_keys, original_keys)
+
+    # The copy still reports the same number of cached tokens, i.e. the layer
+    # state (such as the initialization flag) survives the copy.
+    assert cloned.get_seq_length() == cache.get_seq_length()
+
+    # Appending to the clone (what a model does) leaves the original untouched.
+    cloned.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+    assert dynamic_cache_seq_len(cloned) == 4
+    assert dynamic_cache_seq_len(cache) == 3
+    assert cloned.get_seq_length() == 4
+    assert cache.get_seq_length() == 3
+
+
+def test_get_cache_single_item_returns_independent_copy(kv_memory):
+    # The stored cache must never be handed out: the model appends the
+    # generated tokens to the cache it receives, in place.
+    item = KVCacheItem(memory=make_filled_cache())
+    kv_memory.add([item])
+
+    cached = kv_memory.get_cache([item.id])
+
+    assert cached is not None
+    assert cached is not item.memory
+    assert dynamic_cache_seq_len(cached) == 3
+
+    cached.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+
+    assert dynamic_cache_seq_len(cached) == 4
+    assert dynamic_cache_seq_len(item.memory) == 3
+
+
+def test_clone_dynamic_cache_handles_empty_cache():
+    # An empty cache has to be copied into an independent empty cache: the
+    # model appends its first layer to the cache it receives, and that append
+    # must not reach the stored cache.
+    cache = DynamicCache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert dynamic_cache_layer_count(cloned) == 0
+    assert dynamic_cache_layer_count(cache) == 0
+
+    cloned.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+
+    assert dynamic_cache_layer_count(cloned) == 1
+    assert dynamic_cache_layer_count(cache) == 0
+    assert cache.get_seq_length() == 0
+
+
+def test_get_cache_merge_does_not_alias_inputs(kv_memory):
+    # Merging concatenates the inputs into new tensors; the stored caches
+    # stay untouched.
+    item1 = KVCacheItem(memory=make_filled_cache())
+    item2 = KVCacheItem(memory=make_filled_cache())
+    kv_memory.add([item1, item2])
+
+    merged = kv_memory.get_cache([item1.id, item2.id])
+
+    assert merged is not None
+    assert dynamic_cache_layer_count(merged) == 1
+    # Two caches of 3 tokens each are merged into a single 6-token layer.
+    assert dynamic_cache_seq_len(merged) == 6
+
+    merged.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+
+    assert dynamic_cache_seq_len(item1.memory) == 3
+    assert dynamic_cache_seq_len(item2.memory) == 3
+
+
+def test_clone_dynamic_cache_handles_legacy_layout():
+    # transformers < 4.55 keeps the tensors in the flat "key_cache" and
+    # "value_cache" lists instead of per-layer objects.
+    class LegacyCache:
+        def __init__(self):
+            self.key_cache = [torch.zeros(1, 2, 3, 4)]
+            self.value_cache = [torch.ones(1, 2, 3, 4)]
+
+    cache = LegacyCache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert cloned.key_cache[0] is not cache.key_cache[0]
+    assert cloned.key_cache[0].data_ptr() != cache.key_cache[0].data_ptr()
+    assert torch.equal(cloned.key_cache[0], cache.key_cache[0])
+    assert cloned.value_cache[0].data_ptr() != cache.value_cache[0].data_ptr()
+    assert torch.equal(cloned.value_cache[0], cache.value_cache[0])
+
+
+def test_clone_dynamic_cache_prefers_the_layers_layout():
+    # A cache that exposes the modern "layers" layout next to the legacy
+    # "key_cache"/"value_cache" aliases must be copied through the modern
+    # layout: the model appends its layers to "layers", so that list must not
+    # be shared with the original.
+    class CacheWithLegacyAliases:
+        def __init__(self):
+            self.layers = []
+            self.key_cache = []
+            self.value_cache = []
+
+        def update(self, key_states, value_states, layer_idx):
+            self.layers.append(layer_idx)
+
+    cache = CacheWithLegacyAliases()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert cloned.layers is not cache.layers
+
+    cloned.update(None, None, 0)
+
+    assert cache.layers == []
+
+
+def test_clone_dynamic_cache_copies_layers_with_legacy_attribute_names():
+    # Some transformers versions expose "key_cache"/"value_cache" on the layer
+    # objects instead of "keys"/"values"; both variants have to be copied.
+    class LegacyLayer:
+        def __init__(self):
+            self.key_cache = torch.zeros(1, 2, 3, 4)
+            self.value_cache = torch.ones(1, 2, 3, 4)
+
+    class LayeredCache:
+        def __init__(self):
+            self.layers = [LegacyLayer()]
+
+    cache = LayeredCache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert cloned.layers[0] is not cache.layers[0]
+    assert cloned.layers[0].key_cache.data_ptr() != cache.layers[0].key_cache.data_ptr()
+    assert torch.equal(cloned.layers[0].key_cache, cache.layers[0].key_cache)
+    assert cloned.layers[0].value_cache.data_ptr() != cache.layers[0].value_cache.data_ptr()
+    assert torch.equal(cloned.layers[0].value_cache, cache.layers[0].value_cache)
diff --git a/tests/utils.py b/tests/utils.py
index ec8a32799..0f1cf4842 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -9,6 +9,88 @@
 from pydantic.aliases import PydanticUndefined
 
 
+def make_dynamic_cache(seq_len: int = 3, num_layers: int = 1):
+    """
+    Build a ``DynamicCache`` with populated layers.
+
+    ``transformers`` stores the key/value tensors either in the legacy
+    ``key_cache``/``value_cache`` lists (< 4.55) or in per-layer objects
+    (``layers``, >= 4.55). Building the cache through the public ``update()``
+    API works for both layouts.
+
+    Args:
+        seq_len: Number of cached tokens per layer.
+        num_layers: Number of layers to populate.
+
+    Returns:
+        DynamicCache: A cache with ``num_layers`` populated layers.
+    """
+    import torch
+
+    from transformers import DynamicCache
+
+    cache = DynamicCache()
+    for layer_idx in range(num_layers):
+        cache.update(
+            torch.zeros(1, 2, seq_len, 4),
+            torch.ones(1, 2, seq_len, 4),
+            layer_idx,
+        )
+    return cache
+
+
+def dynamic_cache_key_tensors(cache: Any) -> list[Any]:
+    """
+    Return the key tensor of every layer of a cache, for either cache layout.
+
+    Args:
+        cache: The cache to inspect.
+
+    Returns:
+        list: One key tensor per layer.
+    """
+    layers = getattr(cache, "layers", None)
+    if layers is not None:
+        tensors = []
+        for layer in layers:
+            tensor = getattr(layer, "keys", None)
+            if tensor is None:
+                tensor = getattr(layer, "key_cache", None)
+            tensors.append(tensor)
+        return tensors
+    return list(cache.key_cache)
+
+
+def dynamic_cache_layer_count(cache: Any) -> int:
+    """
+    Return the number of layers of a cache, for either cache layout.
+
+    Args:
+        cache: The cache to inspect.
+
+    Returns:
+        int: Number of layers.
+    """
+    layers = getattr(cache, "layers", None)
+    if layers is not None:
+        return len(layers)
+    return len(cache.key_cache)
+
+
+def dynamic_cache_seq_len(cache: Any, layer_idx: int = 0) -> int:
+    """
+    Return the number of cached tokens held by one layer of a cache.
+
+    Args:
+        cache: The cache to inspect.
+        layer_idx: Index of the layer to look at.
+
+    Returns:
+        int: Number of cached tokens.
+    """
+    return dynamic_cache_key_tensors(cache)[layer_idx].shape[-2]
+
+
 def check_module_base_class(cls: Any) -> None:
     """
     General function to test the correctness of an abstract base class.