From 98ea915dfe7ebda5249e08a5a053305af876f1f9 Mon Sep 17 00:00:00 2001 From: Prashant Rawat Date: Wed, 19 Aug 2026 04:36:40 -0700 Subject: [PATCH] Remap every position in get_freqs, not just the chunk start Summary: RopeWithAttentionSink.get_freqs has two branches. The static-shape branch remaps every position and indexes the table with the result. The dynamic-shape branch remapped only the first position of the chunk and then took a contiguous slice from there: remapped_item = sink_size + (input_pos_item - sink_size) % ring_size freqs_cos = self.freqs_cos.narrow(0, remapped_item, seq_len) narrow hands row remapped_item + j to the j-th token, so the chunk is assigned remap(start) + j where it should be remap(start + j). At sink_size=4, window_size=8 the ring top is 20. A chunk of 5 starting at 18 covers true positions 18..22 and needs rows 18, 19, 4, 5, 6. The slice asks for 18, 19, 20, 21, 22. Rows 20..22 are outside the ring the cache actually uses, so those tokens get frequencies for positions no cache slot holds. The two branches therefore disagree with each other. On master, the same wrapping chunk returns rows 18, 19, 4, 5, 6 under static shapes and raises under dynamic ones, so the frequencies depend on enable_dynamic_shape. Differential Revision: D116138610 --- .../source_transformation/attention_sink.py | 41 ++--- .../test_attention_sink.py | 162 ++++++++++++++++++ 2 files changed, 181 insertions(+), 22 deletions(-) diff --git a/examples/models/llama/source_transformation/attention_sink.py b/examples/models/llama/source_transformation/attention_sink.py index cada2ac4e6d..c2d548fa798 100644 --- a/examples/models/llama/source_transformation/attention_sink.py +++ b/examples/models/llama/source_transformation/attention_sink.py @@ -37,8 +37,10 @@ class RopeWithAttentionSink(Rope): - Window tokens (pos >= sink_size): wrapped into ring buffer range [sink_size, sink_size + ring_size) via modulo - The ring buffer is 2x window_size, so the live window (window_size tokens) - never spans a wrap boundary, preserving correct relative distances in RoPE. + The ring buffer is 2x window_size for write-ahead headroom, not to keep the + live window contiguous -- it can span a wrap. Across a wrap two positions + remap to a difference that is not their true distance, so RoPE preserves + relative distance only within a wrap. """ def __init__( @@ -64,32 +66,27 @@ def get_freqs(self, input_pos: Optional[torch.Tensor], seq_len: int): """ Get rotary embedding frequencies with position remapping. - For dynamic shape mode (input_pos is a single start position), we remap - the start and use narrow. For static shape mode (input_pos is the full - position tensor), we remap all positions and index directly. + For dynamic shape mode input_pos is a single start position, expanded + here into the full chunk; for static shape mode it already is the full + position tensor. Either way every position is remapped and indexed. + Remapping only the start and slicing from there would branch on a + data-dependent value, which blocks export, and would read past the ring + on a chunk that wraps. """ assert input_pos is not None if not self.params.use_kv_cache: return self.freqs_cos[:seq_len], self.freqs_sin[:seq_len] if self.params.enable_dynamic_shape: - # Dynamic shape: input_pos is [start_pos], remap and narrow - input_pos_item = input_pos[-1].item() - if input_pos_item < self.sink_size: - remapped_item = input_pos_item - else: - remapped_item = ( - self.sink_size + (input_pos_item - self.sink_size) % self.ring_size - ) - torch._check_is_size(remapped_item) - torch._check(remapped_item + seq_len <= self.sink_size + self.ring_size) - freqs_cos = self.freqs_cos.narrow(0, remapped_item, seq_len) - freqs_sin = self.freqs_sin.narrow(0, remapped_item, seq_len) - else: - # Static shape: remap full position tensor and index - remapped = self._remap_input_pos(input_pos) - freqs_cos = self.freqs_cos[remapped] - freqs_sin = self.freqs_sin[remapped] + # Dynamic shape: input_pos is [start_pos], expand to the whole chunk + input_pos = input_pos[-1] + torch.arange( + seq_len, device=input_pos.device, dtype=input_pos.dtype + ) + + # Remap the full position tensor and index + remapped = self._remap_input_pos(input_pos) + freqs_cos = self.freqs_cos[remapped] + freqs_sin = self.freqs_sin[remapped] return freqs_cos, freqs_sin diff --git a/examples/models/llama/source_transformation/test_attention_sink.py b/examples/models/llama/source_transformation/test_attention_sink.py index 8cdb00951f2..c4105338f0e 100644 --- a/examples/models/llama/source_transformation/test_attention_sink.py +++ b/examples/models/llama/source_transformation/test_attention_sink.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import copy import unittest import torch @@ -53,6 +54,106 @@ def test_get_freqs_passthrough(self, input_pos, seq_len): torch.testing.assert_close(freqs_sin, expected_sin) +class RopeWithAttentionSinkWrapTest(unittest.TestCase): + """get_freqs over a chunk that crosses the top of the ring. + + The cases above all stay below the first wrap, where remapping is the + identity, so none of them can tell a per-position remap from a remapped + start plus a contiguous slice. These can. + """ + + SINK_SIZE = 4 + WINDOW_SIZE = 8 + + def setUp(self) -> None: + # Ring top is 20. The table is deliberately longer, so a slice running + # past the ring still lands on real rows instead of going out of bounds. + self.params = ModelArgs( + use_kv_cache=True, enable_dynamic_shape=True, max_context_len=64 + ) + self.rope = RopeWithAttentionSink( + params=self.params, + window_size=self.WINDOW_SIZE, + sink_size=self.SINK_SIZE, + ) + self.ring_top = self.SINK_SIZE + 2 * self.WINDOW_SIZE + + def test_a_chunk_crossing_the_ring_top_wraps(self) -> None: + """Ground truth: the rows are spelled out, not derived from the code. + + Every other test here checks the implementation against itself -- the + sweep below uses _remap_input_pos as its own oracle, and the two tests + after it compare two calls of the same code. None of them would notice + a wrong modulus or a wrong sink boundary. This one would. + """ + start, seq_len = 18, 5 + self.assertGreater(start + seq_len, self.ring_top) + + freqs_cos, freqs_sin = self.rope.get_freqs( + input_pos=torch.tensor([start], dtype=torch.int32), seq_len=seq_len + ) + + # Positions 18..22 with sink_size=4 and a ring of 16: 18 and 19 are + # still below the ring top, 20 is the first to come back around. + expected = [18, 19, 4, 5, 6] + torch.testing.assert_close(freqs_cos, self.rope.freqs_cos[expected]) + torch.testing.assert_close(freqs_sin, self.rope.freqs_sin[expected]) + + # A contiguous slice from the remapped start is a different answer, not + # merely a different spelling of this one. + sliced = self.rope.freqs_cos.narrow(0, start, seq_len) + self.assertFalse(torch.allclose(freqs_cos, sliced)) + + def test_a_position_gets_the_same_freqs_whatever_chunk_it_lands_in(self) -> None: + """Chunk size must not change a token's rotation. + + The only test that mixes chunk lengths. The sweep below is seq_len=5 + throughout, and single-token decode is the one shape that is correct + even without this change, so a disagreement between the two is what + the previous code produced and what a caller would actually hit. + """ + # True position 20 is the first position past the ring top. Ask for it + # as the third entry of a chunk, then as a chunk of its own. + in_chunk, _ = self.rope.get_freqs( + input_pos=torch.tensor([18], dtype=torch.int32), seq_len=5 + ) + alone, _ = self.rope.get_freqs( + input_pos=torch.tensor([20], dtype=torch.int32), seq_len=1 + ) + + torch.testing.assert_close(in_chunk[2], alone[0]) + + def test_every_chunk_across_four_ring_cycles_gathers_in_bounds(self) -> None: + """Breadth: every start across four ring cycles, and no index escapes. + + Consistency with _remap_input_pos rather than ground truth, so it + cannot catch a wrong remap -- it is here for the starts the other + tests do not name, and for the bound on the gathered indices. + """ + seq_len = 5 + table_len = self.rope.freqs_cos.shape[0] + wrapping = 0 + for start in range(4 * self.ring_top): + with self.subTest(start=start): + expected = self.rope._remap_input_pos( + torch.arange(start, start + seq_len) + ) + wrapping += int(bool((expected.diff() != 1).any())) + self.assertGreaterEqual(int(expected.min()), 0) + self.assertLess(int(expected.max()), table_len) + + freqs_cos, freqs_sin = self.rope.get_freqs( + input_pos=torch.tensor([start], dtype=torch.int32), seq_len=seq_len + ) + torch.testing.assert_close(freqs_cos, self.rope.freqs_cos[expected]) + torch.testing.assert_close(freqs_sin, self.rope.freqs_sin[expected]) + + # Guard the fixture: 4 of every 16 starts put a chunk of 5 across the + # ring top. If a config change made that 0 the sweep would still pass + # while testing nothing this diff is about. + self.assertEqual(wrapping, 16) + + class CachePositionsManagerWithSinkTest(unittest.TestCase): def test_sink_indices_fixed(self): @@ -380,6 +481,23 @@ def _run_generation(self, model, args, num_tokens): return outputs + def _feed_in_chunks(self, model, tokens, chunk_size): + """Feed a fixed token sequence through the model chunk_size at a time. + + Returns one output per chunk. chunk_size=1 is the decode loop; anything + larger is a chunked prefill, which is the only way to get a multi-token + chunk at a start position other than 0. + """ + outputs = [] + with torch.no_grad(): + for pos in range(0, tokens.shape[1], chunk_size): + result = model( + tokens=tokens[:, pos : pos + chunk_size], + attn_options={"input_pos": torch.tensor([pos], dtype=torch.long)}, + ) + outputs.append(result[0] if isinstance(result, tuple) else result) + return outputs + def test_beyond_context_window_basic(self): """Generate tokens well beyond the KV cache size using standard SDPA.""" sink_size = 4 @@ -416,6 +534,50 @@ def test_beyond_max_context_len(self): "Output contains non-finite values beyond max_context_len", ) + def test_chunked_prefill_across_the_ring_wrap(self): + """Chunked prefill where a chunk spans the ring wrap. + + sink_size=4, window_size=16, so the ring is slots [4, 36). Feeding 5 + tokens at a time puts chunk starts at 0, 5, ..., 95. The chunk at 35 + covers positions 35..39 and needs rows 35, 4, 5, 6, 7; the chunk at 65 + covers 65..69 and needs rows 33, 34, 35, 4, 5. Both span the wrap. + + The other beyond-context-window tests decode one token at a time, and a + chunk of one can never span the wrap however far the position runs, so + none of them reach this. + + The RopeWithAttentionSinkWrapTest cases call get_freqs directly at + window_size=8. This is the only one that goes through the model, and + the only one at window_size=16, so it also covers the mask and the KV + cache agreeing with the remapped frequencies rather than get_freqs + alone. + """ + sink_size = 4 + window_size = 16 + args = self._make_args(max_context_len=64) + + torch.manual_seed(0) + model = self._build_model(args, sink_size, window_size) + tokens = torch.randint(0, args.vocab_size, (1, 100)) + + chunked = self._feed_in_chunks(copy.deepcopy(model), tokens, chunk_size=5) + one_at_a_time = self._feed_in_chunks(copy.deepcopy(model), tokens, chunk_size=1) + + self.assertEqual(len(chunked), 20) + self.assertEqual(len(one_at_a_time), 100) + + # generate_full_logits is off, so each call returns logits for its last + # position only. How the input was chunked must not change the result: + # chunk i ends on the same token as one_at_a_time[5 * i + 4]. + for i, out in enumerate(chunked): + self.assertTrue(torch.isfinite(out).all(), f"chunk {i} is not finite") + torch.testing.assert_close( + out, + one_at_a_time[5 * i + 4], + msg=lambda m, i=i: f"chunk {i}, positions {5 * i}..{5 * i + 4}, " + f"disagrees with feeding the same tokens one at a time:\n{m}", + ) + def test_beyond_context_window_custom_sdpa(self): """Generate tokens beyond context window with custom SDPA + custom KV cache.""" sink_size = 4