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: 8fc024ed3ed670ab
Base: main
Head: fix/activation-cache-aliasing

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

---

1. tests/cache_helpers.py (L54-L63)

The try/except TypeError block wraps both the constructor and the two cache.update() calls. A TypeError raised during update() — for example if the new layers-based API expects different tensor shapes or keyword arguments — would be silently caught and turned into a pytest.skip("DynamicCache(config=...) is not supported"), which is the wrong diagnosis. This masks real bugs in the production update path.

Scope the try to only the constructor call, and let update errors propagate normally.

💡 Suggested Change

Before:

    try:
        cache = DynamicCache(config=HybridConfig())
        if populate:
            keys = torch.zeros(1, 2, 3, 4)
            values = torch.zeros(1, 2, 3, 4)
            cache.update(keys, values, layer_idx=0)
            cache.update(keys, values, layer_idx=1)
    except TypeError:
        pytest.skip("DynamicCache(config=...) is not supported")
    return cache

After:

    try:
        cache = DynamicCache(config=HybridConfig())
    except TypeError:
        pytest.skip("DynamicCache(config=...) is not supported")
    if populate:
        keys = torch.zeros(1, 2, 3, 4)
        values = torch.zeros(1, 2, 3, 4)
        cache.update(keys, values, layer_idx=0)
        cache.update(keys, values, layer_idx=1)
    return cache

---

2. tests/cache_helpers.py (L8-L9)

The hasattr check is correctly done on cache (the instance being populated), so there is no two-object inconsistency — this is fine as written.

💡 Suggested Change

Before:

    cache = DynamicCache()
    keys = torch.zeros(1, 2, 3, 4) if hasattr(cache, "layers") else torch.zeros(1, 2, 3)

After:

    cache = DynamicCache()
    keys = torch.zeros(1, 2, 3, 4) if hasattr(cache, "layers") else torch.zeros(1, 2, 3)

---

3. tests/memories/activation/test_kv.py (L118-L121)

If the aliasing assertion fails, merged_keys.zero_() is never reached, leaving the tensor filled with 99.0 for subsequent tests. Wrap the fill/assert/zero sequence in a try/finally to guarantee cleanup regardless of test outcome.

💡 Suggested Change

Before:

    merged_keys = cache_keys(merged)
    merged_keys.fill_(99.0)
    assert not torch.all(cache_keys(item.memory) == 99.0), "get_cache shares storage with store"
    merged_keys.zero_()

After:

    merged_keys = cache_keys(merged)
    try:
        merged_keys.fill_(99.0)
        assert not torch.all(cache_keys(item.memory) == 99.0), "get_cache shares storage with store"
    finally:
        merged_keys.zero_()

---

4. tests/memories/activation/test_kv.py (L367-L372)

The test asserts exactly 2 clones when a layer exposes all four attributes (keys, values, key_cache, value_cache), relying on the implementation preferring key_cache/value_cache and skipping keys/values. However, the implementation in kv.py uses copy.copy(layer) first, which may call clone() on tensor attributes depending on the tensor subclass's __copy__ behavior. If CloneCountingTensor.__copy__ delegates to clone(), the count could be 4 (from copy.copy) + 2 (from explicit clones) = 6, or some other number, making this assertion fragile. Verify that copy.copy on a CloneCountingTensor-bearing layer does not trigger clone(), or reset the counter after the copy.copy step inside the implementation.

---

5. tests/memories/activation/test_kv.py (L245-L248)

After cloned.layers[1].update(torch.ones(1, 2, 1, 4), ...), the test asserts keys.shape[-2] == 3 while cumulative_length == 4. This implies the sliding-window layer retains only sliding_window - 1 = 3 tokens after eviction — but with sliding_window=4 and 3 existing tokens, adding 1 should reach exactly 4 (the window capacity), so the expected shape should be 4, not 3. If the eviction policy drops the oldest token as soon as the window is full (i.e., keeps at most sliding_window - 1), the assertion is correct but that should be documented. If the policy keeps up to sliding_window tokens, the assertion is wrong and would be a false-passing test on versions where eviction hasn't triggered yet. Confirm the sliding-window eviction semantics against the transformers layer implementation.

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

[yetuge]

Follow-up on the latest CI run: all build-matrix jobs stopped at the Ruff checks before tests. The failure was a formatting-only result naming exactly src/memos/llms/hf.py and tests/memories/activation/test_kv.py; ruff check itself passed. I applied Ruff 0.11.8 formatting to those two files in commit d5649e6 (no behavior changes).

