fix: stop activation KV cache from growing every turn

[yetuge]

Fixes #2301.

Root cause

The stored activation cache and the cache handed to generation were the same object, so every turn permanently grew the store (and the re-dumped memory file):

1. MemOS.chat (mem_os/core.py), mem_os/main.py, mem_chat/simple.py and mos_for_test_scheduler.py 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 (DynamicLayer.update rebinds keys/values on the same object; identity is preserved);
3. KVCacheMemory._concat_caches returns caches[0] unchanged for a single id, so the get_cache() merge path hands out the stored object as well.

No caller reads the cache back expecting it to have grown — the growth is only observable as leaked state (which ActivationMemoryManager then re-dumps to disk).

Fix

Treat 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): independent copy with cloned K/V tensors, compatible with both the legacy key_cache/value_cache structure (transformers <= 4.55, per poetry.lock) and the newer layers structure (>= 4.56, still < 5.0.0 in the supported range);
• KVCacheMemory._concat_caches: the single-cache case now returns a clone instead of the stored object;
• HFLLM.generate / generate_stream: clone the incoming past_key_values once at the boundary. This fixes all four call sites without touching them; the cost is one cache copy per generate call, the same order as the prefill work itself.

Tests

• test_generate_with_cache_does_not_mutate_caller_cache (tests/llms/test_hf.py): mocks a model forward that appends K/V in place, asserts the caller's cache is unchanged — fails on current main, passes with this fix;
• test_get_cache_single_item_returns_independent_copy + test_get_cache_multi_item_merge_does_not_alias_inputs (tests/memories/activation/test_kv.py): get_cache() never aliases the store;
• test_clone_dynamic_cache_copies_legacy_tensors + test_clone_dynamic_cache_handles_layers_structure: cover both cache structures;
• full tests/memories/ + tests/llms/: 117 passed (+4 subtests), no regressions (Python 3.13, torch 2.14 CPU, transformers 4.53.2 per poetry.lock).

cc @issue reporter — thanks for the exceptionally detailed write-up; the line-level analysis made this straightforward to confirm and fix.

[Copilot]

🟡 Changes recommended

Layered cache cloning may remain incorrect on transformers ≥4.56, and streaming generation lacks regression coverage.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Fixes activation KV-cache aliasing that causes stored caches to grow across turns.

Changes:

• Adds cache cloning for legacy and layered cache formats.
• Clones caches at retrieval and HF generation boundaries.
• Adds cache-isolation regression tests.

File summaries

File	Summary
tests/memories/activation/test_kv.py	Tests cloning and merge isolation.
tests/llms/test_hf.py	Tests that generation does not mutate caller caches.
src/memos/memories/activation/kv.py	Adds cache cloning. Critical (3 votes): layered cache cloning may reset required layer state on transformers ≥4.56; clone layer state and test a real update.
src/memos/llms/hf.py	Clones caches before generation. Nit (1 vote): add regression coverage for streaming generation.

Review details

Suppressed comments (2)

src/memos/llms/hf.py:116

• The new generate_stream boundary is not covered by the regression test, which only exercises generate. A future change could preserve the caller cache for non-streaming generation while reintroducing aliasing on the streaming path; add a streaming test that performs the same in-place K/V append and checks the original cache length.

            from memos.memories.activation.kv import clone_dynamic_cache

            yield from self._generate_with_cache_stream(
                prompt, clone_dynamic_cache(past_key_values)
            )

src/memos/memories/activation/kv.py:279

• The new-layer branch only copies keys/values, but this codebase already supports layer objects that expose the alternative key_cache/value_cache names in move_dynamic_cache_htod. For those caches, the cloned layer is appended without any tensors, so generation receives an empty cache or fails instead of using the stored activation memory. Copy both attribute variants.

            if getattr(layer, "keys", None) is not None:
                new_layer.keys = layer.keys.clone()
                new_layer.values = layer.values.clone()

• Files reviewed: 4/4 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Memtensor-AI]

🤖 Open Code Review

Target: PR #2359
Task: 18f1ad369c72d2ca
Base: main
Head: fix/activation-cache-aliasing

🔍 OpenCodeReview found 7 issue(s) in this PR.

---

1. src/memos/memories/activation/kv.py (L290-L291)

