Fix SinkCache RoPE coordinates for absolute generation

## Problem

In Transformers 4.52.1, the standard generation loop advances Qwen2 RoPE
positions absolutely. The upstream SinkCache truncates its KV window but also
rerotates every surviving key into a bounded local slot. Its stored key phase
therefore no longer matches the new query's absolute phase.

## Change

- The custom `generate()` wrapper selects absolute-position mode. A shift
retains each surviving key's original RoPE phase while dropping old K/V.
- Direct callers can select `position_mode="local"` when both `position_ids`
and `cache_position` intentionally use bounded slots. That path computes
rerotation in FP32, normalizes the derived twiddle, then writes BF16 keys.
- Appending exactly to capacity no longer evicts a token early.
- README documents the position contract and the tested Transformers version.

## Verification

- Qwen2.5-1.5B-Instruct, BF16, eager attention, Transformers 4.52.1, RTX 4080;
two cache arms interleaved against the same frozen model and WikiText-2
teacher-forced token stream. Prefill logits matched bitwise.
- 2,600 tokens, four sinks, 256 recent positions, absolute positions:
2,339 evictions; upstream 65,492 rerotations and 27.840 post-fill PPL;
corrected cache zero rerotations and 19.227 PPL. Every reported block
favored the correction.
- On the identical model/data/tokens/prefill with bounded local positions,
correctly rerotating the keys gave 19.268 PPL upstream and 19.236 with
FP32 normalized rerotation. The corrected absolute and correctly local
contracts therefore agree closely.
- At a 1,024-recent-token window and 20,000 tokens with bounded local
positions, upstream and FP32 normalized rerotation were nearly tied:
10.122 versus 10.136 post-fill PPL. The numerical change improves an
isolated retained key's fidelity but is not a demonstrated next-token
quality gain under that local contract.
- Eleven regression tests pass under Transformers 4.52.1, including native
Qwen2 generation, exact capacity fill, and repeated BF16 key shifts.
- A paired 20,000-token, 1,024-recent-token absolute-position follow-up is
queued on the same GPU; I will add its result to this PR when it completes.

## Scope

The paired model-forward test reproduces the relevant Qwen2 generation
position/cache contract while keeping the input tokens identical between
arms. It is not a throughput benchmark. This historical custom cache is not
compatible with the changed Cache API in Transformers 5.x; the README now
states the tested version.

README.md

@@ -11,7 +11,22 @@ This is done by always keeping the first few tokens ("sink tokens") in the KV ca
 amount of attention to them. As it discards past non-sink tokens, the model will lose the ability to generate tokens
 that depend on the context that was discarded. It's also a solution to contain the memory footprint of the KV cache.
 