Local verification: ruff check passed, ruff format --check passed, the applicable pre-commit hooks passed, and python -m pytest tests/memories/activation/test_kv.py tests/llms/test_hf.py -q passed (17 tests). The separate autotest failure still reports the known missing-torch environment issue and was not changed.

[Memtensor-AI]

⚠️ Automated Test Results: ENV ISSUE

The test environment encountered an issue that requires manual attention.

Details: Both tests fail at collection time with ModuleNotFoundError for 'torch', meaning PyTorch is not installed in the test environment. No test logic or application code was actually executed.
Branch: fix/activation-cache-aliasing

[yetuge]

Thanks for the updated Open Code Review in comment 5629951647 (edited 2026-09-14). I rechecked all four findings against the current code and agree that they are actionable:

1. Layer attribute precedence: clone_dynamic_cache now selects the per-layer key_cache/value_cache naming scheme as a unit before falling back to keys/values, matching the precedence in move_dynamic_cache_htod. The inner guards remain independent so the existing keys-only and values-only compatibility case is preserved. I added test_clone_dynamic_cache_prefers_per_layer_cache_attributes for a layer exposing both schemes. I did not copy the suggested snippet literally because its final elif would drop a values-only layer, which the existing regression test intentionally covers.
2. Single-item alias test: fill_ now runs before replacing the list slot with the simulated appended tensor; the clone is reset and the append/shape assertion then runs. This makes the storage-alias check effective.
3. Legacy clone alias test: the same ordering fix is applied before assigning a new tensor to the list slot.
4. Layered value storage: the per-layer test now mutates value_cache in place and verifies the original value tensor is unchanged, symmetrically with key_cache.

The new hybrid-layer regression failed on the previous implementation (10 passed, 1 failed), then passed after the fix. Final local verification on commit 873782a3:

• python -m pytest tests/memories/activation/test_kv.py tests/llms/test_hf.py -q — 18 passed, 1 pre-existing Pydantic warning
• Ruff check, Ruff format check, and the applicable pre-commit hooks — passed

The fix is pushed as a new commit without rewriting history. The separate Python workflow remains action_required because this is a fork workflow awaiting maintainer approval; that is external CI state, not a code failure.

[yetuge]

Correction to the verification paragraph above: the additional latest-edited review findings were implemented in commit aae1cbbe305e20339b2619292f72777ca104d2b7 (following 873782a3). The final local verification on that commit is:

• python -m pytest tests/memories/activation/test_kv.py tests/llms/test_hf.py -q — 20 passed, 1 pre-existing Pydantic warning
• Ruff check, Ruff format check, and the applicable pre-commit hooks — passed

The fork Python workflow for this head is still action_required while awaiting maintainer approval; no code failure is reported by that workflow state.

[yetuge]

Follow-up on the current edited version of Open Code Review comment 5629951647 (updated 2026-09-14T10:47:25Z, now reporting three findings):

1. DynamicLayer state: agreed. clone_dynamic_cache now copies every non-tensor attribute from the source layer, including is_initialized and _seen_tokens, before cloning K/V tensors. test_clone_dynamic_cache_preserves_layer_state models the first update() decision and verifies that existing history is appended rather than replaced.
2. Unknown cache shape: agreed. clone_dynamic_cache now raises AttributeError when neither layers nor key_cache is available, matching _concat_caches instead of silently returning an empty cache. test_clone_dynamic_cache_rejects_unknown_shape covers this contract.
3. Layered key alias probe: agreed. The fill_ check now runs before the test replaces cloned.layers[0].keys, then resets the clone before the slot-replacement assertion.

The earlier edited version's four findings (layer naming precedence, the two legacy alias probes, and the missing value_cache probe) remain covered in the same head. The final local verification on aae1cbbe305e20339b2619292f72777ca104d2b7 is 20 passed with one pre-existing Pydantic warning; Ruff, format, and applicable pre-commit hooks pass. The fork Python workflow remains action_required pending maintainer approval.

[yetuge]

Thanks for the second in-place OCR update. I re-read comment 5629951647 at updated_at=2026-09-14T10:58:00Z and verified both findings against the current code.

