Serialized activation KV cache carries no model identity, so a cache built by one model loads silently into another

[jeojdi1]

Re-checked 2026-08-28 against main at 185ebdb ("Dev v2.0.32", #2295). src/memos/memories/activation/kv.py is unchanged since #797 (2025-12-30) and _concat_caches is byte-identical on main and dev-v2.0.30, so everything below still applies at HEAD. The line numbers cite 9119efe, where it was measured.

Environment

• MemOS at main / 9119efe5554e61a94b669df5eb84cc1b8ef3c0ab
• torch 2.13.0, transformers 5.15.1, CPU, fp32
• Models: Qwen/Qwen2.5-0.5B-Instruct (writer) and Qwen/Qwen2.5-0.5B (reader) — same architecture
  and identical tensor shapes, different weights

Summary

KVCacheMemory.dump writes {"kv_cache_memories": self.kv_cache_memories} and nothing else
(src/memos/memories/activation/kv.py:195-198). The only metadata attached to an item is
{"source_text", "extracted_at"} (kv.py:47-50), and KVCacheItem
(src/memos/memories/activation/item.py:32-43) has no field for the producing model. KVCacheMemory.load
(kv.py:139-181) unpickles and installs whatever it finds without checking anything about the
producer. The only identity check in the cube load path is a config schema version string
(src/memos/mem_cube/general.py:60-66).

A KV cache is only valid for the exact weights, tokenizer and dtype that produced it, and caches are
explicitly designed to move between machines — GeneralMemCube.init_from_remote_repo clones a cube
from https://huggingface.co/datasets/<cube_id> and loads it (mem_cube/general.py:161-184).
Two further details make a mismatch easy to hit without anyone doing anything wrong:

• merge_config_with_default rebuilds the config from the caller's default and preserves only
  {"user_id", "cube_id", "config_filename", "model_schema"} (mem_cube/utils.py:148-153), so the
  model recorded in a cube's own config.json is discarded on that path.
• model_name_or_path is not a weights identity: HFLLM.__init__ loads with torch_dtype="auto"
  and no revision pin (llms/hf.py:44-54), and HFLLMConfig (configs/llm.py:130-138) has no
  dtype, quantization or revision field. The same string can resolve to different weights across
  machines, after a Hub update, or in a different precision.

Minimal reproduction

import pickle, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

WRITER, READER = "Qwen/Qwen2.5-0.5B-Instruct", "Qwen/Qwen2.5-0.5B"  # same shapes, different weights
tok  = AutoTokenizer.from_pretrained(WRITER)
TEXT = "The user's name is Alice and she lives in Berlin. Her favourite colour is green."
PROBES = ["\nQ: Which city?", "\nQ: What colour?", "\nSummarise the user:",
          "\nQ: Who is the user?", "\nWrite a greeting for the user:"]

def load(n): return AutoModelForCausalLM.from_pretrained(n, dtype=torch.float32).eval()
def build(m, t):
    ids = tok(t, return_tensors="pt").input_ids
    with torch.no_grad():
        return m(input_ids=ids, use_cache=True, return_dict=True).past_key_values
def last_logits(m, cache, p):
    import copy
    ids = tok(p, return_tensors="pt", add_special_tokens=False).input_ids
    with torch.no_grad():
        return m(input_ids=ids, use_cache=True, past_key_values=copy.deepcopy(cache),
                 return_dict=True).logits[0, -1].float()

# exactly what dump() writes -- kv.py:195-198
w = load(WRITER)
payload = {"kv_cache_memories": {"abc": {"id": "abc", "memory": build(w, TEXT),
           "metadata": {"source_text": TEXT, "extracted_at": "2026-08-26T00:00:00"}}}}
blob = pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)
del w

item = pickle.loads(blob)["kv_cache_memories"]["abc"]
print("metadata keys:", sorted(item["metadata"]))   # no model / tokenizer / dtype / revision

r = load(READER)                                    # different weights -- loads with no error
own = build(r, TEXT)
for p in PROBES:
    a, b = last_logits(r, item["memory"], p), last_logits(r, own, p)
    kl = torch.nn.functional.kl_div(a.log_softmax(-1), b.log_softmax(-1),
                                    log_target=True, reduction="sum").item()
    print(f"{p!r:36s} KL={kl:7.4f}  top1 {tok.decode([b.argmax().item()])!r}"
          f" vs {tok.decode([a.argmax().item()])!r}")

Observed

metadata keys: ['extracted_at', 'source_text']
'\nQ: Which city?'                   KL= 0.2535  top1 ' ' vs ' '
'\nQ: What colour?'                  KL= 0.3825  top1 ' ' vs ' Alice'
'\nSummarise the user:'              KL= 0.0821  top1 ' Alice' vs ' Alice'
'\nQ: Who is the user?'              KL= 0.9231  top1 '  ' vs ' ('
'\nWrite a greeting for the user:'   KL= 0.1477  top1 ' "' vs ' "'

No exception, no warning, no log line. The next-token distribution shifts and the greedy token flips
on 2 of 5 probes.

Being precise about severity: these two checkpoints are a base model and its own instruction-tuned
fine-tune, which is about the mildest possible mismatch, and even here the outputs diverge. I have
not measured an unrelated same-shape pair, where I would expect the divergence to be larger. In the
other direction, loading the same cache into a different architecture
(HuggingFaceTB/SmolLM2-360M-Instruct) fails, but with an opaque message that does not point at the
cause:

RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 2 but got size 5
              for tensor number 1 in the list.

Expected

Loading a cache produced by different weights, tokenizer, dtype or quantization should fail with a
clear message naming the mismatch (or be refused and rebuilt from source_text), rather than being
used silently or crashing inside torch.cat.

Why it matters

Activation memory is the one artifact in MemOS that is only meaningful for one exact set of weights,
and it is the one artifact with no producer identity recorded — so an upgraded model, a re-quantised
deployment, or a shared cube produces quietly different answers instead of an error.

Suggested fix

Record a fingerprint at extract()/dump() time and verify it at load():

# when building the item
item.metadata["producer"] = {
    "model_name_or_path": self.config.extractor_llm.config.model_name_or_path,
    "config_fingerprint": hashlib.sha256(
        json.dumps(self.llm.model.config.to_diff_dict(), sort_keys=True).encode()).hexdigest(),
    "torch_dtype": str(self.llm.model.dtype),
    "quantization": str(getattr(self.llm.model.config, "quantization_config", None)),
    "tokenizer_fingerprint": hashlib.sha256(
        self.llm.tokenizer.backend_tokenizer.to_str().encode()).hexdigest(),
    "num_layers": len(...), "num_kv_heads": ..., "head_dim": ...,
    "transformers_version": transformers.__version__,
}

On load(), compare against the live model and, on mismatch, log an error and drop that item (or
raise, if strictness is preferred) instead of installing it. metadata is already a free-form dict,
so this is backward compatible — items without a producer key can be accepted with a warning during
a deprecation window. A weights-content hash is stronger still, but the config + tokenizer + dtype +
shape fingerprint above is cheap and catches every case I reproduced, including the
torch_dtype="auto" and unpinned-revision paths.

Two smaller hardening notes on the same function, offered only because they are adjacent — happy to
split them out:

• kv.py:179 catches (EOFError, pickle.UnpicklingError, Exception) and resets to {} with no log,
  so a corrupt or rejected cache is indistinguishable from an empty one. A warning here would make
  any of the above visible in production.
• Cache identity checking is complementary to, not a substitute for, the pickle.load hardening
  already tracked in activation cache uses pickle.load (unsafe deserialization) — correction: the zip-slip vector does NOT reproduce #2203 / PR Fix #2203: skill-memory zip extractall is path-traversal (zip-slip); activation cache uses  #2204.

[Memtensor-AI]

🤖 AutoDev has picked up this issue and started working on it.

Task ID: b7c23a4c04476531
Working branch: bugfix/autodev-2300-20260828172611164
Target branch: latest dev* branch
Workflow: opsp (analysis → coding → testing → PR)

I will post the PR link here once done. If I need more information, I will ask in the comments.

[Memtensor-AI]

✅ AutoDev task b7c23a4c04476531 completed.

Summary: Fixes issue #2300 — serialized activation KV caches now carry a producer fingerprint that is verified at load time, replacing the silent cross-model swap the issue reproduced with a clear diagnostic.

What changed. KVCacheMemory.extract / from_textual_memory and their VLLMKVCacheMemory counterparts now stamp metadata["producer"] with model_name_or_path, backend, config_fingerprint (sha256 of model.config.to_diff_dict()), tokenizer_fingerprint (sha256 of backend_tokenizer.to_str()), torch_dtype, quantization, plus soft hints (architectures, num_hidden_layers, num_kv_heads, head_dim, transformers_version). On load, every item's saved fingerprint is compared against a freshly-computed live one across the strict subset {model_name_or_path, config_fingerprint, tokenizer_fingerprint, torch_dtype, quantization}. Mismatched items are dropped with a single logger.error line naming every offending field and the item id; other items keep loading. Items with no producer key (pre-fix caches) are kept with a logger.warning — the deprecation window the issue requested. The silent except (EOFError, pickle.UnpicklingError, Exception): self.kv_cache_memories = {} at the load site now emits logger.warning(..., exc_info=True) so a corrupt cache is distinguishable from an empty one in production (issue's second suggestion).

Backward compatibility. Fully backward compatible — no config schema change, no on-disk file layout change (still a pickled dict at os.path.join(dir, memory_filename)); metadata simply gains one additive sub-key per item.

Tests. 14 new pytest cases under tests/memories/activation/test_kv.py and new tests/memories/activation/test_vllmkv.py cover: fingerprint capture on extract / from_textual_memory / on a partially-broken introspection source; matched load stays silent; config / tokenizer / dtype mismatch each drops the item with an ERROR line; mixed batch keeps the good item and drops only the bad one; missing-producer pre-fix cache installs with WARNING; corrupt pickle logs WARNING and resets to empty. Two pre-existing tests (test_get_cache_merge, test_delete_and_get_all) were failing on main against transformers 4.57.6 because they used the removed DynamicCache.key_cache API; rewritten to use DynamicCache.update. Full activation module = 25/25 pass; broader tests/memories/ + tests/llms/ = 125/125 pass; no regressions attributable to this change. Ruff clean.

Base branch: main
Branch: bugfix/autodev-2300-20260828172611164
Commit: fe4c4b7
PR: #2305
Assigned to: @Wang-Daoji
Reviewers: @bittergreen