Merged KV caches report get_seq_length() == 0 and are silently discarded on transformers >= 4.57

[jeojdi1]

Environment

• MemOS at main (verified at 9119efe, re-checked at 185ebdb and still current — src/memos/memories/activation/kv.py is unchanged upstream since Scheduler: address some issues to run old scheduler example and kv cache example #797, 2025-12-30)
• torch 2.13.0, transformers 5.15.1, CPU, fp32, Qwen/Qwen2.5-0.5B-Instruct
• Version-gated: applies to transformers >= 4.57.0, which is inside the declared range in pyproject.toml:41 (>=4.51.3,<5.0.0). poetry.lock pins 4.53.2, where this does not occur.

Summary

_concat_caches (src/memos/memories/activation/kv.py:200-255) builds each merged layer with layer_cls() (line 228) and then assigns .keys / .values directly. A layer constructed that way never runs DynamicLayer.lazy_initialization, so is_initialized stays False.

From transformers v4.57.0 onward, DynamicLayer.get_seq_length() short-circuits on exactly 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 attention mask is then sized for the query alone, and DynamicLayer.update calls lazy_initialization, which resets keys/values to empty tensors — every merged token is dropped on the first forward pass.

There is no exception and no warning. Activation memory simply has no effect, and the only symptom is that the model behaves as if the memory had never been loaded. Below v4.57 the gate is self.keys is None, so direct assignment survives and the bug does not appear.

This is very likely the same thing as the test_get_cache_merge failure that PR #2204 describes as a "transformers-API-version issue unrelated to this patch" — it is a version issue, but it is silent data loss rather than a test nuisance.

Minimal reproduction

Needs only torch + transformers + a small model. The merge function is copied verbatim from kv.py so no MemOS install is required. Runs on CPU in under a minute.

import torch, transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

MODEL = "Qwen/Qwen2.5-0.5B-Instruct"

# --- verbatim from src/memos/memories/activation/kv.py:200-255 (`self` removed) ---
def _concat_caches(caches):
    assert caches, "Need at least one cache"
    if len(caches) == 1:
        return caches[0]
    merged = DynamicCache()
    if hasattr(caches[0], "layers"):
        num_layers = len(caches[0].layers)
        if not hasattr(merged, "layers"):
            merged.layers = []
        if num_layers > 0:
            layer_cls = type(caches[0].layers[0])
            while len(merged.layers) < num_layers:
                merged.layers.append(layer_cls())
        for layer in range(num_layers):
            keys = [c.layers[layer].keys for c in caches]
            vals = [c.layers[layer].values for c in caches]
            merged.layers[layer].keys = torch.cat(keys, dim=-2)
            merged.layers[layer].values = torch.cat(vals, dim=-2)
    elif hasattr(caches[0], "key_cache"):
        num_layers = len(caches[0].key_cache)
        for layer in range(num_layers):
            keys = [c.key_cache[layer] for c in caches]
            vals = [c.value_cache[layer] for c in caches]
            merged.key_cache.append(torch.cat(keys, dim=-2))
            merged.value_cache.append(torch.cat(vals, dim=-2))
    else:
        raise AttributeError("no layers / key_cache")
    return merged
# ----------------------------------------------------------------------------------

print("transformers", transformers.__version__)
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32).eval()

def prefill(ids, cache=None):
    with torch.no_grad():
        return model(input_ids=ids, use_cache=True, past_key_values=cache,
                     return_dict=True).past_key_values

A = tok("The user's name is Alice and she lives in Berlin.", return_tensors="pt").input_ids
L = A.shape[-1]

merged = _concat_caches([prefill(A), prefill(A)])
print("get_seq_length():", merged.get_seq_length(), "expected", 2 * L,
      "| is_initialized:", getattr(merged.layers[0], "is_initialized", "n/a"))

q = tok("\nQ: In which city does the user live?", return_tensors="pt",
        add_special_tokens=False).input_ids
print("K rows before forward:", merged.layers[0].keys.shape[-2],
      "-> after forward:", prefill(q, merged).layers[0].keys.shape[-2],
      "(query length was", q.shape[-1], ")")

The same thing through the MemOS API — examples/core_memories/kv_cache_memory.py step 6 already builds this object:

merged = kv_mem.get_cache([item1.id, item2.id])
assert merged.get_seq_length() == merged.layers[0].keys.shape[-2]   # fails on transformers >= 4.57

Observed

transformers 5.15.1
get_seq_length(): 0 expected 24 | is_initialized: False
K rows before forward: 24 -> after forward: 11 (query length was 11 )

24 cached rows go in; 11 come out, which is the query alone. The merged content is gone.

Expected

merged.get_seq_length() should equal L1 + L2, and the merged content should survive the first forward pass.

Why it matters

get_cache() is documented and shipped (examples/core_memories/kv_cache_memory.py:113). On any transformers >= 4.57 install — permitted by the declared dependency range — a user who merges two activation memories gets a cache that contributes nothing, with no error, no warning, and no log line. The failure is invisible: the model just answers as though no memory were attached.

Suggested fix

Initialize the layer through its public lazy_initialization path before assigning, guarded by hasattr so the older key_cache layout is untouched. I have a patch ready with a regression test that fails on current main and passes with the fix, and will open it against this issue.

Note that the test fixture in tests/memories/activation/test_kv.py also needs repairing to see this: make_filled_cache appends to cache.key_cache, which was removed in 4.57, so test_get_cache_merge and test_delete_and_get_all fail with AttributeError on any current install and never reach the layers path where the bug lives.

Related

A second, version-independent problem in the same function — merged fragments are not re-rotated to their landing positions — is tracked separately in draft PR #2304, and is deliberately kept apart from this one because it is a design question rather than a clear-cut bug.

[Memtensor-AI]

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

Task ID: d33e83230ed5145f
Working branch: bugfix/autodev-2313-20260901080550424
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.