1. Legacy cache-level state — confirmed. The legacy branch now copies cache-level non-tensor attributes (including _seen_tokens) with copy.deepcopy, while still rebuilding key_cache/value_cache from cloned tensors. Added test_clone_dynamic_cache_preserves_legacy_cache_state, which seeds _seen_tokens=2, clones, updates the clone to 3, and verifies the stored cache remains at 2.

2. Mutable layer metadata aliasing — confirmed. Layer non-tensor attributes are now copied with copy.deepcopy, so nested mutable state is independent as well as the K/V tensors. Added test_clone_dynamic_cache_copies_mutable_layer_state, which mutates a nested list through the clone and verifies the original layer is unchanged.

The earlier findings remain covered, including cache naming precedence, pre-replacement alias probes for single-item/legacy/layered tests, and symmetric value_cache coverage. The new commit is dac93a1a39a9301008504f770497a99155f795bd (only src/memos/memories/activation/kv.py and tests/memories/activation/test_kv.py), pushed without rewriting history.

Verification on the new head: python -m pytest tests/memories/activation/test_kv.py tests/llms/test_hf.py -q → 22 passed, 1 warning (the existing Pydantic serializer warning); Ruff check/format and applicable pre-commit hooks pass. The fork Python workflow remains externally gated by maintainer approval (action_required), not a code failure.

[yetuge]

Thanks for the latest in-place OCR update. I re-read comment 5629951647 at updated_at=2026-09-14T11:10:36Z and verified both findings against the new head.

1. Mismatched legacy K/V layer counts — confirmed. The legacy clone now uses zip(..., strict=True), so a corrupted or partially populated cache raises immediately instead of silently truncating. Added test_clone_dynamic_cache_rejects_mismatched_legacy_layers, which verifies the ValueError on unequal key/value list lengths.

2. Unknown-shape exception contract — confirmed. An object with neither supported cache representation is an unsupported cache input, so TypeError is more accurate than AttributeError. Both clone_dynamic_cache and _concat_caches now raise TypeError with the same message, and test_clone_dynamic_cache_rejects_unknown_shape asserts that contract.

The earlier findings remain covered: cache naming precedence, pre-replacement alias probes, symmetric value-cache coverage, cache/layer state preservation, deep copying of mutable layer metadata, and explicit rejection of unknown shapes. The new commit is 507de420c8206cbf42750516f17925ccdb366f6e (only src/memos/memories/activation/kv.py and tests/memories/activation/test_kv.py), pushed as a new commit without rewriting history.

Verification on the new head: python -m pytest tests/memories/activation/test_kv.py tests/llms/test_hf.py -q → 23 passed, 1 warning (the existing Pydantic serializer warning); Ruff check/format and applicable pre-commit hooks pass. The fork Python workflow remains action_required pending maintainer approval, not a code failure.

[Memtensor-AI]

⚠️ Automated Test Results: INCONCLUSIVE

Automated tests inconclusive (auto-generated test defect); treated as non-blocking. Manual review recommended. Details: Newly added tests in test_kv.py use a make_filled_cache() helper that directly appends to DynamicCache.key_cache, but the installed transformers version exposes only the layers attribute, so the helper fails before exercising the code under test. The same helper is imported by test_hf.py::test_generate_with_cache_does_not_mutate_caller_cache, causing that test to fail for the same reason.

Branch: fix/activation-cache-aliasing

[yetuge]

Thanks for the latest in-place edit of Open Code Review comment 5629951647. I re-read the current body at updated_at=2026-09-14T11:19:43Z and verified each of its three findings against the new head.

1. test_hf.py cache guard — fixed. The mocked forward now asserts that past_key_values is present before accessing it, so a call-shape regression cannot be swallowed by the surrounding try/finally. The test helper also populates both the legacy key_cache API and the modern DynamicCache.update()/layers API, and the test snapshots both caller K/V shapes before generation.

2. CPython-specific zip(strict=True) message — fixed. test_clone_dynamic_cache_rejects_mismatched_legacy_layers now asserts the controlled ValueError type without matching CPython's implementation-specific message. The strict length check remains in the legacy clone path.

3. Empty layer data-loss — not a defect in this implementation. I reproduced the supported transformers==4.56.2 API with a real hybrid cache. DynamicCache(config=...) intentionally creates lazy DynamicLayer and DynamicSlidingWindowLayer instances whose keys and values are both None before their first update; this is valid cache state, not an unsupported empty layer. The clone now preserves those real layer types, sliding_window, cumulative_length, and the uninitialized None state, covered by test_clone_dynamic_cache_preserves_uninitialized_real_hybrid_layers. An object with neither supported cache representation still raises TypeError.

