fix: re-rotate KV fragments to their landing positions when merging

src/memos/memories/activation/kv.py

@@ -6,13 +6,86 @@
 from transformers import DynamicCache
 
 from memos.configs.memory import KVCacheMemoryConfig
+from memos.log import get_logger
 from memos.dependency import require_python_package
 from memos.llms.factory import LLMFactory
 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__)
+
+
+# RoPE re-rotation for cache merges.
+#
+# transformers applies rotary embedding BEFORE writing to the cache:
+#
+#     query_states, key_states = apply_rotary_pos_emb(query, key, cos, sin)
+#     past_key_value.update(key_states, value_states, ...)
+#
+# so a cached key for token i of a fragment encodes absolute position i OF THAT
+# FRAGMENT. Concatenating fragment 2 behind fragment 1 moves its keys to slots
+# L1..L1+L2-1 while their phase still says 0..L2-1. V is never rotated, which is
+# why a naive merge leaves V bit-exact against a reference while K diverges --
+# the fault is positional phase and nothing else.
+#
+# Fixing it means rotating each fragment forward by the offset it lands at. Per
+# pair the rotation composes exactly, R(a) then R(b) == R(a+b), which is what
+# makes this a rewrite of phase rather than an approximation.
+_COMPOSABLE_ROPE_TYPES = ("default", "linear")
+
+
+def _rope_inv_freq_and_type(model) -> tuple:
+    """(inv_freq, rope_type) for a HF causal LM, or (None, reason) if unavailable.
+
+    Reads the rotary embedding's own buffer rather than recomputing from config,
+    so any scaling already applied at init is respected.
+    """
+    rotary = getattr(getattr(model, "model", None), "rotary_emb", None)
+    if rotary is None:
+        rotary = getattr(model, "rotary_emb", None)
+    if rotary is None or getattr(rotary, "inv_freq", None) is None:
+        return None, "model exposes no rotary_emb.inv_freq"
+    rope_type = getattr(rotary, "rope_type", None)
+    if rope_type is None:
+        params = getattr(getattr(model, "config", None), "rope_parameters", None)
+        rope_type = (params or {}).get("rope_type", "default")
+    return rotary.inv_freq, rope_type
+
+
+def _rope_shift(keys, delta: int, inv_freq):
+    """Rotate already-RoPE'd keys forward by `delta` absolute positions.
+
+    ``keys`` is [..., seq, head_dim] and ``inv_freq`` is [rotary_dim / 2]. Only
+    the first ``rotary_dim`` channels are rotated, so models with a partial
+    rotary factor are handled by construction.
+
+    The delta rotation deliberately uses UNSCALED cos/sin. On a model with
+    ``attention_scaling != 1`` (YaRN-style) the stored key is already ``s * R(m) k``;
+    rotating with the scale reapplied would produce ``s^2``.
+    """
+    import torch
+
+    rotary_dim = int(inv_freq.shape[-1]) * 2
+    if rotary_dim > keys.shape[-1]:
+        raise ValueError(
+            f"inv_freq implies rotary_dim={rotary_dim} but head_dim={keys.shape[-1]}"
+        )
+    rot, passthrough = keys[..., :rotary_dim], keys[..., rotary_dim:]
+
+    freq = inv_freq.to(device=keys.device, dtype=torch.float32) * float(delta)
+    ang = torch.cat([freq, freq], dim=-1)
+    cos, sin = ang.cos().to(keys.dtype), ang.sin().to(keys.dtype)
+
+    half = rotary_dim // 2
+    x1, x2 = rot[..., :half], rot[..., half:]
+    rotated = torch.cat([-x2, x1], dim=-1)
+    out = rot * cos + rotated * sin
+    return out if passthrough.numel() == 0 else torch.cat([out, passthrough], dim=-1)
+
+
 class KVCacheMemory(BaseActMemory):
     """
     Key-Value Cache Memory for activation memories.
@@ -197,6 +270,59 @@ def dump(self, dir: str) -> None:
         with open(file_path, "wb") as f:
             pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
 
+    def _merge_inv_freq(self):
+        """inv_freq to re-rotate merged fragments with, or None to skip.
+
+        Returns None -- and warns -- when the model is unreachable or uses a rope
+        schedule whose rotations do not compose. Skipping reproduces the previous
+        (incorrect) behaviour rather than raising, so this change cannot break a
+        working deployment; the warning makes the case visible instead of silent.
+        """
+        model = getattr(self.llm, "model", None)
+        if model is None:
+            logger.warning(
+                "Merging %s KV cache fragments without positional re-rotation: the "
+                "configured LLM exposes no local model, so the rotary frequencies "
+                "are unavailable. Merged fragments after the first will carry the "
+                "phase of their original positions.",
+                "multiple",
+            )
+            return None
+
+        inv_freq, rope_type = _rope_inv_freq_and_type(model)
+        if inv_freq is None:
+            logger.warning(
+                "Merging KV cache fragments without positional re-rotation: %s.",
+                rope_type,
+            )
+            return None
+        if rope_type not in _COMPOSABLE_ROPE_TYPES:
+            logger.warning(
+                "Merging KV cache fragments without positional re-rotation: rope_type "
+                "%r recomputes its frequencies as the sequence grows, so a delta "
+                "rotation does not compose and re-rotating would be wrong in a "
+                "different way. Build a single cache for this model instead of "
+                "merging fragments.",
+                rope_type,
+            )
+            return None
+        return inv_freq
+
+    @staticmethod
+    def _rerotate_fragments(keys: list, inv_freq) -> list:
+        """Shift each fragment's keys to the offset it lands at after concat.
+
+        Fragment 0 keeps its phase; fragment k is rotated forward by the summed
+        length of everything before it. A no-op when inv_freq is None.
+        """
+        if inv_freq is None:
+            return keys
+        out, offset = [], 0
+        for frag in keys:
+            out.append(frag if offset == 0 else _rope_shift(frag, offset, inv_freq))
+            offset += frag.shape[-2]
+        return out
+
     def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
         """
         Faster concat merge: for each layer, gather all caches' tensors