-This implementation matches the `SinkCache` class present in `transformers<4.53.0`.
+This implementation is based on the `SinkCache` class present in `transformers<4.53.0`.
+The behavior below was verified with `transformers==4.52.1`. The `Cache` API in
+Transformers 5.x is incompatible with this class; the example is not currently
+supported on that version.
+
+### RoPE positions when the cache shifts
+
+The standard `model.generate()` path advances RoPE positions absolutely. When
+the window discards old entries, surviving keys keep their original rotation;
+only the stored K/V entries move. Direct callers that intentionally reuse
+bounded local `position_ids` and `cache_position` instead need the surviving
+keys rerotated to their new local positions. Keep those two position inputs in
+the same coordinate system. The custom `generate()` wrapper sets absolute mode
+explicitly. Direct `SinkCache` users can pass `position_mode="absolute"` or
+`position_mode="local"`; the default `"auto"` infers the mode at first eviction
+from `cache_position` and assumes positions started at zero.
 
 ![Sink Cache diagram from the original paper](https://arxiv.org/html/2309.17453v4/x1.png)
 
@@ -27,7 +42,8 @@ This implementation matches the `SinkCache` class present in `transformers<4.53.
 
 
 ## Additional Arguments
-- `window_length` (`int`, *optional*, defaults to 256): The length of the context window.
+- `window_length` (`int`, *optional*, defaults to 256): The total KV cache capacity,
+  including `num_sink_tokens`. For four sinks and 1,024 recent tokens, set this to 1,028.
 - `num_sink_tokens` (`int`, *optional*, defaults to 4): The number of sink tokens. See the original paper for more information.
 
 
@@ -41,7 +57,7 @@ in `generate.py`, in this repository.
 We can use the custom generation method in this repository like the the base `generate` from `transformers`:
 
 ```py
-# requires `transformers>=4.52.0`
+# verified with `transformers==4.52.1`
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Preparing model, tokenizer, and model inputs

custom_generate/generate.py

@@ -26,13 +26,19 @@ class SinkCache(Cache):
 
     Parameters:
         window_length (`int`):
-            The length of the context window.
+            The total cache capacity, including sink tokens.
         num_sink_tokens (`int`):
             The number of sink tokens. See the original paper for more information.
+        position_mode (`str`, *optional*, defaults to `"auto"`):
+            `"absolute"` keeps retained RoPE keys at their original phases,
+            `"local"` rerotates them after each shift, and `"auto"` infers the
+            mode at the first eviction from `cache_position`.
     """
 
-    def __init__(self, window_length: int, num_sink_tokens: int) -> None:
+    def __init__(self, window_length: int, num_sink_tokens: int, position_mode: str = "auto") -> None:
         super().__init__()
+        if position_mode not in ("auto", "absolute", "local"):
+            raise ValueError("position_mode must be 'auto', 'absolute', or 'local'")
         self.key_cache: List[torch.Tensor] = []
         self.value_cache: List[torch.Tensor] = []
         self.window_length = window_length
@@ -40,6 +46,7 @@ class SinkCache(Cache):
         self.cos_sin_rerotation_cache = {}
         self._cos_cache = None
         self._sin_cache = None
+        self._global_positions = {"auto": None, "absolute": True, "local": False}[position_mode]
 
     @staticmethod
     def _rotate_half(x):
@@ -50,14 +57,18 @@ class SinkCache(Cache):
     def _apply_key_rotary_pos_emb(
         self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
     ) -> torch.Tensor:
+        # A retained key can be rotated hundreds of times. Keep the arithmetic in fp32;
+        # rounding the one-step rotation to bf16 before multiplying makes its norm drift.
+        key_dtype = key_states.dtype
+        key_states = key_states.to(torch.float32)
         rotated_key_states = (key_states * cos) + (self._rotate_half(key_states) * sin)
-        return rotated_key_states
+        return rotated_key_states.to(key_dtype)
 
     def _get_rerotation_cos_sin(
         self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         if key_states.shape[-2] not in self.cos_sin_rerotation_cache:
-            # Upcast to float32 temporarily for better accuracy
+            # Keep the rerotation coefficients in float32 through key rotation.
             cos = cos.to(torch.float32)
             sin = sin.to(torch.float32)
 
@@ -68,10 +79,16 @@ class SinkCache(Cache):
             shifted_sin = sin[self.num_sink_tokens : -key_states.shape[-2]]
             rerotation_cos = original_cos * shifted_cos + original_sin * shifted_sin
             rerotation_sin = -original_sin * shifted_cos + original_cos * shifted_sin
+            # The supplied table may already be rounded to bf16 (or include a
+            # RoPE attention scale). Its derived delta is not exactly unit length.
+            # Reusing that delta on every eviction compounds the magnitude error.
+            inv_norm = torch.rsqrt(rerotation_cos.square() + rerotation_sin.square())
+            rerotation_cos = rerotation_cos * inv_norm
+            rerotation_sin = rerotation_sin * inv_norm
 
             self.cos_sin_rerotation_cache[key_states.shape[-2]] = (
-                rerotation_cos.to(key_states.dtype).unsqueeze(0),
-                rerotation_sin.to(key_states.dtype).unsqueeze(0),
+                rerotation_cos.unsqueeze(0),
+                rerotation_sin.unsqueeze(0),
             )
         return self.cos_sin_rerotation_cache[key_states.shape[-2]]
 
@@ -140,19 +157,29 @@ class SinkCache(Cache):
             self.key_cache.append(key_states)
             self.value_cache.append(value_states)
 
-        elif key_states.shape[-2] + self.get_seq_length(layer_idx) < self.window_length:
+        elif key_states.shape[-2] + self.get_seq_length(layer_idx) <= self.window_length:
             # Growing cache
             self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
             self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
 
         else:
             # Shifting cache
+            # GenerationMixin advances cache_position globally. With absolute
+            # RoPE positions, retained keys already have the correct phase;
+            # rerotating them would put them in a different coordinate system
+            # from the new query. Explicit bounded-local position streams still
+            # need rerotation after each eviction.
+            if layer_idx == 0 and self._global_positions is None:
+                cache_position = cache_kwargs.get("cache_position")
+                self._global_positions = bool(
+                    cache_position is not None and cache_position[-1].item() >= self.window_length
+                )
             keys_to_keep = self.key_cache[layer_idx][
                 :, :, -self.window_length + self.num_sink_tokens + key_states.shape[-2] :
             ]
 
             # On RoPE models, we need to recompute the Key rotation as the tokens are shifted
-            if using_rope:
+            if using_rope and not self._global_positions:
                 rerotation_cos, rerotation_sin = self._get_rerotation_cos_sin(
                     key_states, self._cos_cache[: self.window_length], self._sin_cache[: self.window_length]
                 )
@@ -222,7 +249,9 @@ def generate(model, window_length=256, num_sink_tokens=4, **kwargs):
     # 2.a. prepare the cache, if it was not passed.
     past_key_values = kwargs.pop("past_key_values", None)
     if past_key_values is None:
-        past_key_values = SinkCache(window_length=window_length, num_sink_tokens=num_sink_tokens)
+        past_key_values = SinkCache(
+            window_length=window_length, num_sink_tokens=num_sink_tokens, position_mode="absolute"
+        )
     elif not isinstance(past_key_values, SinkCache):
         raise ValueError(f"`past_key_values` must be a `SinkCache` instance, got a {type(past_key_values)} instance")
 

tests/test_sink_cache_precision.py

@@ -0,0 +1,243 @@
+import torch
+from transformers import Qwen2Config, Qwen2ForCausalLM
+
+from custom_generate.generate import SinkCache, generate
+
+
+def test_bf16_rerotation_is_unit_length_and_keeps_cache_dtype():
+    window_length = 1024
+    head_dim = 128
+    positions = torch.arange(window_length, dtype=torch.float32)
+    inv_freq = 1_000_000.0 ** (-torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+    angles = positions[:, None] * inv_freq[None, :]
+    cos = torch.cat((angles.cos(), angles.cos()), dim=-1).to(torch.bfloat16)
+    sin = torch.cat((angles.sin(), angles.sin()), dim=-1).to(torch.bfloat16)
+
+    cache = SinkCache(window_length=window_length, num_sink_tokens=4)
+    keys = torch.ones((1, 2, 1, head_dim), dtype=torch.bfloat16)
+    rerotation_cos, rerotation_sin = cache._get_rerotation_cos_sin(keys, cos, sin)
+
+    assert rerotation_cos.dtype == torch.float32
+    assert rerotation_sin.dtype == torch.float32
+    squared_norm = rerotation_cos.square() + rerotation_sin.square()
+    torch.testing.assert_close(squared_norm, torch.ones_like(squared_norm), atol=5e-7, rtol=0)
+
+    shifted = cache._apply_key_rotary_pos_emb(
+        keys.expand(1, 2, window_length - 5, head_dim), rerotation_cos, rerotation_sin
+    )
+    assert shifted.dtype == torch.bfloat16
+    assert torch.isfinite(shifted).all()
+
+
+def test_repeated_shift_preserves_a_key_close_to_direct_rope():
+    # A resident key walks from slot 512 to slot 4. The reference rotates its
+    # untouched pre-RoPE key directly at the final slot, without repeated writes.
+    torch.manual_seed(7)
+    capacity, sinks, head_dim = 1028, 4, 128
+    positions = torch.arange(capacity, dtype=torch.float32)
+    inv_freq = 1_000_000.0 ** (-torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+    angles = positions[:, None] * inv_freq[None, :]
+    cos = torch.cat((angles.cos(), angles.cos()), dim=-1).to(torch.bfloat16)
+    sin = torch.cat((angles.sin(), angles.sin()), dim=-1).to(torch.bfloat16)
+    raw = torch.randn(64, head_dim).to(torch.bfloat16)
+
+    def rotate_half(value):
+        half = value.shape[-1] // 2
+        return torch.cat((-value[..., half:], value[..., :half]), dim=-1)
+
+    def rotate_once(value, position):
+        return (value * cos[position] + rotate_half(value) * sin[position]).to(torch.bfloat16)
+
+    cache = SinkCache(window_length=capacity, num_sink_tokens=sinks)
+    twiddle_cos, twiddle_sin = cache._get_rerotation_cos_sin(
+        torch.empty(1, 1, 1, head_dim, dtype=torch.bfloat16), cos, sin
+    )
+    current = rotate_once(raw, 512)
+    for destination in range(511, sinks - 1, -1):
+        offset = destination - sinks
+        current = cache._apply_key_rotary_pos_emb(
+            current, twiddle_cos[0, offset], twiddle_sin[0, offset]
+        )
+    direct = rotate_once(raw, sinks).float()
+    relative_error = ((current.float() - direct).norm(dim=-1) / direct.norm(dim=-1)).mean()
+    norm_ratio = (current.float().norm(dim=-1) / direct.norm(dim=-1)).mean()
+    assert relative_error < 0.16
+    assert 0.95 < norm_ratio < 1.05
+
+
+def test_update_applies_rope_shift_to_retained_keys():
+    cache = SinkCache(window_length=8, num_sink_tokens=1)
+    positions = torch.arange(8, dtype=torch.float32)
+    angles = positions[:, None] * torch.tensor([1.0, 0.1])
+    cos = torch.cat((angles.cos(), angles.cos()), dim=-1).to(torch.bfloat16)
+    sin = torch.cat((angles.sin(), angles.sin()), dim=-1).to(torch.bfloat16)
+    old_keys = torch.arange(32, dtype=torch.float32).reshape(1, 1, 8, 4).to(torch.bfloat16)
+    values = torch.zeros_like(old_keys)
+    cache.update(old_keys, values, 0, {"cos": cos, "sin": sin})
+    old_window = cache.key_cache[0][..., 2:, :].clone()
+
+    new_key = torch.ones((1, 1, 1, 4), dtype=torch.bfloat16)
+    cache.update(new_key, torch.zeros_like(new_key), 0, {"cos": cos, "sin": sin})
+
+    assert cache.get_seq_length() == 8
+    assert cache.key_cache[0].dtype == torch.bfloat16
+    assert torch.equal(cache.key_cache[0][..., :1, :], old_keys[..., :1, :])
+    assert torch.equal(cache.key_cache[0][..., -1:, :], new_key)
+    assert not torch.equal(cache.key_cache[0][..., 1:-1, :], old_window)
+
+
+def test_absolute_positions_keep_retained_keys_in_original_rope_frame():
+    cache = SinkCache(window_length=8, num_sink_tokens=1)
+    old_keys = torch.arange(32, dtype=torch.float32).reshape(1, 1, 8, 4).to(torch.bfloat16)
+    values = torch.zeros_like(old_keys)
+    cos = torch.arange(12, dtype=torch.float32)[:, None].expand(12, 4).cos().to(torch.bfloat16)
+    sin = torch.arange(12, dtype=torch.float32)[:, None].expand(12, 4).sin().to(torch.bfloat16)
+    cache.update(old_keys, values, 0, {"cos": cos, "sin": sin, "cache_position": torch.arange(8)})
+
+    new_key = torch.ones((1, 1, 1, 4), dtype=torch.bfloat16)
+    cache.update(new_key, torch.zeros_like(new_key), 0,
+                 {"cos": cos, "sin": sin, "cache_position": torch.tensor([8])})
+
+    assert cache._global_positions is True
+    assert cache.get_seq_length() == 8
+    assert torch.equal(cache.key_cache[0][..., :1, :], old_keys[..., :1, :])
+    assert torch.equal(cache.key_cache[0][..., 1:-1, :], old_keys[..., 2:, :])
+    assert torch.equal(cache.key_cache[0][..., -1:, :], new_key)
+    assert not cache.cos_sin_rerotation_cache
+
+
+def test_explicit_absolute_mode_does_not_guess_from_position_index():
+    cache = SinkCache(window_length=8, num_sink_tokens=1, position_mode="absolute")
+    old_keys = torch.arange(32, dtype=torch.float32).reshape(1, 1, 8, 4)
+    cache.update(old_keys, old_keys, 0)
+    new_key = torch.ones((1, 1, 1, 4))
+    cos = torch.ones((8, 4))
+    sin = torch.zeros((8, 4))
+    cache.update(new_key, new_key, 0,
+                 {"cos": cos, "sin": sin, "cache_position": torch.tensor([7])})
+    assert cache._global_positions is True
+    assert torch.equal(cache.key_cache[0][..., 1:-1, :], old_keys[..., 2:, :])
+    assert not cache.cos_sin_rerotation_cache
+
+
+def test_exactly_filling_capacity_does_not_evict():
+    cache = SinkCache(window_length=8, num_sink_tokens=1)
+    old_keys = torch.arange(28, dtype=torch.float32).reshape(1, 1, 7, 4)
+    cache.update(old_keys, old_keys, 0)
+    new_key = torch.ones((1, 1, 1, 4))
+    cache.update(new_key, new_key, 0)
+    assert cache.get_seq_length() == 8
+    assert torch.equal(cache.key_cache[0][..., :7, :], old_keys)
+    assert torch.equal(cache.key_cache[0][..., 7:, :], new_key)
+
+
+def test_qwen2_native_forward_supplies_rope_and_runs_shifts():
+    config = Qwen2Config(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+        num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=64,
+        rope_theta=1_000_000.0,
+    )
+    config._attn_implementation = "eager"
+    model = Qwen2ForCausalLM(config).eval()
+    cache = SinkCache(window_length=8, num_sink_tokens=1)
+    with torch.inference_mode():
+        model(
+            input_ids=torch.arange(8).unsqueeze(0), past_key_values=cache,
+            cache_position=torch.arange(8), use_cache=True,
+        )
+        assert not cache.cos_sin_rerotation_cache
+        for token in range(8, 24):
+            model(
+                input_ids=torch.tensor([[token % 64]]), past_key_values=cache,
+                cache_position=torch.tensor([7]), use_cache=True,
+            )
+
+    assert cache.get_seq_length() == 8
+    assert 1 in cache.cos_sin_rerotation_cache
+    assert len(cache.key_cache) == config.num_hidden_layers
+
+
+def test_qwen2_generate_uses_absolute_positions_without_rerotation():
+    config = Qwen2Config(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+        num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=64,
+        rope_theta=1_000_000.0,
+    )
+    config._attn_implementation = "eager"
+    model = Qwen2ForCausalLM(config).eval()
+    cache = SinkCache(window_length=8, num_sink_tokens=1)
+    with torch.inference_mode():
+        outputs = model.generate(
+            input_ids=torch.arange(8).unsqueeze(0), max_new_tokens=5,
+            do_sample=False, past_key_values=cache, use_cache=True,
+        )
+    assert outputs.shape == (1, 13)
+    assert cache.get_seq_length() == 8
+    assert cache._global_positions is True
+    assert not cache.cos_sin_rerotation_cache
+
+
+def test_custom_generate_wrapper_selects_absolute_mode():
+    config = Qwen2Config(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+        num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=64,
+        rope_theta=1_000_000.0,
+    )
+    config._attn_implementation = "eager"
+    model = Qwen2ForCausalLM(config).eval()
+    with torch.inference_mode():
+        outputs = generate(
+            model, window_length=8, num_sink_tokens=1,
+            input_ids=torch.arange(8).unsqueeze(0), max_new_tokens=5,
+            do_sample=False, return_dict_in_generate=True,
+        )
+    cache = outputs.past_key_values
+    assert isinstance(cache, SinkCache)
+    assert cache._global_positions is True
+    assert not cache.cos_sin_rerotation_cache
+
+
+def test_explicit_absolute_positions_match_qwen2_default_forward():
+    config = Qwen2Config(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+        num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=64,
+        rope_theta=1_000_000.0,
+    )
+    config._attn_implementation = "eager"
+    model = Qwen2ForCausalLM(config).eval()
+    caches = [SinkCache(window_length=8, num_sink_tokens=1) for _ in range(2)]
+    with torch.inference_mode():
+        for cache in caches:
+            model(input_ids=torch.arange(8).unsqueeze(0), past_key_values=cache,
+                  cache_position=torch.arange(8), use_cache=True)
+        explicit = model(input_ids=torch.tensor([[8]]), past_key_values=caches[0],
+                         cache_position=torch.tensor([8]), position_ids=torch.tensor([[8]]),
+                         use_cache=True).logits
+        default = model(input_ids=torch.tensor([[8]]), past_key_values=caches[1],
+                        cache_position=torch.tensor([8]), use_cache=True).logits
+    torch.testing.assert_close(explicit, default, atol=0, rtol=0)
+
+
+def test_explicit_absolute_and_auto_modes_match_on_standard_stream():
+    config = Qwen2Config(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+        num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=64,
+        rope_theta=1_000_000.0,
+    )
+    config._attn_implementation = "eager"
+    model = Qwen2ForCausalLM(config).eval()
+    caches = [SinkCache(8, 1, "auto"), SinkCache(8, 1, "absolute")]
+    with torch.inference_mode():
+        for cache in caches:
+            model(input_ids=torch.arange(8).unsqueeze(0), past_key_values=cache,
+                  cache_position=torch.arange(8), use_cache=True)
+        for token in range(8, 20):
+            logits = []
+            for cache in caches:
+                logits.append(model(input_ids=torch.tensor([[token % 64]]),
+                                    past_key_values=cache, cache_position=torch.tensor([token]),
+                                    position_ids=torch.tensor([[token]]), use_cache=True).logits)
+            torch.testing.assert_close(logits[0], logits[1], atol=0, rtol=0)
+    assert caches[0]._global_positions is True
+    assert not caches[0].cos_sin_rerotation_cache
+    assert not caches[1].cos_sin_rerotation_cache