diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py index 1981b958f..1f406dd62 100644 --- a/src/memos/memories/activation/kv.py +++ b/src/memos/memories/activation/kv.py @@ -6,6 +6,7 @@ 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 @@ -13,6 +14,78 @@ 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)) diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py index 6490d687f..e5b550e4f 100644 --- a/tests/memories/activation/test_kv.py +++ b/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]}" + )