Fix #2300: Serialized activation KV cache carries no model identity, so a cache built by on

src/memos/memories/activation/kv.py

@@ -1,18 +1,148 @@
+import hashlib
+import json
 import os
 import pickle
 
 from datetime import datetime
+from typing import Any
 
 from transformers import DynamicCache
 
 from memos.configs.memory import KVCacheMemoryConfig
 from memos.dependency import require_python_package
 from memos.llms.factory import LLMFactory
+from memos.log import get_logger
 from memos.memories.activation.base import BaseActMemory
 from memos.memories.activation.item import KVCacheItem
 from memos.memories.textual.item import TextualMemoryItem
 
 
+logger = get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Producer fingerprint helpers (see openspec change 2300-... for the contract)
+# ---------------------------------------------------------------------------
+
+# The strict subset of fingerprint fields we compare on load. A field mismatches
+# only if both sides are non-None and unequal — unknown fields never fail.
+_STRICT_FIELDS: tuple[str, ...] = (
+    "model_name_or_path",
+    "config_fingerprint",
+    "tokenizer_fingerprint",
+    "torch_dtype",
+    "quantization",
+)
+
+
+def _compute_producer_fingerprint(llm: Any) -> dict[str, Any]:
+    """Return an identity fingerprint for the LLM that produced a KV cache.
+
+    Every field is best-effort — a missing attribute becomes ``None`` rather
+    than an exception. The loader treats ``None`` as unknown, not as mismatch,
+    so a partial fingerprint stays useful.
+    """
+    fp: dict[str, Any] = {
+        "model_name_or_path": None,
+        "backend": None,
+        "config_fingerprint": None,
+        "tokenizer_fingerprint": None,
+        "torch_dtype": None,
+        "quantization": None,
+        "architectures": None,
+        "num_hidden_layers": None,
+        "num_kv_heads": None,
+        "head_dim": None,
+        "transformers_version": None,
+    }
+
+    llm_config = getattr(llm, "config", None)
+    if llm_config is not None:
+        fp["model_name_or_path"] = getattr(llm_config, "model_name_or_path", None)
+        fp["backend"] = type(llm_config).__name__
+
+    model = getattr(llm, "model", None)
+    if model is not None:
+        try:
+            model_config = getattr(model, "config", None)
+            if model_config is not None:
+                to_diff_dict = getattr(model_config, "to_diff_dict", None)
+                if callable(to_diff_dict):
+                    diff = to_diff_dict()
+                    fp["config_fingerprint"] = hashlib.sha256(
+                        json.dumps(diff, sort_keys=True, default=str).encode()
+                    ).hexdigest()
+                quant = getattr(model_config, "quantization_config", None)
+                fp["quantization"] = str(quant) if quant is not None else None
+                arch = getattr(model_config, "architectures", None)
+                if arch:
+                    fp["architectures"] = list(arch)
+                fp["num_hidden_layers"] = getattr(model_config, "num_hidden_layers", None)
+                _num_kv = getattr(model_config, "num_key_value_heads", None)
+                fp["num_kv_heads"] = (
+                    _num_kv
+                    if _num_kv is not None
+                    else getattr(model_config, "num_attention_heads", None)
+                )
+                head_dim = getattr(model_config, "head_dim", None)
+                if head_dim is None:
+                    hidden = getattr(model_config, "hidden_size", None)
+                    n_heads = getattr(model_config, "num_attention_heads", None)
+                    if hidden and n_heads:
+                        head_dim = hidden // n_heads
+                fp["head_dim"] = head_dim
+            dtype = getattr(model, "dtype", None)
+            if dtype is not None:
+                fp["torch_dtype"] = str(dtype)
+        except Exception:
+            logger.warning(
+                "Failed to introspect HF model config for KV cache fingerprint",
+                exc_info=True,
+            )
+
+    tokenizer = getattr(llm, "tokenizer", None)
+    if tokenizer is not None:
+        try:
+            backend_tokenizer = getattr(tokenizer, "backend_tokenizer", None)
+            if backend_tokenizer is not None and hasattr(backend_tokenizer, "to_str"):
+                fp["tokenizer_fingerprint"] = hashlib.sha256(
+                    backend_tokenizer.to_str().encode()
+                ).hexdigest()
+        except Exception:
+            logger.warning(
+                "Failed to compute tokenizer fingerprint for KV cache",
+                exc_info=True,
+            )
+
+    try:
+        import transformers
+
+        fp["transformers_version"] = transformers.__version__
+    except Exception:
+        logger.debug("Could not determine transformers version for fingerprint", exc_info=True)
+
+    return fp
+
+
+def _fingerprint_mismatch_reasons(saved: dict[str, Any] | None, live: dict[str, Any]) -> list[str]:
+    """Return a list of human-readable mismatch reasons.
+
+    Empty list means "safe to install" (either every field matched or one side
+    was unknown). Only strict fields participate in the decision.
+    """
+    if not isinstance(saved, dict):
+        return []
+    reasons: list[str] = []
+    for field in _STRICT_FIELDS:
+        s_val = saved.get(field)
+        l_val = live.get(field)
+        if s_val is None or l_val is None:
+            continue
+        if s_val != l_val:
+            reasons.append(f"{field}: saved={s_val!r} live={l_val!r}")
+    return reasons
+
+
 class KVCacheMemory(BaseActMemory):
     """
     Key-Value Cache Memory for activation memories.
@@ -29,6 +159,10 @@ def __init__(self, config: KVCacheMemoryConfig) -> None:
         self.llm = LLMFactory.from_config(config.extractor_llm)
         self.kv_cache_memories: dict[str, KVCacheItem] = {}
 
+    def _producer_fingerprint(self) -> dict[str, Any]:
+        """Fingerprint of the LLM currently wired into this memory."""
+        return _compute_producer_fingerprint(self.llm)
+
     def extract(self, text: str) -> KVCacheItem:
         """Extract memory based on the text.
 