Using zip here silently truncates the iteration if key_cache and value_cache have different lengths, producing a clone that is missing the trailing layers with no error or warning. Use zip(..., strict=True) (Python 3.10+) or an explicit length-equality assert before the loop to catch this early.

Suggested fix:

assert len(cache.key_cache) == len(cache.value_cache), (
    f"key_cache/value_cache length mismatch: "
    f"{len(cache.key_cache)} vs {len(cache.value_cache)}"
)
for keys, values in zip(cache.key_cache, cache.value_cache):

---

2. src/memos/memories/activation/kv.py (L269-L271)

DynamicCache from transformers maintains a _seen_tokens counter (an integer tracking how many tokens have been processed so far). Neither the layers branch nor the key_cache branch copies it. The cloned cache will start with _seen_tokens = 0, which can corrupt position IDs and attention masks during generation when the model uses this value to offset positions. Copy it explicitly after constructing cloned:

cloned = DynamicCache()
if hasattr(cache, "_seen_tokens"):
    cloned._seen_tokens = cache._seen_tokens

---

3. tests/llms/test_hf.py (L225-L226)

The test only verifies the original cache was not mutated, but never asserts that the clone it handed to the model actually received the mutations. A broken clone_dynamic_cache that returns an empty DynamicCache() (with no key_cache entries at all) would cause kv.key_cache[0] in the forward mock to raise IndexError, which would be caught by the outer try/finally, and generate would raise — but the mutation-check assertions at the end would still never be reached (the test would error-out rather than fail on the right thing). Add an assertion that verifies the clone was mutated as expected, e.g.:

# Verify the cloned cache that was handed to the model did grow
# (confirms clone_dynamic_cache actually copies content, not an empty shell)
clone_received_mutations = False
for call in self.mock_model.call_args_list:
    kv = call[1].get("past_key_values")
    if kv is not None and kv.key_cache and kv.key_cache[0].shape[-2] > 2:
        clone_received_mutations = True
        break
self.assertTrue(clone_received_mutations, "clone handed to model must have been mutated by forward")
self.assertEqual(kv_cache.key_cache[0].shape, (1, 2, 3))
self.assertEqual(kv_cache.value_cache[0].shape, (1, 2, 3))

Without this, the test only guards one half of the invariant.

---

4. tests/memories/activation/test_kv.py (L99-L105)

The in-place mutation check (fill_(99.0)) after the index reassignment does not test what it claims. After merged.key_cache[0] = torch.cat(...), the slot holds a brand-new tensor that cannot share storage with item.memory.key_cache[0] by construction — the torch.cat always allocates fresh memory. The subsequent fill_ therefore always passes regardless of whether clone_dynamic_cache truly deep-copies the original tensor, making the assertion vacuous and unable to catch a shallow-clone bug.

To actually verify storage independence, run the fill_ on the original cloned tensor before replacing it:

merged = kv_memory.get_cache([item.id])
assert merged is not item.memory

# Verify the clone is a real copy, not a view sharing storage.
merged.key_cache[0].fill_(99.0)
assert not torch.all(item.memory.key_cache[0] == 99.0), "get_cache shares storage with store"
assert item.memory.key_cache[0].shape == (1, 2, 3)

# Separately verify that appending to the handed-out cache doesn't grow the store.
merged.key_cache.append(torch.ones(1, 1, 3))
assert len(item.memory.key_cache) == 1

---

5. tests/memories/activation/test_kv.py (L127-L133)

Same structural problem as in test_get_cache_single_item_returns_independent_copy: after cloned.key_cache[0] = torch.ones(1, 5, 3), the slot holds a freshly allocated tensor that trivially cannot share storage with cache.key_cache[0]. The fill_(99.0) check that follows is guaranteed to pass and does not exercise the case it is documented to catch — a clone whose tensors share underlying storage with the original.

Move the fill_ call to before the index reassignment so it operates on the actually-cloned tensor:

cloned = clone_dynamic_cache(cache)
assert cloned is not cache
assert cloned.key_cache[0] is not cache.key_cache[0]
assert torch.equal(cloned.key_cache[0], cache.key_cache[0])

# In-place mutation must not leak: operates on the cloned tensor directly.
cloned.key_cache[0].fill_(99.0)
assert not torch.all(cache.key_cache[0] == 99.0), "clone shares storage with original"

# Reference swap check (independent of storage).
cloned.key_cache[0] = torch.ones(1, 5, 3)
assert cache.key_cache[0].shape == (1, 2, 3)

---

6. tests/memories/activation/test_kv.py (L158-L164)

Same vacuous fill_ pattern as the other two tests: cloned.layers[0].keys is replaced with a fresh torch.ones tensor on the line above, so the fill_(99.0) call operates on that new tensor, not on the originally cloned tensor. This means the assertion cannot catch a clone whose tensors share storage with the original — it trivially passes regardless.

Run fill_ on the cloned tensor before replacing it:

cloned = clone_dynamic_cache(cache)
assert isinstance(cloned, DynamicCache)
assert len(cloned.layers) == 1
assert cloned.layers[0].keys is not cache.layers[0].keys
assert torch.equal(cloned.layers[0].keys, cache.layers[0].keys)

# In-place mutation on the actually-cloned tensor.
cloned.layers[0].keys.fill_(99.0)
assert not torch.all(cache.layers[0].keys == 99.0), "clone shares tensor storage with original"

# Reference swap check is separate.
cloned.layers[0].keys = torch.ones(2, 2, 3)
assert cache.layers[0].keys.shape == (1, 2, 3)

---

7. tests/memories/activation/test_kv.py (L88)

Only one blank line separates test_from_textual_memory from test_get_cache_single_item_returns_independent_copy. PEP 8 requires two blank lines between top-level function definitions.

💡 Suggested Change

Before:

def test_get_cache_single_item_returns_independent_copy(kv_memory):

After:

def test_get_cache_single_item_returns_independent_copy(kv_memory):

Generated by cloud-assistant via Open Code Review.

[Memtensor-AI]

⚠️ Automated Test Results: INCONCLUSIVE

Automated tests inconclusive (auto-generated test defect); treated as non-blocking. Manual review recommended. Details: Both new test files import torch unconditionally at the module level, causing collection-time failures in an environment where torch is not installed. The tests never execute.

Branch: fix/activation-cache-aliasing

[yetuge]

Thanks for the automated review — all 4 findings are addressed in 2808420:

1. test_hf.py: the mock now reads past_key_values via .get() with a positional fallback, and uses deterministic argmax logits (-1e9 everywhere except a non-EOS token) so the generation loop always runs all max_tokens turns instead of risking an early EOS sample (~1% chance).
2. kv.py: clone_dynamic_cache now guards keys and values independently in the legacy-layer path, so a layer with only one side populated no longer raises AttributeError; added test_clone_dynamic_cache_layers_guard_keys_and_values_independently to cover it.
3. & 4. test_kv.py: both clone tests now also mutate the cloned tensors in place (fill_(99.0)) and assert the original is untouched, so a future regression to a storage-sharing clone is caught rather than only slot rebinding.

All 16 tests in the two affected files pass locally.

[Memtensor-AI]

⚠️ Automated Test Results: ENV ISSUE

The test environment encountered an issue that requires manual attention.

Details: Both test files fail at collection time because PyTorch (torch) is not installed in the test execution environment. No test logic ran at all.
Branch: fix/activation-cache-aliasing

[yetuge]

Second review round addressed in 04e3cc4:

1. clone_dynamic_cache now also copies per-layer key_cache/value_cache attributes (some transformers versions carry that shape instead of keys/values, as move_dynamic_cache_htod handles), with a dedicated test asserting storage independence.
2. test_hf mock: removed the unreachable positional fallback — .get() stays with a comment explaining that _prefill always passes the cache by keyword.
3. test_get_cache_single_item_returns_independent_copy now also mutates in place (fill_) so a storage-sharing clone cannot slip through.

On the ENV ISSUE flag: the two new test files import torch at module level because they exercise real tensor semantics (shape growth, in-place mutation, storage sharing); with torch available they pass locally (17 passed on Python 3.13 / CPU torch). No test logic is skipped — the collection failure only occurs in environments where torch is absent. If the CI image can't install torch, an alternative is skipping these two files via a pytest collection hook there, but that would leave the regression unguarded, so I'd rather keep them and let the env provide torch.

[Memtensor-AI]

❌ Automated Test Results: FAILED

Auto-fix retry 1/2 triggered.

Failed tests:

• collection
• collection

Error details

Tests failed. Failed cases: collection, collection [advisory, non-gating] AI-generated tests on branch test/auto-gen-18f1ad369c72d2ca-20260911185234: 69/80 passed, 11 failed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/activation-cache-aliasing