@@ -208,6 +334,7 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
         if len(caches) == 1:
             return caches[0]
 
+        inv_freq = self._merge_inv_freq()
         merged = DynamicCache()
 
         # Check for new structure (layers)
@@ -231,6 +358,8 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
                 # gather all K and V for this layer
                 keys = [c.layers[layer].keys for c in caches]
                 vals = [c.layers[layer].values for c in caches]
+                # re-rotate each fragment to the position it will occupy
+                keys = self._rerotate_fragments(keys, inv_freq)
                 # single concat per layer
                 merged.layers[layer].keys = torch.cat(keys, dim=-2)
                 merged.layers[layer].values = torch.cat(vals, dim=-2)
@@ -243,6 +372,8 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
                 # gather all K and V for this layer
                 keys = [c.key_cache[layer] for c in caches]
                 vals = [c.value_cache[layer] for c in caches]
+                # re-rotate each fragment to the position it will occupy
+                keys = self._rerotate_fragments(keys, inv_freq)
                 # single concat per layer
                 merged.key_cache.append(torch.cat(keys, dim=-2))
                 merged.value_cache.append(torch.cat(vals, dim=-2))

tests/memories/activation/test_kv.py

@@ -1,3 +1,5 @@
+import logging
+
 from unittest.mock import MagicMock
 
 import pytest
@@ -33,11 +35,11 @@ def kv_memory(dummy_config):
         yield KVCacheMemory(dummy_config)
 
 
-def make_filled_cache():
-    # Create a DynamicCache with at least one dummy tensor layer
+def make_filled_cache(seq_len: int = 3, n_layers: int = 1):
+    """Create a DynamicCache with dummy tensors, on any transformers version."""
     cache = DynamicCache()
-    cache.key_cache.append(torch.zeros(1, 2, 3))
-    cache.value_cache.append(torch.zeros(1, 2, 3))
+    for layer_idx in range(n_layers):
+        cache.update(torch.zeros(1, 2, seq_len, 4), torch.zeros(1, 2, seq_len, 4), layer_idx)
     return cache
 
 
@@ -59,8 +61,11 @@ def test_get_cache_merge(kv_memory):
     merged = kv_memory.get_cache([item1.id, item2.id])
     assert isinstance(merged, DynamicCache)
     # Check the number of layers in merged key/value cache
-    assert len(merged.key_cache) == 1
-    assert len(merged.value_cache) == 1
+    if hasattr(merged, "layers"):
+        assert len(merged.layers) == 1
+    else:
+        assert len(merged.key_cache) == 1
+        assert len(merged.value_cache) == 1
 
 
 def test_delete_and_get_all(kv_memory):
@@ -84,3 +89,116 @@ class DummyTextualMemory:
     item = kv_memory.from_textual_memory(DummyTextualMemory())
     assert isinstance(item, KVCacheItem)
     assert item.metadata["bar"] == 1
+
+
+# --------------------------------------------------------------------------
+# RoPE re-rotation on merge
+# --------------------------------------------------------------------------
+
+def _rope_reference(k_raw, positions, inv_freq):
+    """Rotate raw keys to absolute `positions` using transformers' own kernel.
+
+    Grounding the test against the library's implementation rather than a second
+    copy of our own arithmetic -- otherwise the test only proves we are
+    self-consistent.
+    """
+    from transformers.models.qwen2.modeling_qwen2 import apply_rotary_pos_emb
+
+    ang = positions[:, None].float() * inv_freq[None, :].float()
+    ang = torch.cat([ang, ang], dim=-1)
+    cos, sin = ang.cos()[None], ang.sin()[None]
+    kt = k_raw[None, None]
+    _, out = apply_rotary_pos_emb(kt, kt, cos, sin)
+    return out[0, 0]
+
+
+def test_merge_rerotates_second_fragment_to_its_new_position(kv_memory):
+    """A merged fragment must carry the phase of where it LANDS, not where it was built.
+
+    transformers rotates keys before writing them to the cache, so a cached key
+    for token i of a fragment encodes absolute position i *of that fragment*.
+    Concatenating fragment 2 behind fragment 1 moves its keys to slots
+    L1..L1+L2-1 while their phase still says 0..L2-1.
+
+    This builds two fragments the way the model would, merges them, and compares
+    against keys rotated directly at the positions they end up occupying, using
+    transformers' own ``apply_rotary_pos_emb`` as the reference.
+
+    Scope: this checks POSITIONAL PHASE only. Merging fragments remains an
+    approximation of a single cache regardless, because each fragment's hidden
+    states were computed without the others in context -- but that is a separate
+    and much smaller error than a wholesale phase mismatch.
+    """
+    from unittest.mock import MagicMock
+
+    torch.manual_seed(0)
+    head_dim, n_heads, l1, l2 = 8, 2, 5, 4
+    inv_freq = 1_000_000.0 ** (-torch.arange(0, head_dim, 2).float() / head_dim)
+
+    k_raw_1 = torch.randn(n_heads, l1, head_dim)
+    k_raw_2 = torch.randn(n_heads, l2, head_dim)
+
+    # what each fragment stores: rotated at ITS OWN positions, from 0
+    frag1 = _rope_reference(k_raw_1, torch.arange(l1), inv_freq)[None]
+    frag2 = _rope_reference(k_raw_2, torch.arange(l2), inv_freq)[None]
+
+    # what a correct merge must produce: fragment 2 rotated at l1..l1+l2-1
+    want = torch.cat(
+        [
+            _rope_reference(k_raw_1, torch.arange(l1), inv_freq)[None],
+            _rope_reference(k_raw_2, torch.arange(l1, l1 + l2), inv_freq)[None],
+        ],
+        dim=-2,
+    )
+
+    c1, c2 = DynamicCache(), DynamicCache()
+    c1.update(frag1, torch.zeros_like(frag1), 0)
+    c2.update(frag2, torch.zeros_like(frag2), 0)
+
+    # a model stub exposing only what the merge needs
+    rotary = MagicMock()
+    rotary.inv_freq = inv_freq
+    rotary.rope_type = "default"
+    model = MagicMock()
+    model.model.rotary_emb = rotary
+    kv_memory.llm.model = model
+
+    merged = kv_memory._concat_caches([c1, c2])
+    got = merged.layers[0].keys if hasattr(merged, "layers") else merged.key_cache[0]
+
+    err = (got.float() - want.float()).abs().max().item()
+    assert err < 1e-4, (
+        f"merged keys differ from keys rotated at their landing positions by "
+        f"{err:.3e}; fragment 2 is carrying the phase of positions 0..{l2 - 1} "
+        f"instead of {l1}..{l1 + l2 - 1}"
+    )
+
+
+def test_merge_refuses_to_rerotate_non_composing_rope(kv_memory, caplog):
+    """A rope schedule that recomputes with length must not be silently re-rotated.
+
+    Dynamic/NTK schedules change the frequencies as the sequence grows, so
+    R(a) then R(b) != R(a+b) and a delta rotation would be wrong in a *different*
+    way. Skip and say so, rather than guessing.
+    """
+    from unittest.mock import MagicMock
+
+    rotary = MagicMock()
+    rotary.inv_freq = torch.ones(4)
+    rotary.rope_type = "dynamic"
+    model = MagicMock()
+    model.model.rotary_emb = rotary
+    kv_memory.llm.model = model
+
+    c1, c2 = make_filled_cache(seq_len=3), make_filled_cache(seq_len=2)
+    with caplog.at_level(logging.WARNING, logger="memos.memories.activation.kv"):
+        merged = kv_memory._concat_caches([c1, c2])
+
+    # Assert on the tensor, not get_seq_length(): on transformers >= 4.57 a merged
+    # cache reports 0 regardless, which is a separate bug fixed on its own branch.
+    got = merged.layers[0].keys if hasattr(merged, "layers") else merged.key_cache[0]
+    assert got.shape[-2] == 5, "the fragments should still be concatenated"
+    assert any("does not compose" in r.getMessage() for r in caplog.records), (
+        f"no warning about the non-composing schedule; got "
+        f"{[r.getMessage() for r in caplog.records]}"
+    )