@@ -46,7 +180,11 @@ def extract(self, text: str) -> KVCacheItem:
         # Create a KVCacheItem with the extracted cache
         cache_item = KVCacheItem(
             memory=kv_cache,
-            metadata={"source_text": text, "extracted_at": datetime.now().isoformat()},
+            metadata={
+                "source_text": text,
+                "extracted_at": datetime.now().isoformat(),
+                "producer": self._producer_fingerprint(),
+            },
         )
 
         return cache_item
@@ -134,7 +272,53 @@ def from_textual_memory(self, mem: TextualMemoryItem) -> KVCacheItem:
         """
         # Build KV cache from the textual memory content
         kv_cache = self.llm.build_kv_cache(mem.memory)
-        return KVCacheItem(memory=kv_cache, metadata=mem.metadata.model_dump())
+        metadata = mem.metadata.model_dump()
+        metadata["producer"] = self._producer_fingerprint()
+        return KVCacheItem(memory=kv_cache, metadata=metadata)
+
+    def _verify_and_filter(self, memories: dict[str, KVCacheItem]) -> dict[str, KVCacheItem]:
+        """Compare each item's saved producer fingerprint against the live LLM.
+
+        Drops items whose fingerprint disagrees on any strict field. Items
+        without a fingerprint are kept with a warning (backward compatibility
+        with pre-2.0.30 caches).
+        """
+        if not memories:
+            return {}
+
+        # Live fingerprint may itself fail to compute for an unusual backend.
+        # In that case skip verification with a warning — do not regress
+        # loading reliability.
+        try:
+            live_fp = self._producer_fingerprint()
+        except Exception:
+            logger.warning(
+                "Failed to compute live producer fingerprint; loading KV cache without verification",
+                exc_info=True,
+            )
+            return memories
+
+        verified: dict[str, KVCacheItem] = {}
+        for item_id, item in memories.items():
+            item_metadata = getattr(item, "metadata", None) or {}
+            saved_fp = item_metadata.get("producer") if isinstance(item_metadata, dict) else None
+            if saved_fp is None:
+                logger.warning(
+                    "KV cache item %s has no producer fingerprint (pre-2.0.30 cache); loading unchecked",
+                    item_id,
+                )
+                verified[item_id] = item
+                continue
+            reasons = _fingerprint_mismatch_reasons(saved_fp, live_fp)
+            if reasons:
+                logger.error(
+                    "KV cache item %s dropped: producer fingerprint mismatch (%s)",
+                    item_id,
+                    "; ".join(reasons),
+                )
+                continue
+            verified[item_id] = item
+        return verified
 
     def load(self, dir: str) -> None:
         """Load memories from os.path.join(dir, self.config.memory_filename)
@@ -163,21 +347,30 @@ def load(self, dir: str) -> None:
                     memories = data["kv_cache_memories"]
                     if isinstance(memories, list):
                         # Convert list to dict format
-                        self.kv_cache_memories = {item.id: item for item in memories}
+                        candidate = {item.id: item for item in memories}
                     else:
-                        self.kv_cache_memories = memories
+                        candidate = memories
+                    self.kv_cache_memories = self._verify_and_filter(candidate)
                 else:
                     # Reset to empty if no memories in data
                     self.kv_cache_memories = {}
             elif isinstance(data, list):
                 # Backward compatibility: convert list to dict
-                self.kv_cache_memories = {item.id: item for item in data}
+                candidate = {item.id: item for item in data}
+                self.kv_cache_memories = self._verify_and_filter(candidate)
             else:
                 # Reset to empty if data format is unexpected
                 self.kv_cache_memories = {}
 
-        except (EOFError, pickle.UnpicklingError, Exception):
-            # If loading fails, start with empty memories
+        except Exception:
+            # Corrupt or incompatible cache — log the reason so the failure is
+            # distinguishable from an empty cache in production. Loader stays
+            # resilient by resetting to an empty dict.
+            logger.warning(
+                "Failed to load KV cache from %s; resetting to empty",
+                file_path,
+                exc_info=True,
+            )
             self.kv_cache_memories = {}
 
     def dump(self, dir: str) -> None: