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

[jeojdi1]

Draft. This is the largest of the four and it touches a design question, so per CONTRIBUTING.md ("For larger changes, please open an Issue or Discussion first to align with maintainers") I've now filed the problem as #2315 and am keeping this as a draft for direction rather than as a finished proposal. #2315 lists three possible fixes, cheapest first; this PR implements the third and most involved one. If maintainers prefer option 1 (fail loudly) or option 2 (re-prefill from source_text), say so on the issue and I'll close this in favour of a much smaller patch.

Description

transformers applies rotary embedding before writing to the cache:

# modeling_qwen2.py
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
...
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 encodes absolute position i of that fragment.

_concat_caches joins fragments with a bare torch.cat(dim=-2). That moves fragment 2's 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.

V is never rotated, which is what makes the diagnosis unambiguous: on a naive merge, the V halves match a correct reference bit for bit while the K halves diverge. That isolates the fault to positional phase and nothing else.

This PR rotates each fragment forward by the summed length of everything before it. Per pair the rotation composes exactly — R(a) then R(b) == R(a+b) — so it is a rewrite of phase, not an approximation.

Four cases are handled deliberately, each of which would be wrong if done naively:

case	handling
attention_scaling != 1 (YaRN-style)	delta rotation uses unscaled cos/sin — the stored key is already s·R(m)·k, so reapplying the scale gives s²
partial rotary dims	rotary_dim derived from inv_freq, so these work by construction
dynamic / NTK schedules	frequencies are recomputed as the sequence grows, so rotations don't compose — skipped with a warning rather than guessed at
non-HF backend, no reachable model	skips and warns; reproduces current behaviour, so this cannot break a working deployment

Scope — what this does not claim

This corrects 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. That is a separate and much smaller error than a wholesale phase mismatch, but it is real and I don't want to overstate the fix.

Related Issue (Required): #2315

Type of change

• Bug fix (non-breaking change which fixes an issue)

How Has This Been Tested?

• Unit Test

test_merge_rerotates_second_fragment_to_its_new_position builds two fragments the way the model would, merges them, and compares against keys rotated directly at the positions they land on — using transformers' own apply_rotary_pos_emb as the reference, deliberately not a second copy of this PR's arithmetic, since grounding against your own implementation only proves self-consistency.

Fails-before / passes-after on transformers 5.16.1:

# upstream
E  AssertionError: merged keys differ from keys rotated at their landing
   positions by 3.473e+00; fragment 2 is carrying the phase of positions
   0..3 instead of 5..8

# with this change
6 passed

test_merge_refuses_to_rerotate_non_composing_rope asserts a dynamic schedule is skipped with a warning and the fragments are still concatenated.

Checklist

• I have performed a self-review of my own code
• I have commented my code in hard-to-understand areas
• I have added tests that prove my fix is effective
• I have created related documentation issue/PR in MemOS-Docs (would suggest a note on merge semantics if this lands)
• I have linked the issue to this PR
• I have mentioned the person who will review this PR

A note on the target branch

CONTRIBUTING.md says to open PRs against dev, but no dev branch exists — only main and dev-v2.0.28 … dev-v2.0.32. This is against main (185ebdb, "Dev v2.0.32"). Happy to retarget.

Relationship to the other KV cache PRs

This is one of four independent fixes to src/memos/memories/activation/kv.py, each kept to a single logical change per CONTRIBUTING.md:

• fix: stop merged KV caches being silently discarded on transformers >= 4.57 #2314 — merged caches report length 0 and are discarded on transformers >= 4.57 (Merged KV caches report get_seq_length() == 0 and are silently discarded on transformers >= 4.57 #2313)
• fix: record and check model identity on serialized KV caches #2302 — serialized caches carry no model identity (Serialized activation KV cache carries no model identity, so a cache built by one model loads silently into another #2300)
• fix: stop get_cache handing out the stored KV cache by reference #2303 — get_cache hands out the stored cache by reference, so it grows every turn (The stored activation DynamicCache is mutated in place by generation, so activation memory grows every turn #2301)
• this PR (draft) — fragments are not re-rotated to their landing positions (KVCacheMemory._concat_caches merges KV fragments without re-rotating them to their new positions #2315)

#2313 and #2315 are both defects in _concat_caches but are independent: #2313 is version-gated and clear-cut, this one is version-independent and carries a design decision. They were split so the first is not held up by the second.

[jeojdi1]

Filed the underlying problem as #2315 and linked it above, per CONTRIBUTING.md's "open an Issue or Discussion first for larger changes".

The issue writes up the defect on its own terms, with a standalone reproduction and prior-art check, and lists three possible fixes cheapest-first: (1) make _concat_caches fail loudly for multi-fragment merges and document that activation memories must be merged at the text level, (2) re-prefill from the fragments' source_text, or (3) re-rotate, which is what this PR implements.

Keeping this as a draft. If you'd prefer option 1 or 2, say so on #2315 and I'll close this in favour of a much smaller patch — I'd rather not push a ~200-line rotary change if a two-line guard is the answer you want.

Note also that #2315 and #2313 are both defects in this same function but are independent, and were split deliberately: #2313 is version-gated and clear-cut (its fix is #2314), whereas this one carries a design decision.