fix: stop merged KV caches being silently discarded on transformers >= 4.57

[jeojdi1]

Description

Fixes #2313.

_concat_caches builds each merged layer with layer_cls() and then assigns .keys / .values directly. A layer constructed that way is never marked initialized, and on transformers >= 4.57 DynamicLayer.get_seq_length() short-circuits on that flag:

def get_seq_length(self) -> int:
    if not self.is_initialized or self.keys.numel() == 0:
        return 0
    return self.keys.shape[-2]

So a correctly concatenated cache reports a length of 0, the model treats it as empty, and every merged token is dropped on the first forward pass. No exception, no warning — activation memory simply has no effect, and the only symptom is that the model behaves as if the memory were never loaded.

The fix initializes the layer through its public lazy_initialization path before assigning, guarded by hasattr so the older key_cache layout is untouched.

This PR also repairs the test fixture, which is what hid the bug. make_filled_cache appended to cache.key_cache, removed in 4.57, so test_get_cache_merge and test_delete_and_get_all failed with AttributeError on any current install. PR #2204's description notes these as "pre-existing test_kv.py failures … transformers-API-version issues unrelated to this patch" — they are version issues, but rebuilding the fixture on the public update API both makes them pass and makes them exercise the layers path, which the old fixture never did. That is what surfaces the real bug.

Related Issue (Required): #2313

Type of change

• Bug fix (non-breaking change which fixes an issue)

How Has This Been Tested?

• Unit Test

Added test_concat_caches_preserves_seq_length, which builds two caches through the public update API, merges them, and asserts the merged length is the sum.

Verified against the requirement that a bug fix ships a regression test that fails on the old code and passes on the new, on transformers 5.16.1:

# upstream src/memos/memories/activation/kv.py
FAILED tests/memories/activation/test_kv.py::test_concat_caches_preserves_seq_length
1 failed

# with this change
1 passed

Whole file, before and after:

before:  2 failed, 3 passed   (test_get_cache_merge, test_delete_and_get_all)
after:   5 passed

Checklist

• I have performed a self-review of my own code
• I have commented my code in hard-to-understand areas
• I have added tests that prove my fix is effective
• I have created related documentation issue/PR in MemOS-Docs (not applicable — no user-facing API change)
• I have linked the issue to this PR
• I have mentioned the person who will review this PR

A note on the target branch

CONTRIBUTING.md:283 says to branch off dev and open PRs against dev, but there is no plain dev branch on this repository — the branches are main plus dev-v2.0.28 … dev-v2.0.32. This PR targets main, which is the default branch, where I verified the behaviour, and where 11 of the last 15 merged PRs landed. src/memos/memories/activation/kv.py is byte-identical on main and dev-v2.0.30, so the rebase is trivial either way — happy to retarget to whichever branch you prefer, just say which.

Relationship to the other KV cache PRs

This is one of three independent fixes to src/memos/memories/activation/kv.py, each kept to a single logical change per CONTRIBUTING.md. They touch adjacent lines and are easiest to review in this order:

• this PR — merged caches report length 0 and are discarded (Merged KV caches report get_seq_length() == 0 and are silently discarded on transformers >= 4.57 #2313)
• fix: record and check model identity on serialized KV caches #2302 — serialized caches carry no model identity (Serialized activation KV cache carries no model identity, so a cache built by one model loads silently into another #2300)
• fix: stop get_cache handing out the stored KV cache by reference #2303 — get_cache hands out the stored cache by reference, so it grows every turn (The stored activation DynamicCache is mutated in place by generation, so activation memory grows every turn #2301)
• fix: re-rotate KV fragments to their landing positions when merging #2304 (draft) — fragments are not re-rotated to their landing positions

Happy to rebase them into whatever order suits review.

[Memtensor-AI]

🤖 Open Code Review

Target: PR #2314
Task: 67a84472b3458e97
Base: main
Head: fix/merged-kv-cache-reports-zero-length

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

---

1. src/memos/memories/activation/kv.py (L250-L255)

When is_initialized is False but lazy_initialization is absent (e.g., a different cache implementation or a future transformers API change), both conditions fail and execution silently falls through to the bare direct assignments on lines 254–255. The comment in this very block explains that direct assignment while is_initialized is False causes silent data loss — so this fallthrough silently reproduces the bug being fixed. Consider logging a warning on the else path to make this observable.

💡 Suggested Change

Before:

                if not getattr(merged_layer, "is_initialized", True) and hasattr(
                    merged_layer, "lazy_initialization"
                ):
                    merged_layer.lazy_initialization(merged_keys, merged_values)
                merged_layer.keys = merged_keys
                merged_layer.values = merged_values

After:

                if not getattr(merged_layer, "is_initialized", True):
                    if hasattr(merged_layer, "lazy_initialization"):
                        merged_layer.lazy_initialization(merged_keys, merged_values)
                    else:
                        logger.warning(
                            "Merged cache layer is uninitialized and has no "
                            "lazy_initialization; get_seq_length() may return 0 "
                            "and merged tokens may be silently discarded."
                        )
                merged_layer.keys = merged_keys
                merged_layer.values = merged_values

---

2. src/memos/memories/activation/kv.py (L253-L255)

lazy_initialization(merged_keys, merged_values) receives the real tensors and almost certainly stores them internally (setting is_initialized = True and populating self.keys/self.values). The unconditional assignments that immediately follow bypass any property setters or internal bookkeeping the method applies, and are at best redundant. If the direct assignments are still required after the call, the intent should be made explicit with a comment; otherwise they should be placed in an else branch so they only run when lazy_initialization was not called.

💡 Suggested Change

Before:

                    merged_layer.lazy_initialization(merged_keys, merged_values)
                merged_layer.keys = merged_keys
                merged_layer.values = merged_values

After:

                    merged_layer.lazy_initialization(merged_keys, merged_values)
                else:
                    merged_layer.keys = merged_keys
                    merged_layer.values = merged_values

Generated by cloud-assistant via Open Code Review.

[Memtensor-AI]

⚠️ Automated Test Results: ENV ISSUE

The test environment encountered an issue that requires manual attention.

Details: Test collection failed because the torch module is not installed in the test environment, preventing the test file from being imported.
Branch: fix/merged-kv-cache-reports-zero-length