KVCacheMemory._concat_caches merges KV fragments without re-rotating them to their new positions

[jeojdi1]

Environment

• MemOS at main (measured at 9119efe, re-checked at 185ebdb — 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, so this still applies at HEAD)
• torch 2.13.0, transformers 5.15.1, CPU, fp32, Qwen/Qwen2.5-0.5B-Instruct
• Version-independent: unlike Merged KV caches report get_seq_length() == 0 and are silently discarded on transformers >= 4.57 #2313, this affects every supported transformers version.

Summary

_concat_caches (src/memos/memories/activation/kv.py:200-255) merges cache fragments with a bare torch.cat(..., dim=-2) per layer, and nothing re-rotates them.

Rotary embedding is applied to keys before they are written to the cache. In transformers v4.53.2 — the version pinned in poetry.lock:5760-5761 — src/transformers/models/qwen2/modeling_qwen2.py:

159:  query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
...
164:  key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

So a cached key for token i of a fragment carries the phase of absolute position i of that fragment. Concatenating fragment 2 behind fragment 1 moves those keys into slots L1 … L1+L2-1 while their phase still says 0 … L2-1. Attention encodes relative distance in the angle between q and k, so after a merge the model's sense of how far apart tokens are is wrong for every fragment after the first, and the error grows the further into the merged cache you go. The merged cache does not correspond to any prefill of any text.

V is never rotated, and that makes the diagnosis unambiguous: on a naive merge the V halves match a correct reference bit for bit while the K halves diverge. That localises the entire discrepancy to positional phase — it is not a numerical-tolerance question.

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)])   # same fragment twice
gold   = prefill(A, prefill(A))                     # same token ids, correct positions
mk, gk = merged.layers[0].keys, gold.layers[0].keys

print("K first  half == reference:", torch.equal(mk[:, :, :L, :], gk[:, :, :L, :]))
print("K second half == reference:", torch.equal(mk[:, :, L:, :], gk[:, :, L:, :]))
print("max|K_merged - K_ref| 2nd half:", (mk[:, :, L:, :] - gk[:, :, L:, :]).abs().max().item())
print("merged K: slots[0,L) bit-identical to slots[L,2L) in every layer:",
      all(torch.equal(l.keys[:, :, :L, :], l.keys[:, :, L:, :]) for l in merged.layers))
print("reference K: same check:",
      all(torch.equal(l.keys[:, :, :L, :], l.keys[:, :, L:, :]) for l in gold.layers))
print("V (never rotated) second half == reference:",
      torch.equal(merged.layers[0].values[:, :, L:, :], gold.layers[0].values[:, :, L:, :]))

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])
L = kv_mem.get(item1.id).memory.layers[0].keys.shape[-2]
k = merged.layers[0].keys
assert not torch.equal(k[:, :, :L, :], k[:, :, L:2 * L, :])   # fails when item1 and item2 are the same text

Observed

transformers 5.15.1
K first  half == reference: True
K second half == reference: False
max|K_merged - K_ref| 2nd half: 29.280248641967773
merged K: slots[0,L) bit-identical to slots[L,2L) in every layer: True
reference K: same check: False
V (never rotated) second half == reference: True

Two identical fragments placed at different offsets come out bit-identical in the merged cache, while a correct prefill of the same tokens makes them differ. That is the positional information being absent.

Expected

Keys of the second fragment should equal the keys the model would produce for those same tokens at their new positions — so the second half should match the reference, and two identical fragments at different offsets must not be bit-identical.

Why it matters

get_cache() is documented and shipped (examples/core_memories/kv_cache_memory.py:113). A user who merges two activation memories gets a cache whose positions are wrong for everything after the first fragment, with no error and no warning.

Prior art checked

• _concat_caches has been touched once, by PR fix bug when calling _concat_caches in kv.py #177 → fix: fix bug when calling _concat_caches in kv.py (from pr#177) #205 ("fix bug when calling _concat_caches in kv.py", merged 2025-08-01 as 5bfac009ff). That fixed only an AttributeError from a transformers API change and did not address position or phase.
• No issue or PR in this repository mentions RoPE, rotary, or positions in this file — title searches for rope, position, concat and kv cache merge return nothing relevant.

Suggested fix

Three options, cheapest first:

1. Fail loudly. If a correct merge is not intended to be supported, make _concat_caches raise for len(caches) > 1 (or warn prominently) and document that activation memories must be merged at the text level and re-prefilled. Two lines, and it removes the silent wrongness, which is the expensive part.
2. Re-prefill. Implement multi-id get_cache as: concatenate the fragments' metadata["source_text"] and call self.llm.build_kv_cache(...) once. Correct by construction; costs one prefill of text you already have.
3. Re-rotate. Rotate each fragment after the first from its local positions to its destination positions. This is what I've drafted in fix: re-rotate KV fragments to their landing positions when merging #2304 — the per-pair rotation composes exactly, so it is a phase rewrite rather than an approximation, and it handles attention_scaling, partial rotary dims, and refuses dynamic/NTK schedules (whose frequencies recompute with length, so rotations don't compose).

Scope note, stated deliberately: option 3 corrects positional phase only. Merging fragments remains an approximation of a single prefill regardless, because each fragment's hidden states were computed without the others in context. That residual is real and much smaller than a wholesale phase mismatch, and I don't want to overstate what a fix buys.

Related

A second, version-gated defect in the same function — the merged cache reporting get_seq_length() == 0 and being silently discarded on transformers >= 4.57 — is tracked separately in #2313 / #2314. The two are independent; #2313 is the clear-cut bug, this one carries a design decision, which is why they were split.