fix: record and check model identity on serialized KV caches

[jeojdi1]

Description

A KV cache is the internal activation state of one specific set of weights — it is not portable data. KVCacheMemory.dump() currently writes only {"kv_cache_memories": ...}, and load() restores it unconditionally. There is nothing recording which model produced a cache and nothing checking it on the way back in.

The consequence is silent: a cache dumped under one model and loaded under another is accepted with no error, and the model simply produces different tokens. On a close fine-tune pair I measured a next-token KL shift of 0.08–0.92 with the top-1 token flipping on 2 of 5 probes. A distant architecture does raise, but only as an opaque RuntimeError about tensor sizes, which does not tell the user what actually went wrong.

This PR:

• records model_identity in the dumped payload (model_name_or_path, best-effort from the configured extractor LLM);
• checks it on load() and warns on mismatch, naming both models and saying what to do about it.

It is a warning rather than an exception on purpose: caches dumped before this field existed carry no identity, and refusing to load them would break every existing store. _model_identity() returns None when the LLM config exposes no name, so a dump never fails because identity could not be determined, and a None on either side skips the check.

Related Issue (Required): #2300

Note on overlap: issue #2203 (closed) and PR #2204 cover pickle.load in this same load() path as an unsafe-deserialization sink (CWE-502). This change is orthogonal — it adds a payload field and a check, and does not touch how the payload is deserialized — but it does edit adjacent lines, so it may need a trivial rebase depending on merge order.

Type of change

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

How Has This Been Tested?

• Unit Test

Three tests added:

• test_dump_records_model_identity — the field is written.
• test_load_warns_on_model_mismatch — dump under org/model-a, load under org/model-b, assert the warning names the original model. Fails on the current code, which emits nothing.
• test_load_is_quiet_when_model_matches — no warning on the matching case, so this is not a new source of noise.

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 (happy to, if you consider the new payload field user-facing)
• I have linked the issue to this PR (Serialized activation KV cache carries no model identity, so a cache built by one model loads silently into another #2300)
• I have mentioned the person who will review this PR

A note on the target branch

CONTRIBUTING.md says to open PRs against dev, but no dev branch exists — only main and dev-v2.0.28 … dev-v2.0.32. This is against main (185ebdb, "Dev v2.0.32"). Happy to retarget.

[Memtensor-AI]

🤖 Open Code Review

Target: PR #2302
Task: 022280e030428671
Base: main
Head: fix/kv-cache-carries-no-model-identity

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

---

1. src/memos/memories/activation/kv.py (L238-L241)

The saved parameter is typed dict | None but is sourced directly from data.get('model_identity') after a pickle.load, so its actual runtime type is unconstrained. If the stored value is anything other than dict or None (e.g., a plain string "gpt2" from a corrupted or hand-edited file), the guard saved is None is False, execution falls through, and saved.get(...) raises AttributeError, crashing the entire load path with an unhelpful traceback — the opposite of the intended graceful degradation.

Replace the guard with an isinstance check:

💡 Suggested Change

Before:

        current = self._model_identity()
        if saved is None or current is None:
            return
        if saved.get("model_name_or_path") != current.get("model_name_or_path"):

After:

        current = self._model_identity()
        if not isinstance(saved, dict) or current is None:
            return
        if saved.get("model_name_or_path") != current.get("model_name_or_path"):

---

2. src/memos/memories/activation/kv.py (L220-L226)

When none of the three hardcoded attribute names matches cfg, the method returns None. _check_model_identity then hits if saved is None or current is None: return and silently skips the entire identity check — including for mismatched models. This means the mismatch warning is permanently suppressed for any LLM config that exposes its model path under a different attribute name (e.g. model_id, checkpoint), with no log message or other indication that the check was skipped. Consider logging a debug/warning when cfg is non-None but no known attribute was found, so that gaps in coverage are discoverable.

---

3. tests/memories/activation/test_kv.py (L139)

The negative assertion filters on the literal substring "was built with model", while the positive counterpart (test_load_warns_on_model_mismatch) uses a completely different predicate — checking that both model name strings appear in the message. These two tests are supposed to be inverses of each other, but their predicates are not symmetric.

The production message is "KV cache was built with model %r but is being loaded into %r". If that phrasing ever changes (e.g., to "KV cache produced by model %r"), this filter would match zero records regardless of how many warnings were actually emitted, making the assertion pass vacuously. test_load_warns_on_model_mismatch would fail correctly on the same change — but this test would silently give false confidence.

Suggestion: mirror the predicate from the positive test so both break together on a format change:

assert not any("org/model-a" in r.getMessage() for r in caplog.records)

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 module from being imported at all. [advisory, non-gating] AI-generated tests on branch test/auto-gen-b6472f335cd88b41-20260829014814: 82/82 passed — these do NOT affect the PR verdict; review the branch manually.
Branch: fix/kv-cache-carries-no-model-identity

[jeojdi1]

Thanks for the review — taking one of these and pushing back on the other.

1. import pickle at module level — agreed, fixed in 56fe586.

2. hasattr(merged, "layers") is dead code — I think this one is backwards, and applying the suggestion would break the test.

DynamicCache does expose .layers, and key_cache is the attribute that no longer exists. In transformers v5.16.1, src/transformers/cache_utils.py:

1306:  self.layers = layers if layers is not None else []
1322:  return len(self.layers)          # __len__
1375:  while len(self.layers) <= layer_idx:
1383:  keys, values = self.layers[layer_idx].update(...)

and key_cache appears zero times in that file. The rename landed in 4.57; .key_cache / .value_cache were the pre-4.57 layout. So the suggested assert len(merged.key_cache) == 1 raises AttributeError on any current install — which is exactly the failure mode this PR series is fixing, and the reason test_get_cache_merge and test_delete_and_get_all currently fail on main (see #2313).

The hasattr branch is there deliberately because pyproject.toml:41 declares transformers >=4.51.3,<5.0.0, which spans both layouts, so the test has to work on either.

That said — your parenthetical suggestion of len(merged) == 1 is strictly better than the hasattr branch: __len__ returns len(self.layers) and is stable across both layouts. Happy to push that if a maintainer prefers it.

[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. [advisory, non-gating] AI-generated tests on branch test/auto-gen-022280e030428671-20260901164722: 0/54 passed, 54 failed — these do NOT affect the PR verdict; review the branch manually.
Branch: fix/kv-cache-carries-no-model-identity