The underlying constructor failure is fixed in commit 93ec8b1e: clone_dynamic_cache copies the layer instance without invoking a constructor that may require sliding_window, then deep-copies non-tensor state and clones tensor attributes while preserving the existing naming precedence. The real populated hybrid regression also verifies that updating the clone does not mutate the original.

The separate AutoTest comment 5663159190 is classified as INCONCLUSIVE / non-blocking because its generated helper used the removed legacy direct-list API. Both test_kv.py and test_hf.py now use the public update API on modern transformers while retaining the legacy branch.

Verification:

• Red before the source fix on transformers 4.56.2: 22 passed, 1 skipped, 2 failed with DynamicSlidingWindowLayer.__init__() missing sliding_window.
• Project environment transformers 4.53.2: python -m pytest tests/memories/activation/test_kv.py tests/llms/test_hf.py -q — 23 passed, 2 skipped.
• Modern transformers 4.56.2: the same command — 24 passed, 1 skipped.
• Ruff check/format and all applicable pre-commit hooks — passed.

The new commit was pushed without rewriting history. The fork's Python workflow for the new head is run 34860047308, currently action_required pending maintainer approval; AutoTest is pending on Open Code Review, so those are external gating states rather than local test failures. Earlier findings remain covered by the preceding commits and replies.

[yetuge]

Thanks for the latest Open Code Review pass. I re-read the current in-place version of comment 5629951647 at updated_at=2026-09-14T15:11:58Z (fingerprint sha256:aa2a2e4a38649680a09e32f334757e4e316b87bbe29aa6624e3d9aee55c11ba6) and checked all seven findings against head c1bd4a55.

1. The HF isolation test should assert clone identity — fixed. The mock now captures the cache passed to forward and asserts it is present and is not the caller's kv_cache, in addition to the caller K/V shape checks. This directly proves HFLLM.generate() hands generation an independent cache.

2. The HF legacy helper should use update() — fixed. The shared make_filled_cache() helper now populates both legacy and modern transformers caches through the public DynamicCache.update() API, so legacy bookkeeping is initialized consistently instead of being bypassed by direct list append.

3. The version-branching helpers should be shared — fixed. The helpers are now centralized in tests/cache_helpers.py; both test_kv.py and test_hf.py import the same implementation, including the real hybrid-cache factory and value-layer count accessor.

4. Legacy top-level tensor state should be copied — fixed. The legacy clone branch now clones every cache-level tensor attribute other than the K/V lists, while still deep-copying non-tensor state. test_clone_dynamic_cache_copies_legacy_tensor_state verifies storage independence on the legacy DynamicCache API.

5. copy.copy(layer) is supposedly wasted — not changed deliberately. It is the constructor-safe compatibility boundary: it preserves any instance state held in slots or supplied by a custom __copy__, while avoiding DynamicSlidingWindowLayer.__init__() and its required sliding_window. The following vars() loop overwrites every dictionary attribute with an independent clone/deep copy. Replacing this with object.__new__ would be less compatible with layer types that do not keep all state in __dict__; this is an allocation tradeoff, not a correctness failure.

6. The value-layer count assertion was weakened — fixed. test_get_cache_merge now checks both cache_layer_count(merged) == 1 and the explicit symmetric cache_value_layer_count(merged) == 1, as well as the value tensor itself.

7. The hybrid test should skip intermediate versions without config= — fixed. The real hybrid-cache helper now catches TypeError only around DynamicCache(config=...) and skips with a clear reason. Supported 4.56.2 continues to execute the real hybrid regression rather than skipping.

The previous constructor regression remains covered: the pre-fix 4.56.2 run failed at DynamicSlidingWindowLayer.__init__() because the old code called type(layer)(); the fixed clone preserves the real layer instance and metadata. The current verification is:

• Project transformers 4.53.2: python -m pytest tests/memories/activation/test_kv.py tests/llms/test_hf.py -q — 24 passed, 2 skipped.
• Modern transformers 4.56.2: the same command — 24 passed, 2 skipped.
• 26 tests collected; Ruff check/format and all applicable pre-commit hooks — passed.

The separate AutoTest comment 5663159190 remains INCONCLUSIVE / non-blocking: its generated test directly appended to the removed legacy key_cache API. The shared helper now uses DynamicCache.update() with the legacy-compatible path and the modern layers path.

