[TEST] Attention sink

[kirklandsign]

Summary

Test plan

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/17252

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures, 27 Cancelled Jobs, 1 Pending, 4 Unrelated Failures

As of commit bbca983 with merge base 1eb3f9d ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner / linux-job (gh)
  >>> Lint for extension/llm/runner/text_llm_runner.cpp:
• pull / test-parakeet-xnnpack-linux / linux-job (gh)
  RuntimeError: Command docker exec -t d17c402bb121efd8a7751121e9284e6703d9b4c24ff24de7e453ed6d3a07c1c1 /exec failed with exit code 1
• Test Metal Backend / export-model-metal-artifact (nvidia, parakeet-tdt, quantized-int4-metal) / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• Test Metal Backend / test-metal-backend-modules / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 127

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• pull / android / run-emulator (gh)
  ##[error]The operation was canceled.
• pull / test-coreml-bc-macos (macos-m1-stable) / macos-job (gh)
• pull / test-coreml-bc-macos (macos-m2-stable) / macos-job (gh)
• pull / test-llama-runner-linux (fp32, xnnpack+custom+qe, linux.arm64.2xlarge, executorch-ubuntu-22.04-gc... / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-llama-runner-linux (fp32, xnnpack+custom+quantize_kv, linux.arm64.2xlarge, executorch-ubuntu... / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-llama-runner-linux (fp32, xnnpack+quantize_kv, linux.arm64.2xlarge, executorch-ubuntu-22.04-... / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-llama-runner-qnn-linux (fp32, qnn_16a16w, qnn) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-llama-runner-qnn-linux (fp32, qnn_8a8w, qnn) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-lora-linux / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-models-linux (emformer_join, portable, linux.4xlarge.memory) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-models-linux (ic4, portable, linux.4xlarge.memory) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-models-linux (llama3_2_vision_encoder, portable, linux.4xlarge.memory) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-models-linux (phi_4_mini, portable, linux.4xlarge.memory) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-models-linux (resnet50, portable, linux.2xlarge) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-models-linux-basic (vit, portable, cmake, linux.arm64.2xlarge, executorch-ubuntu-22.04-gcc11... / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-multimodal-linux (gemma3-4b) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-sqnr-static-llm-qnn-linux (smollm2_135m) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-static-llama-qnn-linux (stories_110m) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / unittest / linux / linux-job (gh)
  ##[error]The operation was canceled.
• pull / unittest / macos / macos-job (gh)
• pull / unittest-arm-backend-with-no-deps (test_pytest_models_tosa) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / unittest-arm-backend-with-no-deps (test_pytest_ops_tosa) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / unittest-buck / macos / macos-job (gh)
• pull / unittest-editable / linux / linux-job (gh)
  ##[error]The operation was canceled.
• pull / unittest-editable / macos / macos-job (gh)
• pull / unittest-nxp-neutron / linux-job (gh)
  ##[error]The operation was canceled.
• Test CUDA Builds (gh)

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / test-samsung-models-linux / linux-job (gh) (similar failure)
  ##[error]The operation was canceled.
• pull / unittest / windows / windows-job (gh) (matched win rule in flaky-rules.json)
  ##[error]The operation was canceled.
• pull / unittest-editable / windows / windows-job (gh) (matched win rule in flaky-rules.json)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-samsung-quantmodels-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[Copilot]

Pull request overview

This pull request implements a ring-buffer-based attention sink mechanism for LLM inference, enabling models to generate beyond their fixed context window by maintaining a sliding window of recent tokens plus a fixed set of "sink" tokens (typically the initial prompt tokens).

Changes:

• Rewrites attention_sink.py to use a torch.export-compatible ring buffer approach instead of explicit token eviction
• Updates C++ text runner to support sliding window validation and extended generation
• Adds comprehensive unit tests for the ring buffer implementation
• Updates evaluation code to support the new ring buffer approach
• Adds configuration for attention sink models

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
examples/models/llama/source_transformation/attention_sink.py	Complete rewrite to ring-buffer implementation with torch.export compatibility
examples/models/llama/source_transformation/test_attention_sink_ring_buffer.py	New comprehensive test suite for ring buffer attention sink
extension/llm/runner/text_llm_runner.cpp	Updates validation logic for sliding window and adds TEMPORARY unlimited generation workaround
examples/models/llama/eval_llama_lib.py	Updates assertions for ring buffer compatibility
examples/models/llama/source_transformation/custom_kv_cache.py	Adds skip logic for KVCacheWithAttentionSink to prevent incorrect replacement
examples/models/llama/model.py	Updates max_context_len validation and sets default RoPE length to 131072
examples/models/llama/config/llama_attention_sink.yaml	New configuration file for attention sink models
examples/models/llama/BUCK	Adds test target for attention sink ring buffer tests

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The entire file content has been compressed onto a single line, removing all newlines. This makes the code completely unreadable and unmaintainable. The file should be properly formatted with appropriate line breaks, indentation, and spacing according to C++ coding standards.

[Copilot]

The comment contains Chinese characters ("滑动窗口大小" meaning "sliding window size" and "每次驱逐 1 个 token" meaning "evict 1 token each time"). Comments should be in English for consistency with the rest of the codebase and to ensure all team members can understand the documentation.

Suggested change

	# window_size=124: 滑动窗口大小
	# eviction_batch_size=1: 每次驱逐 1 个 token
	# window_size=124: sliding window size
	# eviction_batch_size=1: evict 1 token at a time

[Copilot]

The TEMPORARY comment indicates this is test code that should not be merged to production. Setting max_new_tokens to 1000000 when it's <= 0 bypasses important validation logic and could lead to excessive memory usage or infinite generation loops. This workaround should either be removed or properly documented with a plan for resolution before merging.

[Copilot]

The PR title "[TEST] Attention sink" and empty PR description suggest this is experimental or work-in-progress code. The presence of TEMPORARY workarounds in the C++ code and significant implementation changes without proper documentation indicate this PR may not be ready for production merge. Consider adding a proper description explaining the purpose, approach, testing strategy, and any known limitations before merging.

[Copilot]

The entire file content has been compressed onto a single line, removing all newlines. This makes the code completely unreadable and unmaintainable. The file should be properly formatted with appropriate line breaks, indentation, and spacing according to Python coding standards (PEP 8).

Suggested change

	# Copyright (c) Meta Platforms, Inc. and affiliates.# All rights reserved.## This source code is licensed under the BSD-style license found in the# LICENSE file in the root directory of this source tree."""Tests for the ring-buffer based attention sink implementation.This tests the torch.export-compatible implementation that uses a ring bufferfor the sliding window rather than explicit token eviction.Usage: # Run with pytest python -m pytest examples/models/llama/source_transformation/test_attention_sink_ring_buffer.py -v # Or run directly python examples/models/llama/source_transformation/test_attention_sink_ring_buffer.py"""import unittestimport torchfrom executorch.examples.models.llama.model_args import ModelArgsfrom executorch.examples.models.llama.source_transformation.attention_sink import ( CachePositionsManagerWithSink, KVCacheWithAttentionSink, RopeWithAttentionSink, _create_causal_mask_for_attention_sink,)class CachePositionsManagerWithSinkTest(unittest.TestCase): """Test the cache positions manager for ring buffer indexing.""" def setUp(self): self.cache_size = 32 # Total cache size (e.g., sink_size + window_size * 2) self.manager = CachePositionsManagerWithSink(self.cache_size) def test_initial_positions_are_zero(self): """Cache positions should start as zeros.""" expected = torch.zeros(self.cache_size, dtype=torch.long) torch.testing.assert_close(self.manager.cache_positions, expected) def test_simple_update(self): """Test simple sequential update without wraparound.""" input_pos = torch.tensor([0], dtype=torch.long) seq_len = 5 indices = self.manager.calculate_positions_and_update_indices(input_pos, seq_len) # Should return indices 0, 1, 2, 3, 4 expected_indices = torch.tensor([0, 1, 2, 3, 4], dtype=torch.long) torch.testing.assert_close(indices, expected_indices) # Cache positions at those indices should be the original positions for i in range(5): self.assertEqual(self.manager.cache_positions[i].item(), i) def test_wraparound(self): """Test ring buffer wraparound.""" # Fill cache to position 30 input_pos = torch.tensor([0], dtype=torch.long) self.manager.calculate_positions_and_update_indices(input_pos, 30) # Add 5 more tokens at position 30 - should wrap around input_pos = torch.tensor([30], dtype=torch.long) indices = self.manager.calculate_positions_and_update_indices(input_pos, 5) # Indices should wrap: 30 % 32 = 30, 31, 0, 1, 2 expected_indices = torch.tensor([30, 31, 0, 1, 2], dtype=torch.long) torch.testing.assert_close(indices, expected_indices) def test_cache_positions_track_original_positions(self): """Cache positions should track which original position is at each index.""" # Fill with positions 0-31 input_pos = torch.tensor([0], dtype=torch.long) self.manager.calculate_positions_and_update_indices(input_pos, 32) # Now add position 32 which wraps to index 0 input_pos = torch.tensor([32], dtype=torch.long) self.manager.calculate_positions_and_update_indices(input_pos, 1) # Index 0 should now contain original position 32 self.assertEqual(self.manager.cache_positions[0].item(), 32)class CausalMaskTest(unittest.TestCase): """Test the causal mask generation for attention sink.""" def test_mask_allows_sink_tokens(self): """Sink tokens should always be visible (mask = 0).""" cache_size = 32 sink_size = 4 window_size = 14 # cache_size = sink_size + window_size * 2 # Create cache positions where positions 0-3 are sink tokens cache_positions = torch.arange(cache_size, dtype=torch.long) start_pos = 20 # Query at position 20 seq_len = 1 mask = _create_causal_mask_for_attention_sink( cache_positions, window_size, sink_size, start_pos, seq_len ) # Sink tokens (indices 0-3, original positions 0-3) should have mask = 0 for i in range(sink_size): self.assertEqual(mask[0, i].item(), 0.0, f"Sink token at index {i} should be visible") def test_mask_blocks_future_tokens(self): """Future tokens should be masked (-inf).""" cache_size = 32 sink_size = 4 window_size = 14 cache_positions = torch.arange(cache_size, dtype=torch.long) start_pos = 10 seq_len = 1 mask = _create_causal_mask_for_attention_sink( cache_positions, window_size, sink_size, start_pos, seq_len ) # Future tokens (positions > 10) should have mask = -inf for i in range(11, cache_size): self.assertEqual(mask[0, i].item(), float('-inf'), f"Future token at position {i} should be masked") def test_mask_respects_window(self): """Tokens outside the window should be masked.""" cache_size = 32 sink_size = 4 window_size = 5 # Only allow 5 recent tokens cache_positions = torch.arange(cache_size, dtype=torch.long) start_pos = 20 seq_len = 1 mask = _create_causal_mask_for_attention_sink( cache_positions, window_size, sink_size, start_pos, seq_len ) # Positions 16-20 should be visible (within window of 5) for pos in range(16, 21): self.assertEqual(mask[0, pos].item(), 0.0, f"Position {pos} should be visible (in window)") # Position 15 should be masked (outside window, not a sink) self.assertEqual(mask[0, 15].item(), float('-inf'), f"Position 15 should be masked (outside window)")class KVCacheWithAttentionSinkTest(unittest.TestCase): """Test the KV cache with attention sink.""" def setUp(self): torch.manual_seed(42) self.window_size = 14 self.sink_size = 4 self.n_heads = 8 self.head_dim = 64 self.max_batch_size = 1 # Create model args with enough context for RoPE self.params = ModelArgs( use_kv_cache=True, enable_dynamic_shape=True, max_context_len=1024, # Large enough for RoPE n_heads=self.n_heads, n_kv_heads=self.n_heads, dim=self.n_heads * self.head_dim, ) self.rope = RopeWithAttentionSink( params=self.params, window_size=self.window_size, sink_size=self.sink_size, eviction_batch_size=1, ) self.kv_cache = KVCacheWithAttentionSink( n_heads=self.n_heads, head_dim=self.head_dim, enable_dynamic_shape=True, rope=self.rope, max_batch_size=self.max_batch_size, window_size=self.window_size, sink_size=self.sink_size, eviction_batch_size=1, dtype=torch.float32, ) def test_cache_size(self): """Cache should be sink_size + window_size * 2.""" expected_size = self.sink_size + self.window_size * 2 # 4 + 28 = 32 self.assertEqual(self.kv_cache.k_cache.size(2), expected_size) self.assertEqual(self.kv_cache.v_cache.size(2), expected_size) def test_is_ring_buffer(self): """Cache should be marked as ring buffer.""" self.assertTrue(self.kv_cache.is_ring_buffer) def test_update_stores_kv(self): """Update should store key-value pairs.""" k = torch.randn(1, self.n_heads, 5, self.head_dim) v = torch.randn(1, self.n_heads, 5, self.head_dim) input_pos = torch.tensor([0], dtype=torch.long) k_out, v_out = self.kv_cache.update(input_pos, k, v) # First 5 positions should contain our values torch.testing.assert_close(k_out[:, :, :5, :], k) torch.testing.assert_close(v_out[:, :, :5, :], v) def test_evict_tokens_returns_zero(self): """Ring buffer implementation doesn't shift, so evict returns 0.""" input_pos = torch.tensor([100], dtype=torch.long) shift = self.kv_cache.evict_tokens(input_pos, 10) self.assertEqual(shift, 0) def test_extended_generation(self): """Test that cache works for positions beyond cache size.""" cache_size = self.kv_cache.k_cache.size(2) # Fill cache with initial tokens for pos in range(cache_size + 50): k = torch.randn(1, self.n_heads, 1, self.head_dim) v = torch.randn(1, self.n_heads, 1, self.head_dim) input_pos = torch.tensor([pos], dtype=torch.long) k_out, v_out = self.kv_cache.update(input_pos, k, v) # Should not raise any errors self.assertEqual(k_out.shape, self.kv_cache.k_cache.shape) self.assertEqual(v_out.shape, self.kv_cache.v_cache.shape)class RopeWithAttentionSinkTest(unittest.TestCase): """Test RoPE for attention sink.""" def setUp(self): torch.manual_seed(42) self.params = ModelArgs( use_kv_cache=True, enable_dynamic_shape=True, max_context_len=1024, n_heads=8, dim=512, ) self.rope = RopeWithAttentionSink( params=self.params, window_size=100, sink_size=4, eviction_batch_size=1, ) def test_get_freqs_uses_original_position(self): """RoPE frequencies should use the original position.""" input_pos = torch.tensor([50], dtype=torch.long) seq_len = 5 freqs_cos, freqs_sin = self.rope.get_freqs(input_pos, seq_len) # Should get frequencies for positions 50-54 expected_cos = self.rope.freqs_cos[50:55] expected_sin = self.rope.freqs_sin[50:55] torch.testing.assert_close(freqs_cos, expected_cos) torch.testing.assert_close(freqs_sin, expected_sin) def test_rerotate_k(self): """Test re-rotation of k from one position to another.""" batch_size = 1 seq_len = 8 n_heads = self.params.n_heads head_dim = self.params.dim // n_heads k = torch.randn(batch_size, seq_len, n_heads, head_dim) q = torch.randn(batch_size, seq_len, n_heads, head_dim) # Rotate k at position 100 original_pos = 100 freqs_cos, freqs_sin = self.rope.get_freqs( torch.tensor([original_pos], dtype=torch.long), seq_len ) _, rotated_k = self.rope.forward(q, k, freqs_cos, freqs_sin) # Re-rotate to position 50 new_pos = 50 rerotated_k = self.rope.rerotate_k(rotated_k, original_pos, new_pos) # This should be equivalent to directly rotating k at position 50 freqs_cos_new, freqs_sin_new = self.rope.get_freqs( torch.tensor([new_pos], dtype=torch.long), seq_len ) _, expected_k = self.rope.forward(q, k, freqs_cos_new, freqs_sin_new) torch.testing.assert_close(rerotated_k, expected_k, rtol=1e-4, atol=1e-4)class CausalMaskWithWraparoundTest(unittest.TestCase): """Test causal mask with ring buffer wraparound.""" def test_mask_after_wraparound(self): """Test mask after cache has wrapped around.""" cache_size = 16 sink_size = 4 window_size = 6 # cache_size = sink_size + window_size * 2 # Simulate cache after generating beyond cache_size: # The ring buffer wraps, so indices 0-15 contain positions that wrap # At position 50, with cache_size=16, the cache contains: # positions 50-15=35 to 49 at various indices cache_positions = torch.zeros(cache_size, dtype=torch.long) # Fill with positions that would exist after generating 50 tokens # idx = pos % cache_size, so: # pos 34-49 occupy indices 2-15 and 0-1 for pos in range(34, 50): idx = pos % cache_size cache_positions[idx] = pos start_pos = 49 # Query at position 49 seq_len = 1 mask = _create_causal_mask_for_attention_sink( cache_positions, window_size, sink_size, start_pos, seq_len ) # Positions within window (49-6+1=44 to 49) should be visible visible_count = 0 for i in range(cache_size): pos = cache_positions[i].item() if pos >= 44 and pos <= 49: # In window self.assertEqual(mask[0, i].item(), 0.0, f"Position {pos} at idx {i} should be visible (in window)") visible_count += 1 # Should have some visible tokens self.assertGreater(visible_count, 0, "Should have visible tokens in window") def test_mask_with_sink_size_zero(self): """Test pure sliding window (sink_size=0).""" cache_size = 16 sink_size = 0 window_size = 8 cache_positions = torch.arange(cache_size, dtype=torch.long) start_pos = 10 seq_len = 1 mask = _create_causal_mask_for_attention_sink( cache_positions, window_size, sink_size, start_pos, seq_len ) # Positions 3-10 should be visible (within window of 8) for pos in range(3, 11): self.assertEqual(mask[0, pos].item(), 0.0, f"Position {pos} should be visible") # Positions 0-2 should be masked (outside window) for pos in range(0, 3): self.assertEqual(mask[0, pos].item(), float('-inf'), f"Position {pos} should be masked (outside window)")class PrefillTest(unittest.TestCase): """Test prefill scenarios.""" def setUp(self): torch.manual_seed(42) self.window_size = 14 self.sink_size = 4 self.n_heads = 8 self.head_dim = 64 self.params = ModelArgs( use_kv_cache=True, enable_dynamic_shape=True, max_context_len=1024, n_heads=self.n_heads, n_kv_heads=self.n_heads, dim=self.n_heads * self.head_dim, ) self.rope = RopeWithAttentionSink( params=self.params, window_size=self.window_size, sink_size=self.sink_size, eviction_batch_size=1, ) self.kv_cache = KVCacheWithAttentionSink( n_heads=self.n_heads, head_dim=self.head_dim, enable_dynamic_shape=True, rope=self.rope, max_batch_size=1, window_size=self.window_size, sink_size=self.sink_size, eviction_batch_size=1, dtype=torch.float32, ) def test_prefill_entire_cache(self): """Test prefill that fills entire cache.""" cache_size = self.kv_cache.k_cache.size(2) k = torch.randn(1, self.n_heads, cache_size, self.head_dim) v = torch.randn(1, self.n_heads, cache_size, self.head_dim) input_pos = torch.tensor([0], dtype=torch.long) k_out, v_out = self.kv_cache.update(input_pos, k, v) # All positions should be filled torch.testing.assert_close(k_out, k) torch.testing.assert_close(v_out, v) def test_prefill_larger_than_cache_raises_error(self): """Test that prefill larger than cache size raises an assertion error.""" cache_size = self.kv_cache.k_cache.size(2) seq_len = cache_size + 10 k = torch.randn(1, self.n_heads, seq_len, self.head_dim) v = torch.randn(1, self.n_heads, seq_len, self.head_dim) input_pos = torch.tensor([0], dtype=torch.long) # This should raise an assertion error since seq_len > cache_size with self.assertRaises(AssertionError): self.kv_cache.update(input_pos, k, v) def test_prefill_followed_by_decode(self): """Test prefill followed by decode steps.""" cache_size = self.kv_cache.k_cache.size(2) # Prefill with 20 tokens k_prefill = torch.randn(1, self.n_heads, 20, self.head_dim) v_prefill = torch.randn(1, self.n_heads, 20, self.head_dim) input_pos = torch.tensor([0], dtype=torch.long) self.kv_cache.update(input_pos, k_prefill, v_prefill) # Decode 5 more tokens for i in range(5): k_decode = torch.randn(1, self.n_heads, 1, self.head_dim) v_decode = torch.randn(1, self.n_heads, 1, self.head_dim) input_pos = torch.tensor([20 + i], dtype=torch.long) k_out, v_out = self.kv_cache.update(input_pos, k_decode, v_decode) # Verify cache positions are updated expected_pos = 20 + i cache_idx = expected_pos % cache_size self.assertEqual( self.kv_cache.cache_positions_manager.cache_positions[cache_idx].item(), expected_pos )class EnableAttentionSinkTest(unittest.TestCase): """Test the enable_attention_sink transformation.""" def setUp(self): torch.manual_seed(42) self.params = ModelArgs( use_kv_cache=True, enable_dynamic_shape=True, max_context_len=1024, n_heads=8, n_kv_heads=8, dim=512, n_layers=2, vocab_size=100, ) def test_enable_attention_sink_transforms_model(self): """Test that enable_attention_sink properly transforms the model.""" from executorch.examples.models.llama.llama_transformer import construct_transformer from executorch.examples.models.llama.source_transformation.attention_sink import ( enable_attention_sink, ) # Create a simple transformer with torch.device("meta"): model = construct_transformer(self.params) model.to_empty(device="cpu") # Apply attention sink transformation model = enable_attention_sink( module=model, params=self.params, sink_size=4, window_size=100, eviction_batch_size=1, ) # Check that KV caches are replaced for layer in model.layers: kv_cache = layer.attention.kv_cache self.assertIsInstance(kv_cache, KVCacheWithAttentionSink) self.assertEqual(kv_cache.sink_size, 4) self.assertEqual(kv_cache.window_size, 100) self.assertTrue(kv_cache.is_ring_buffer) def test_enable_attention_sink_replaces_rope(self): """Test that RoPE is replaced with RopeWithAttentionSink.""" from executorch.examples.models.llama.llama_transformer import construct_transformer from executorch.examples.models.llama.source_transformation.attention_sink import ( enable_attention_sink, ) with torch.device("meta"): model = construct_transformer(self.params) model.to_empty(device="cpu") model = enable_attention_sink( module=model, params=self.params, sink_size=4, window_size=100, eviction_batch_size=1, ) # Check that rope is replaced for layer in model.layers: rope = layer.attention.rope self.assertIsInstance(rope, RopeWithAttentionSink)class IntegrationTest(unittest.TestCase): """Integration tests for end-to-end scenarios.""" def setUp(self): torch.manual_seed(42) def test_cache_positions_consistency(self): """Test that cache positions remain consistent during generation.""" cache_size = 32 sink_size = 4 window_size = 14 n_heads = 8 head_dim = 64 params = ModelArgs( use_kv_cache=True, enable_dynamic_shape=True, max_context_len=1024, n_heads=n_heads, n_kv_heads=n_heads, dim=n_heads * head_dim, ) rope = RopeWithAttentionSink( params=params, window_size=window_size, sink_size=sink_size, eviction_batch_size=1, ) kv_cache = KVCacheWithAttentionSink( n_heads=n_heads, head_dim=head_dim, enable_dynamic_shape=True, rope=rope, max_batch_size=1, window_size=window_size, sink_size=sink_size, eviction_batch_size=1, dtype=torch.float32, ) # Generate 100 tokens for pos in range(100): k = torch.randn(1, n_heads, 1, head_dim) v = torch.randn(1, n_heads, 1, head_dim) input_pos = torch.tensor([pos], dtype=torch.long) kv_cache.update(input_pos, k, v) # Create mask and verify it's valid mask = kv_cache.create_causal_mask_for_ring_buffer(pos, 1) # Mask should not be all -inf (would mean no tokens to attend to) non_inf_count = (mask != float('-inf')).sum().item() self.assertGreater(non_inf_count, 0, f"At pos {pos}, mask should have visible tokens") # For positions >= sink_size, sinks should always be visible if pos >= sink_size: for i in range(sink_size): cache_pos = kv_cache.cache_positions_manager.cache_positions[i].item() if cache_pos < sink_size: # This is actually a sink token self.assertEqual(mask[0, i].item(), 0.0, f"Sink at idx {i} should be visible at pos {pos}")if __name__ == '__main__': unittest.main()
	# Copyright (c) Meta Platforms, Inc. and affiliates.
	# All rights reserved.
	#
	# This source code is licensed under the BSD-style license found in the
	# LICENSE file in the root directory of this source tree.
	"""
	Tests for the ring-buffer based attention sink implementation.
	This tests the torch.export-compatible implementation that uses a ring buffer
	for the sliding window rather than explicit token eviction.
	Usage:
	# Run with pytest
	python -m pytest examples/models/llama/source_transformation/test_attention_sink_ring_buffer.py -v
	# Or run directly
	python examples/models/llama/source_transformation/test_attention_sink_ring_buffer.py
	"""
	import unittest
	import torch
	from executorch.examples.models.llama.model_args import ModelArgs
	from executorch.examples.models.llama.source_transformation.attention_sink import (
	CachePositionsManagerWithSink,
	KVCacheWithAttentionSink,
	RopeWithAttentionSink,
	_create_causal_mask_for_attention_sink,
	)
	class CachePositionsManagerWithSinkTest(unittest.TestCase):
	"""Test the cache positions manager for ring buffer indexing."""
	def setUp(self) -> None:
	# Total cache size (e.g., sink_size + window_size * 2)
	self.cache_size = 32
	self.manager = CachePositionsManagerWithSink(self.cache_size)
	def test_initial_positions_are_zero(self) -> None:
	"""Cache positions should start as zeros."""
	expected = torch.zeros(self.cache_size, dtype=torch.long)
	torch.testing.assert_close(self.manager.cache_positions, expected)
	def test_simple_update(self) -> None:
	"""Test simple sequential update without wraparound."""
	input_pos = torch.tensor([0], dtype=torch.long)
	seq_len = 5
	indices = self.manager.calculate_positions_and_update_indices(
	input_pos, seq_len
	)
	# Should return indices 0, 1, 2, 3, 4
	expected_indices = torch.tensor([0, 1, 2, 3, 4], dtype=torch.long)
	torch.testing.assert_close(indices, expected_indices)
	# Cache positions at those indices should be the original positions
	for i in range(5):
	self.assertEqual(self.manager.cache_positions[i].item(), i)
	def test_wraparound(self) -> None:
	"""Test ring buffer wraparound."""
	# Fill cache to position 30
	input_pos = torch.tensor([0], dtype=torch.long)
	self.manager.calculate_positions_and_update_indices(input_pos, 30)
	# Add 5 more tokens at position 30 - should wrap around
	input_pos = torch.tensor([30], dtype=torch.long)
	indices = self.manager.calculate_positions_and_update_indices(
	input_pos, 5
	)
	# Indices should wrap: 30 % 32 = 30, 31, 0, 1, 2
	expected_indices = torch.tensor([30, 31, 0, 1, 2], dtype=torch.long)
	torch.testing.assert_close(indices, expected_indices)
	def test_cache_positions_track_original_positions(self) -> None:
	"""Cache positions should track which original position is at each index."""
	# Fill with positions 0-31
	input_pos = torch.tensor([0], dtype=torch.long)
	self.manager.calculate_positions_and_update_indices(input_pos, 32)
	# Now add position 32 which wraps to index 0
	input_pos = torch.tensor([32], dtype=torch.long)
	self.manager.calculate_positions_and_update_indices(input_pos, 1)
	# Index 0 should now contain original position 32
	self.assertEqual(self.manager.cache_positions[0].item(), 32)
	class CausalMaskTest(unittest.TestCase):
	"""Test the causal mask generation for attention sink."""
	def test_mask_allows_sink_tokens(self) -> None:
	"""Sink tokens should always be visible (mask = 0)."""
	cache_size = 32
	sink_size = 4
	# cache_size = sink_size + window_size * 2
	window_size = 14
	# Create cache positions where positions 0-3 are sink tokens
	cache_positions = torch.arange(cache_size, dtype=torch.long)
	start_pos = 20 # Query at position 20
	seq_len = 1
	mask = _create_causal_mask_for_attention_sink(
	cache_positions, window_size, sink_size, start_pos, seq_len
	)
	# Sink tokens (indices 0-3, original positions 0-3) should have mask = 0
	for i in range(sink_size):
	self.assertEqual(
	mask[0, i].item(),
	0.0,
	f"Sink token at index {i} should be visible",
	)
	def test_mask_blocks_future_tokens(self) -> None:
	"""Future tokens should be masked (-inf)."""
	cache_size = 32
	sink_size = 4
	window_size = 14
	cache_positions = torch.arange(cache_size, dtype=torch.long)
	start_pos = 10
	seq_len = 1
	mask = _create_causal_mask_for_attention_sink(
	cache_positions, window_size, sink_size, start_pos, seq_len
	)
	# Future tokens (positions > 10) should have mask = -inf
	for i in range(11, cache_size):
	self.assertEqual(
	mask[0, i].item(),
	float("-inf"),
	f"Future token at position {i} should be masked",
	)
	def test_mask_respects_window(self) -> None:
	"""Tokens outside the window should be masked."""
	cache_size = 32
	sink_size = 4
	# Only allow 5 recent tokens
	window_size = 5
	cache_positions = torch.arange(cache_size, dtype=torch.long)
	start_pos = 20
	seq_len = 1
	mask = _create_causal_mask_for_attention_sink(
	cache_positions, window_size, sink_size, start_pos, seq_len
	)
	# Positions 16-20 should be visible (within window of 5)
	for pos in range(16, 21):
	self.assertEqual(
	mask[0, pos].item(),
	0.0,
	f"Position {pos} should be visible (in window)",
	)
	# Position 15 should be masked (outside window, not a sink)
	self.assertEqual(
	mask[0, 15].item(),
	float("-inf"),
	"Position 15 should be masked (outside window)",
	)
	class KVCacheWithAttentionSinkTest(unittest.TestCase):
	"""Test the KV cache with attention sink."""
	def setUp(self) -> None:
	torch.manual_seed(42)
	self.window_size = 14
	self.sink_size = 4
	self.n_heads = 8
	self.head_dim = 64
	self.max_batch_size = 1
	# Create model args with enough context for RoPE
	self.params = ModelArgs(
	use_kv_cache=True,
	enable_dynamic_shape=True,
	max_context_len=1024, # Large enough for RoPE
	n_heads=self.n_heads,
	n_kv_heads=self.n_heads,
	dim=self.n_heads * self.head_dim,
	)
	self.rope = RopeWithAttentionSink(
	params=self.params,
	window_size=self.window_size,
	sink_size=self.sink_size,
	eviction_batch_size=1,
	)
	self.kv_cache = KVCacheWithAttentionSink(
	n_heads=self.n_heads,
	head_dim=self.head_dim,
	enable_dynamic_shape=True,
	rope=self.rope,
	max_batch_size=self.max_batch_size,
	window_size=self.window_size,
	sink_size=self.sink_size,
	eviction_batch_size=1,
	dtype=torch.float32,
	)
	def test_cache_size(self) -> None:
	"""Cache should be sink_size + window_size * 2."""
	# 4 + 28 = 32
	expected_size = self.sink_size + self.window_size * 2
	self.assertEqual(self.kv_cache.k_cache.size(2), expected_size)
	self.assertEqual(self.kv_cache.v_cache.size(2), expected_size)
	def test_is_ring_buffer(self) -> None:
	"""Cache should be marked as ring buffer."""
	self.assertTrue(self.kv_cache.is_ring_buffer)
	def test_update_stores_kv(self) -> None:
	"""Update should store key-value pairs."""
	k = torch.randn(1, self.n_heads, 5, self.head_dim)
	v = torch.randn(1, self.n_heads, 5, self.head_dim)
	input_pos = torch.tensor([0], dtype=torch.long)
	k_out, v_out = self.kv_cache.update(input_pos, k, v)
	# First 5 positions should contain our values
	torch.testing.assert_close(k_out[:, :, :5, :], k)
	torch.testing.assert_close(v_out[:, :, :5, :], v)
	def test_evict_tokens_returns_zero(self) -> None:
	"""Ring buffer implementation doesn't shift, so evict returns 0."""
	input_pos = torch.tensor([100], dtype=torch.long)
	shift = self.kv_cache.evict_tokens(input_pos, 10)
	self.assertEqual(shift, 0)
	def test_extended_generation(self) -> None:
	"""Test that cache works for positions beyond cache size."""
	cache_size = self.kv_cache.k_cache.size(2)
	# Fill cache with initial tokens
	for pos in range(cache_size + 50):
	k = torch.randn(1, self.n_heads, 1, self.head_dim)
	v = torch.randn(1, self.n_heads, 1, self.head_dim)
	input_pos = torch.tensor([pos], dtype=torch.long)
	k_out, v_out = self.kv_cache.update(input_pos, k, v)
	# Should not raise any errors
	self.assertEqual(k_out.shape, self.kv_cache.k_cache.shape)
	self.assertEqual(v_out.shape, self.kv_cache.v_cache.shape)
	class RopeWithAttentionSinkTest(unittest.TestCase):
	"""Test RoPE for attention sink."""
	def setUp(self) -> None:
	torch.manual_seed(42)
	self.params = ModelArgs(
	use_kv_cache=True,
	enable_dynamic_shape=True,
	max_context_len=1024,
	n_heads=8,
	dim=512,
	)
	self.rope = RopeWithAttentionSink(
	params=self.params,
	window_size=100,
	sink_size=4,
	eviction_batch_size=1,
	)
	def test_get_freqs_uses_original_position(self) -> None:
	"""RoPE frequencies should use the original position."""
	input_pos = torch.tensor([50], dtype=torch.long)
	seq_len = 5
	freqs_cos, freqs_sin = self.rope.get_freqs(input_pos, seq_len)
	# Should get frequencies for positions 50-54
	expected_cos = self.rope.freqs_cos[50:55]
	expected_sin = self.rope.freqs_sin[50:55]
	torch.testing.assert_close(freqs_cos, expected_cos)
	torch.testing.assert_close(freqs_sin, expected_sin)
	def test_rerotate_k(self) -> None:
	"""Test re-rotation of k from one position to another."""
	batch_size = 1
	seq_len = 8
	n_heads = self.params.n_heads
	head_dim = self.params.dim // n_heads
	k = torch.randn(batch_size, seq_len, n_heads, head_dim)
	q = torch.randn(batch_size, seq_len, n_heads, head_dim)
	# Rotate k at position 100
	original_pos = 100
	freqs_cos, freqs_sin = self.rope.get_freqs(
	torch.tensor([original_pos], dtype=torch.long), seq_len
	)
	_, rotated_k = self.rope.forward(q, k, freqs_cos, freqs_sin)
	# Re-rotate to position 50
	new_pos = 50
	rerotated_k = self.rope.rerotate_k(rotated_k, original_pos, new_pos)
	# This should be equivalent to directly rotating k at position 50
	freqs_cos_new, freqs_sin_new = self.rope.get_freqs(
	torch.tensor([new_pos], dtype=torch.long), seq_len
	)
	_, expected_k = self.rope.forward(q, k, freqs_cos_new, freqs_sin_new)
	torch.testing.assert_close(
	rerotated_k,
	expected_k,
	rtol=1e-4,
	atol=1e-4,
	)
	class CausalMaskWithWraparoundTest(unittest.TestCase):
	"""Test causal mask with ring buffer wraparound."""
	def test_mask_after_wraparound(self) -> None:
	"""Test mask after cache has wrapped around."""
	cache_size = 16
	sink_size = 4
	# cache_size = sink_size + window_size * 2
	window_size = 6
	# Simulate cache after generating beyond cache_size:
	# The ring buffer wraps, so indices 0-15 contain positions that wrap
	# At position 50, with cache_size=16, the cache contains:
	# positions 50-15=35 to 49 at various indices
	cache_positions = torch.zeros(cache_size, dtype=torch.long)
	# Fill with positions that would exist after generating 50 tokens
	# idx = pos % cache_size, so:
	# pos 34-49 occupy indices 2-15 and 0-1
	for pos in range(34, 50):
	idx = pos % cache_size
	cache_positions[idx] = pos
	start_pos = 49 # Query at position 49
	seq_len = 1
	mask = _create_causal_mask_for_attention_sink(
	cache_positions, window_size, sink_size, start_pos, seq_len
	)
	# Positions within window (49-6+1=44 to 49) should be visible
	visible_count = 0
	for i in range(cache_size):
	pos = cache_positions[i].item()
	if 44 <= pos <= 49: # In window
	self.assertEqual(
	mask[0, i].item(),
	0.0,
	f"Position {pos} at idx {i} should be visible (in window)",
	)
	visible_count += 1
	# Should have some visible tokens
	self.assertGreater(
	visible_count, 0, "Should have visible tokens in window"
	)
	def test_mask_with_sink_size_zero(self) -> None:
	"""Test pure sliding window (sink_size=0)."""
	cache_size = 16
	sink_size = 0
	window_size = 8
	cache_positions = torch.arange(cache_size, dtype=torch.long)
	start_pos = 10
	seq_len = 1
	mask = _create_causal_mask_for_attention_sink(
	cache_positions, window_size, sink_size, start_pos, seq_len
	)
	# Positions 3-10 should be visible (within window of 8)
	for pos in range(3, 11):
	self.assertEqual(
	mask[0, pos].item(),
	0.0,
	f"Position {pos} should be visible",
	)
	# Positions 0-2 should be masked (outside window)
	for pos in range(0, 3):
	self.assertEqual(
	mask[0, pos].item(),
	float("-inf"),
	f"Position {pos} should be masked (outside window)",
	)
	class PrefillTest(unittest.TestCase):
	"""Test prefill scenarios."""
	def setUp(self) -> None:
	torch.manual_seed(42)
	self.window_size = 14
	self.sink_size = 4
	self.n_heads = 8
	self.head_dim = 64
	self.params = ModelArgs(
	use_kv_cache=True,
	enable_dynamic_shape=True,
	max_context_len=1024,
	n_heads=self.n_heads,
	n_kv_heads=self.n_heads,
	dim=self.n_heads * self.head_dim,
	)
	self.rope = RopeWithAttentionSink(
	params=self.params,
	window_size=self.window_size,
	sink_size=self.sink_size,
	eviction_batch_size=1,
	)
	self.kv_cache = KVCacheWithAttentionSink(
	n_heads=self.n_heads,
	head_dim=self.head_dim,
	enable_dynamic_shape=True,
	rope=self.rope,
	max_batch_size=1,
	window_size=self.window_size,
	sink_size=self.sink_size,
	eviction_batch_size=1,
	dtype=torch.float32,
	)
	def test_prefill_entire_cache(self) -> None:
	"""Test prefill that fills entire cache."""
	cache_size = self.kv_cache.k_cache.size(2)
	k = torch.randn(1, self.n_heads, cache_size, self.head_dim)
	v = torch.randn(1, self.n_heads, cache_size, self.head_dim)
	input_pos = torch.tensor([0], dtype=torch.long)
	k_out, v_out = self.kv_cache.update(input_pos, k, v)
	# All positions should be filled
	torch.testing.assert_close(k_out, k)
	torch.testing.assert_close(v_out, v)
	def test_prefill_larger_than_cache_raises_error(self) -> None:
	"""Test that prefill larger than cache size raises an assertion error."""
	cache_size = self.kv_cache.k_cache.size(2)
	seq_len = cache_size + 10
	k = torch.randn(1, self.n_heads, seq_len, self.head_dim)
	v = torch.randn(1, self.n_heads, seq_len, self.head_dim)
	input_pos = torch.tensor([0], dtype=torch.long)
	# This should raise an assertion error since seq_len > cache_size
	with self.assertRaises(AssertionError):
	self.kv_cache.update(input_pos, k, v)
	def test_prefill_followed_by_decode(self) -> None:
	"""Test prefill followed by decode steps."""
	cache_size = self.kv_cache.k_cache.size(2)
	# Prefill with 20 tokens
	k_prefill = torch.randn(1, self.n_heads, 20, self.head_dim)
	v_prefill = torch.randn(1, self.n_heads, 20, self.head_dim)
	input_pos = torch.tensor([0], dtype=torch.long)
	self.kv_cache.update(input_pos, k_prefill, v_prefill)
	# Decode 5 more tokens
	for i in range(5):
	k_decode = torch.randn(1, self.n_heads, 1, self.head_dim)
	v_decode = torch.randn(1, self.n_heads, 1, self.head_dim)
	input_pos = torch.tensor([20 + i], dtype=torch.long)
	k_out, v_out = self.kv_cache.update(input_pos, k_decode, v_decode)
	# Verify cache positions are updated
	expected_pos = 20 + i
	cache_idx = expected_pos % cache_size
	self.assertEqual(
	self.kv_cache.cache_positions_manager.cache_positions[
	cache_idx
	].item(),
	expected_pos,
	)
	class EnableAttentionSinkTest(unittest.TestCase):
	"""Test the enable_attention_sink transformation."""
	def setUp(self) -> None:
	torch.manual_seed(42)
	self.params = ModelArgs(
	use_kv_cache=True,
	enable_dynamic_shape=True,
	max_context_len=1024,
	n_heads=8,
	n_kv_heads=8,
	dim=512,
	n_layers=2,
	vocab_size=100,
	)
	def test_enable_attention_sink_transforms_model(self) -> None:
	"""Test that enable_attention_sink properly transforms the model."""
	from executorch.examples.models.llama.llama_transformer import (
	construct_transformer,
	)
	from executorch.examples.models.llama.source_transformation.attention_sink import (
	enable_attention_sink,
	)
	# Create a simple transformer
	with torch.device("meta"):
	model = construct_transformer(self.params)
	model.to_empty(device="cpu")
	# Apply attention sink transformation
	model = enable_attention_sink(
	module=model,
	params=self.params,
	sink_size=4,
	window_size=100,
	eviction_batch_size=1,
	)
	# Check that KV caches are replaced
	for layer in model.layers:
	kv_cache = layer.attention.kv_cache
	self.assertIsInstance(kv_cache, KVCacheWithAttentionSink)
	self.assertEqual(kv_cache.sink_size, 4)
	self.assertEqual(kv_cache.window_size, 100)
	self.assertTrue(kv_cache.is_ring_buffer)
	def test_enable_attention_sink_replaces_rope(self) -> None:
	"""Test that RoPE is replaced with RopeWithAttentionSink."""
	from executorch.examples.models.llama.llama_transformer import (
	construct_transformer,
	)
	from executorch.examples.models.llama.source_transformation.attention_sink import (
	enable_attention_sink,
	)
	with torch.device("meta"):
	model = construct_transformer(self.params)
	model.to_empty(device="cpu")
	model = enable_attention_sink(
	module=model,
	params=self.params,
	sink_size=4,
	window_size=100,
	eviction_batch_size=1,
	)
	# Check that rope is replaced
	for layer in model.layers:
	rope = layer.attention.rope
	self.assertIsInstance(rope, RopeWithAttentionSink)
	class IntegrationTest(unittest.TestCase):
	"""Integration tests for end-to-end scenarios."""
	def setUp(self) -> None:
	torch.manual_seed(42)
	def test_cache_positions_consistency(self) -> None:
	"""Test that cache positions remain consistent during generation."""
	cache_size = 32
	sink_size = 4
	window_size = 14
	n_heads = 8
	head_dim = 64
	params = ModelArgs(
	use_kv_cache=True,
	enable_dynamic_shape=True,
	max_context_len=1024,
	n_heads=n_heads,
	n_kv_heads=n_heads,
	dim=n_heads * head_dim,
	)
	rope = RopeWithAttentionSink(
	params=params,
	window_size=window_size,
	sink_size=sink_size,
	eviction_batch_size=1,
	)
	kv_cache = KVCacheWithAttentionSink(
	n_heads=n_heads,
	head_dim=head_dim,
	enable_dynamic_shape=True,
	rope=rope,
	max_batch_size=1,
	window_size=window_size,
	sink_size=sink_size,
	eviction_batch_size=1,
	dtype=torch.float32,
	)
	# Generate 100 tokens
	for pos in range(100):
	k = torch.randn(1, n_heads, 1, head_dim)
	v = torch.randn(1, n_heads, 1, head_dim)
	input_pos = torch.tensor([pos], dtype=torch.long)
	kv_cache.update(input_pos, k, v)
	# Create mask and verify it's valid
	mask = kv_cache.create_causal_mask_for_ring_buffer(pos, 1)
	# Mask should not be all -inf (would mean no tokens to attend to)
	non_inf_count = (mask != float("-inf")).sum().item()
	self.assertGreater(
	non_inf_count,
	0,
	f"At pos {pos}, mask should have visible tokens",
	)
	# For positions >= sink_size, sinks should always be visible
	if pos >= sink_size:
	for i in range(sink_size):
	cache_pos = (
	kv_cache.cache_positions_manager.cache_positions[i].item()
	)
	# This is actually a sink token
	if cache_pos < sink_size:
	self.assertEqual(
	mask[0, i].item(),
	0.0,
	f"Sink at idx {i} should be visible at pos {pos}",
	)
	if __name__ == "__main__":
	unittest.main()

[Copilot]

The entire file content has been compressed onto a single line, removing all newlines. This makes the code completely unreadable and unmaintainable. The file should be properly formatted with appropriate line breaks, indentation, and spacing according to Python coding standards (PEP 8).

[Copilot]

The entire file content has been compressed onto a single line, removing all newlines. This makes the code completely unreadable and unmaintainable. The file should be properly formatted with appropriate line breaks, indentation, and spacing according to Python coding standards (PEP 8).

[Copilot]

The hardcoded value of 131072 (128k) seems arbitrary and may not be appropriate for all model configurations. This could cause excessive memory usage for RoPE frequency tables in models that don't need such large context. Consider making this configurable through the llm_config or deriving it from model-specific parameters rather than using a hardcoded default.

Suggested change

	# Use 131072 (128k) as default for Llama 3.2 models or the original model max if larger.
	default_rope_length = max(131072, model_args.max_context_len)
	# Derive the RoPE max context length from configuration rather than using a fixed
	# default. Use the larger of export.max_context_length and the original
	# model_args.max_context_len.
	default_rope_length = max(
	self.llm_config.export.max_context_length,
	model_args.max_context_len,
	)

[Copilot]

Typo in the assertion message: "secpific" should be "specific".

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 9 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The TEMPORARY override logic (lines 159-167) should not be merged to production. This code allows "effectively unlimited" generation (max_new_tokens = 1000000) which could cause memory issues or infinite loops. If this is needed for testing, it should be guarded by a flag or removed before merging to production branches.

Suggested change

	// TEMPORARY: For sliding window / infinite context testing,
	// override max_new_tokens to allow unlimited generation
	if (max_new_tokens <= 0) {
	max_new_tokens = 1000000; // Effectively unlimited
	}
	// If user specified seq_len, use that instead
	if (config.seq_len > 0 && config.seq_len > max_new_tokens) {
	max_new_tokens = config.seq_len;
	}
	// If user specified seq_len, use that instead.
	if (config.seq_len > 0) {
	max_new_tokens = static_cast<int>(config.seq_len);
	}

[Copilot]

Chinese characters are used in comments. Comments should be in English for consistency with the rest of the codebase. Please translate:

• Line 11: "# window_size=124: 滑动窗口大小" should be "# window_size=124: Sliding window size"
• Line 12: "# eviction_batch_size=1: 每次驱逐 1 个 token" should be "# eviction_batch_size=1: Evict 1 token at a time"

Suggested change

	# window_size=124: 滑动窗口大小
	# eviction_batch_size=1: 每次驱逐 1 个 token
	# window_size=124: Sliding window size
	# eviction_batch_size=1: Evict 1 token at a time

[Copilot]

The TEMPORARY override logic has problematic behavior. The condition on line 165 checks if config.seq_len > max_new_tokens, but if max_new_tokens was just set to 1000000 on line 162, this condition will rarely be true unless seq_len is also extremely large. This means the seq_len override may not work as intended when max_new_tokens was initially <= 0.

Consider revising the logic to check config.seq_len before applying the "unlimited" override, or reconsider the order of these checks. For example:

1. First apply seq_len override if config.seq_len > 0
2. Then apply the unlimited fallback only if the result is still <= 0

Suggested change

	// TEMPORARY: For sliding window / infinite context testing,
	// override max_new_tokens to allow unlimited generation
	if (max_new_tokens <= 0) {
	max_new_tokens = 1000000; // Effectively unlimited
	}
	// If user specified seq_len, use that instead
	if (config.seq_len > 0 && config.seq_len > max_new_tokens) {
	max_new_tokens = config.seq_len;
	}
	// If user specified seq_len, use that instead
	if (config.seq_len > 0 && config.seq_len > max_new_tokens) {
	max_new_tokens = config.seq_len;
	}
	// TEMPORARY: For sliding window / infinite context testing,
	// override max_new_tokens to allow unlimited generation
	if (max_new_tokens <= 0) {
	max_new_tokens = 1000000; // Effectively unlimited
	}

[Copilot]

The function attention_sink_forward is defined but never used. The comment on line 403-404 states "Don't replace forward - let the original AttentionMHA.forward handle it", but the docstring on line 419 still says "Replace Attention's forward with attention_sink_forward". This is inconsistent and creates confusion.

Either remove the unused attention_sink_forward function (lines 307-349) if it's truly not needed, or use it if it's meant to replace the forward method. If keeping it for future use, add a comment explaining why it's currently unused.

[Copilot]

The docstring is outdated and inconsistent with the actual implementation. Line 419 states "Replace Attention's forward with attention_sink_forward", but the code on lines 403-404 explicitly does NOT replace the forward method. Update the docstring to accurately reflect what the function does.

Suggested change

	- Replace Attention's forward with attention_sink_forward
	- Configure Attention/SDPA to use attention masks compatible with the ring buffer and attention sink

[Copilot]

Two imports from executorch.examples.models.llama.attention are unused: CachePositionsManager and RingKVCache (lines 21 and 23). These should be removed to keep the imports clean.

Suggested change

	CachePositionsManager,
	KVCache,
	RingKVCache,
	KVCache,

[Copilot]

Variable k_out is not used.

Suggested change

	k_out, v_out = self.kv_cache.update(input_pos, k_decode, v_decode)
	self.kv_cache.update(input_pos, k_decode, v_decode)

[Copilot]

Variable v_out is not used.

Suggested change

	# Verify that the returned k/v match the decode inputs
	self.assertEqual(k_out.shape, k_decode.shape)
	self.assertEqual(v_out.shape, v_decode.shape)
	torch.testing.assert_close(k_out, k_decode)
	torch.testing.assert_close(v_out, v_decode)

[Copilot]

Variable cache_size is not used.

Suggested change

cache_size = 32

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 11 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The docstring states "There mainly three steps" with a grammatical error. It should be "There are mainly three steps". Additionally, step 3 mentions "Replace Attention's forward with attention_sink_forward", but the implementation at lines 389-390 explicitly does NOT replace the forward method. The docstring should be updated to accurately reflect what the function does.

[Copilot]

CachePositionsManagerWithSink initializes cache_positions to all zeros, which differs from the original CachePositionsManager that uses sentinel values (-1) for invalid/unfilled positions (see attention.py lines 219-225). This means before the cache is fully populated, positions that haven't been written will have a value of 0, which could be confused with actual position 0. The mask creation function checks cache_positions >= 0 to determine validity, so unfilled positions with value 0 will be incorrectly treated as valid position 0. This should use -1 for unfilled positions, or implement the same sentinel logic as the base CachePositionsManager.

Suggested change

	# Use zeros like original CachePositionsManager
	self.register_buffer(
	"cache_positions",
	torch.zeros((self.max_context_length,), dtype=torch.long, device="cpu"),
	# Use -1 as sentinel for unfilled positions, matching CachePositionsManager
	self.register_buffer(
	"cache_positions",
	torch.full(
	(self.max_context_length,),
	-1,
	dtype=torch.long,
	device="cpu",
	),

[Copilot]

The calculate_positions_and_update_indices method doesn't implement the sentinel value logic that the original CachePositionsManager has (see attention.py lines 219-225). Before the cache wraps around, positions that haven't been written yet should be marked as invalid with -1, not left as 0. The current implementation doesn't clear previously unfilled positions to -1 when updating, which can lead to stale or incorrect values in cache_positions being treated as valid. This should implement the same sentinel logic as the base class where positions that haven't been written yet are set to -1.

[Copilot]

Type mismatch in create_causal_mask_for_ring_buffer: the parameter start_pos is declared as torch.Tensor (line 239), but it's passed directly to _create_causal_mask_for_attention_sink (line 249) which expects an int according to its signature (line 103). This will cause a type error. The start_pos should either be extracted to an int with .item() before passing, or the function signature should be updated to accept torch.Tensor consistently.

[Copilot]

The function attention_sink_forward is defined but appears to be unused based on the comment at lines 389-390 which states "Don't replace forward". This creates dead code in the codebase. Either this function should be used to replace the forward method, or it should be removed if it's no longer needed. The docstring at line 405 also still mentions "Replace Attention's forward with attention_sink_forward", which is inconsistent with the actual implementation.

[Copilot]

The update method in KVCacheWithAttentionSink always uses index_copy_ without checking the enable_dynamic_shape flag that was passed during initialization. The base KVCache class (in attention.py lines 88-118) has different behavior depending on this flag. For consistency with the parent class and to support both dynamic and static shape modes, this method should check self.enable_dynamic_shape and implement both code paths like the base class does.

Suggested change

	start_pos = input_pos[0].item()
	torch._check_is_size(start_pos)
	self.k_cache.index_copy_(2, indices, k_val)
	self.v_cache.index_copy_(2, indices, v_val)
	if self.enable_dynamic_shape:
	# Dynamic-shape-friendly bulk update using index_copy_
	start_pos = input_pos[0].item()
	torch._check_is_size(start_pos)
	self.k_cache.index_copy_(2, indices, k_val)
	self.v_cache.index_copy_(2, indices, v_val)
	else:
	# Static-shape-friendly path: per-position assignment using the same indices.
	# This mirrors the base KVCache behavior of avoiding dynamic index_copy_
	# in static-shape mode while preserving ring buffer semantics.
	for i in range(seq_len):
	cache_idx = indices[i].item()
	torch._check_is_size(cache_idx)
	# k_cache/v_cache shape is [batch, heads, context, head_dim]
	# k_val/v_val shape is [batch, heads, seq_len, head_dim]
	self.k_cache[:, :, cache_idx, :] = k_val[:, :, i, :]
	self.v_cache[:, :, cache_idx, :] = v_val[:, :, i, :]

[Copilot]

Variable k_out is not used.

Suggested change

	k_out, v_out = self.kv_cache.update(input_pos, k_decode, v_decode)
	self.kv_cache.update(input_pos, k_decode, v_decode)

[Copilot]

Variable v_out is not used.

Suggested change

	k_out, v_out = self.kv_cache.update(input_pos, k_decode, v_decode)
	self.kv_cache.update(input_pos, k_decode, v_decode)

[Copilot]

Variable cache_size is not used.

Suggested change

cache_size = 32

[Copilot]

Import of 'CachePositionsManager' is not used.
Import of 'RingKVCache' is not used.

Suggested change

	CachePositionsManager,
	KVCache,
	RingKVCache,
	KVCache,

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 7 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Missing opening backtick for the code block. The code block should start with triple backticks before the shell command.

[Copilot]

The TEMPORARY max_new_tokens override with a hardcoded value of 1000000 is concerning for production code. This effectively removes any generation length limits and could cause memory issues or runaway generation. Consider either removing this temporary code if testing is complete, or adding proper configuration for unlimited generation with appropriate safeguards.

Suggested change

	// TEMPORARY: For sliding window / infinite context testing,
	// override max_new_tokens to allow unlimited generation
	if (max_new_tokens <= 0) {
	max_new_tokens = 1000000; // Effectively unlimited
	}

[Copilot]

The BUCK file still references "source_transformation/test_attention_sink.py" at line 276, but this file is being deleted in this PR. This will cause build failures. The test target "attention_sink_test" should either be removed or updated to reference the correct test file.

[Copilot]

The comment mentions that forward should be replaced with attention_sink_forward (line 410), but the actual code at line 394 explicitly states "Don't replace forward". This is a discrepancy between the docstring and the implementation. The docstring should be updated to reflect that the forward method is NOT being replaced.

Suggested change

	- Replace Attention's forward with attention_sink_forward
	- Keep Attention's original forward; KVCacheWithAttentionSink and masks provide Attention Sink behavior.

[Copilot]

The function attention_sink_forward is defined but never used. According to the comment at line 394, the forward method is not being replaced, and instead the original AttentionMHA.forward handles the ring buffer case. This unused function adds confusion and maintenance burden. Consider either removing this function or adding a comment explaining why it's kept (e.g., for potential future use or documentation purposes).

[Copilot]

Typo in comment: "There mainly three steps" should be "There are mainly three steps".

Suggested change

	There mainly three steps:
	There are mainly three steps:

[Copilot]

Missing closing backtick for the code block. The code block starting at line 23 is not properly closed before this line, which will cause the entire section to be rendered incorrectly.