Commit c1bd4a55 was pushed as a new commit without rewriting history. Earlier OCR findings remain covered by the preceding commits and replies. The fork Python workflow for this head is externally gated as action_required pending maintainer approval; AutoTest/OCR status is external gating, not a local test failure.

[Memtensor-AI]

✅ Automated Test Results: PASSED

All tests passed (22/22 executed, 4 skipped). memos_python_core/changed-repo-python: 22 passed, 4 skipped. Duration: 10s [advisory, non-gating] AI-generated tests on branch test/auto-gen-0cb44b46b1ce0f46-20260914233703: 49/49 passed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/activation-cache-aliasing

[yetuge]

Thanks for the latest in-place update of Open Code Review comment 5629951647. I re-read its current six findings at updated_at=2026-09-14T15:22:52Z (fingerprint sha256:0a014e984eaa22e8894fa4aba2603bae25436d5937d56b836bc940b95429f6a0) and checked them against head 39ecbea9dd60a59f063d9fd5b7e186e9d82817c8.

1. Duplicate value-layer helper/assertion — fixed. Removed the redundant cache_value_layer_count helper and duplicate layer-count assertion from tests/cache_helpers.py and tests/memories/activation/test_kv.py. The test still asserts that the merged value tensor is non-empty.

2. Module-level torch import — not changed deliberately. PyTorch is part of the optional all extra, and this module already follows the lazy-import pattern for torch in its tensor/cache operations. Keeping the function-local import preserves importability for installations that do not install the optional extra.

3. Duplicate K/V cloning — fixed. In src/memos/memories/activation/kv.py, the generic layer-state loop now skips keys, values, key_cache, and value_cache; the existing naming-precedence branch is the only code that clones those tensors. copy.copy(layer) remains for constructor-safe compatibility with modern layers and custom/slots state. test_clone_dynamic_cache_clones_layer_kv_tensors_once uses clone-counting tensors and verifies exactly the selected K/V pair is cloned.

4. HF layers branch — not changed. This branch is live on the supported modern environment: with transformers==4.56.2, the shared helper's DynamicCache() has a layers list and the branch executes. The real hybrid/lazy-layer behavior is also covered directly by the activation-cache tests, so splitting the HF test would add scope without fixing a dead path.

5. Failure-path cleanup — not changed. The mutation test uses a function-scoped fixture and its cache is not shared with later tests; a failed assertion terminates that test immediately. There is no production-state or cross-test pollution to correct here.

6. _seen_tokens guard — fixed. test_clone_dynamic_cache_preserves_legacy_cache_state now checks for _seen_tokens itself before setting or asserting it, so versions exposing another cache shape skip cleanly.

Verification:

• RED on unmodified c1bd4a5: the new clone-count regression failed with 6 clones versus the expected 2.
• Project environment (transformers==4.53.2): python -m pytest tests/memories/activation/test_kv.py tests/llms/test_hf.py -q — 25 passed, 2 skipped.
• Isolated transformers==4.56.2: the same command — 25 passed, 2 skipped; skips are only the absent _seen_tokens guard and legacy-only test, while modern layers/hybrid tests execute.
• Ruff check, Ruff format check, and applicable pre-commit hooks — passed.

Commit 39ecbea9dd60a59f063d9fd5b7e186e9d82817c8 was pushed to yetuge:fix/activation-cache-aliasing as one new Conventional Commit, without rewriting history. The separate AutoTest result comment 5666698150 is result-only and required no code action. Earlier findings remain covered by the preceding commits and replies.

[Memtensor-AI]

✅ Automated Test Results: PASSED

All tests passed (23/23 executed, 4 skipped). memos_python_core/changed-repo-python: 23 passed, 4 skipped. Duration: 10s

Branch: fix/activation-cache-aliasing

[yetuge]

Thanks for the follow-up review. I re-checked the four findings against the current head and the current test matrix.

1. Fixed in commit ddf7ad133b724a124708ad566022265cad600501. The try/except TypeError in tests/cache_helpers.py now covers both DynamicCache(config=...) construction and the populated cache.update(...) calls. Added test_make_real_hybrid_cache_skips_update_typeerror to lock in the skip behavior.

2. I did not change this. The reported legacy shape mismatch is not reproducible: the legacy tensors are (1, 2, 3), the appended tensors are (1, 1, 3), and concatenation is along dim=-2 (the middle dimension), so the non-concatenated dimensions match. The full legacy-compatible project test run passes.

3. I did not restore a duplicate value-layer-count helper/assertion. cache_layer_count(merged) == 1 is the single layer-count invariant, while cache_values(merged) is not None preserves the value-cache presence check. Reintroducing a second count helper would duplicate the same contract without adding an independent invariant.

4. I did not add cleanup-only zero_() calls. The test creates a local fake cache/layer and does not share it with later tests or fixtures; the mutation is immediately followed by the independence assertions, and no state escapes the test.

Verification for this follow-up change:

• Project transformers 4.53.2: 26 passed, 2 skipped.
• Isolated transformers 4.56.2: 26 passed, 2 skipped; modern layers and hybrid-cache paths executed.
• Ruff check and format check passed; applicable pre-commit hooks passed.
• The new regression test was red before the fix (TypeError escaped from cache.update) and green afterward in both environments (1 passed, 20 deselected).

The fix is pushed as the new commit above to fork/fix/activation-cache-aliasing. The latest AutoTest result comment (5673860422) is also passing.

[Memtensor-AI]

✅ Automated Test Results: PASSED

All tests passed (24/24 executed, 4 skipped). memos_python_core/changed-repo-python: 24 passed, 4 skipped. Duration: 10s

Branch: fix/activation-cache-aliasing

[yetuge]

Thanks for the additional review. I re-checked the seven findings against the current head and kept this round within the review-scope limit by closing only the safe, locally testable correctness issue.

1. Fixed in commit 499705a0c3f8bea866f3776d9ac788b6e7f6a447. clone_dynamic_cache now unconditionally resets cloned.layers before appending source layers. I added test_clone_dynamic_cache_replaces_preexisting_destination_layers; it was red on the previous head (2 destination layers instead of 1) and is green after the fix.

2. Not changed. In the supported versions tested here, transformers 4.53.2 uses the legacy cache API, while 4.56.2's real DynamicLayer and DynamicSlidingWindowLayer instances expose keys/values and do not expose key_cache/value_cache. The current non-None predicate also intentionally preserves populated keys/values for asymmetric or partially initialized test doubles; switching to hasattr alone could select a null per-layer cache and discard populated keys. The existing real hybrid and uninitialized-layer tests pass, so this proposed compatibility change is not proven safe within this review round.

3. Not changed. The #2301 reference is intentional: this PR's regression test covers the linked issue The stored activation DynamicCache is mutated in place by generation, so activation memory grows every turn #2301, while the change is delivered in PR fix: stop activation KV cache from growing every turn #2359. No traceability defect is present.

4. Not changed. torch.zeros_like(keys) would be a small fixture simplification, but the current literal shapes are equal and the helper is covered in both supported environments. This is not a correctness failure and is deferred under the review-scope limit.

5. Not changed. The regression test passed in both environments and the current pytest.skip.Exception assertion verifies that the helper emits the intended skip rather than silently accepting a different exception. Replacing it with a broader BaseException check would weaken the test; extracting a new abstraction is outside this minimal fix.

6. Not changed. Collection succeeds in both tested environments, and the legacy-only marker is a test compatibility gate with no production side effect. Moving it into the body is a style/defensive concern rather than an observed failure in this PR.

7. Not changed. The zeroing calls operate on local test tensors after storage-independence assertions; no fixture or state is shared with later tests. Removing them is cleanup-only and does not alter the tested contract.

Verification for this follow-up:

• Project transformers 4.53.2: 27 passed, 2 skipped.
• Isolated transformers 4.56.2: 27 passed, 2 skipped; real modern layers and hybrid paths executed.
• The new destination-layer regression was red before the fix (1 failed, 21 deselected) and green after it in the project environment; the complete target suite is green in both environments.
• Ruff check and format check passed; applicable pre-commit hooks passed.

The fix was pushed as the new commit above to fork/fix/activation-cache-aliasing. The latest AutoTest result comment (5674052227) is result-only and does not add a code action.

[Memtensor-AI]

✅ Automated Test Results: PASSED

All tests passed (25/25 executed, 4 skipped). memos_python_core/changed-repo-python: 25 passed, 4 skipped. Duration: 10s [advisory, non-gating] AI-generated tests on branch test/auto-gen-8fc024ed3ed670ab-20260915111818: 14/15 passed, 1 failed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/activation-cache-aliasing