Integrate attention sink into ET LLM export and runner (#18860)

[kirklandsign]

Summary:

Add custom op support, export pipeline integration, and C++ runner fixes
for the attention sink ring buffer implementation.

• CustomKVCacheWithAttentionSink: custom op variant using update_cache_with_indices
  for scatter-write performance. Replaces KVCacheWithAttentionSink during export.
• CustomRingKVCache replacement: handle RingKVCache -> CustomRingKVCache in the
  replacement pass, and set SDPACustom.use_attention_mask=True for ring buffer models.
• Export transform ordering: replace SDPA before KV cache so that
  _replace_kv_cache_with_custom_kv_cache can set use_attention_mask=True on the
  already-existing SDPACustom (previously the ordering was reversed, causing the
  mask flag to be overwritten by a new SDPACustom).
• C++ runner: add max_seq_len prefill check; make context length check conditional
  for sliding window models (max_seq_len < max_context_len) since they handle
  position wrapping internally via ring buffer.

Reviewed By: lucylq

Differential Revision: D100216686

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 3 Pending

As of commit df7ede5 with merge base 3180927 ():

NEW FAILURE - The following job has failed:

• Cadence Build & Test / cpu-test / test-aot / test-aot (gh)
  backends/cadence/aot/tests/test_replace_ops_passes.py::TestReplaceOpsPasses::test_replace_transposed_conv_with_linear_0

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

[meta-codesync]

@kirklandsign has exported this pull request. If you are a Meta employee, you can view the originating Diff in D100216686.

[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

Integrates the “attention sink” ring-buffer KV-cache path into the Llama export pipeline and runtime runner, including custom-op KV cache replacements and updated runner-side constraints for sliding-window models.

Changes:

• Add CustomKVCacheWithAttentionSink and extend KV-cache replacement to handle KVCacheWithAttentionSink and RingKVCache, including enabling attention-mask mode on SDPACustom.
• Reorder export transforms so SDPA replacement happens before KV-cache replacement (to preserve use_attention_mask=True for ring-buffer models).
• Update C++ text runner to (a) enforce a single-prefill-chunk max_seq_len limit and (b) relax/adjust context-limit logic for sliding-window models; add new end-to-end tests for custom SDPA + custom KV cache.

Reviewed changes

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

File	Description
extension/llm/runner/text_llm_runner.cpp	Adds max-seq-len prefill guard and modifies max-new-tokens budgeting for sliding-window models.
examples/models/llama/source_transformation/test_attention_sink.py	Adds end-to-end generation tests for custom SDPA + custom KV cache and sink_size=0 edge case.
examples/models/llama/source_transformation/custom_kv_cache.py	Introduces CustomKVCacheWithAttentionSink and expands KV-cache replacement logic for attention-sink/ring KV caches.
examples/models/llama/export_llama_lib.py	Reorders source transforms so SDPA replacement precedes KV-cache replacement to preserve mask settings.

---

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

[Copilot]

effective_pos is set to 0 for sliding-window models, but resolve_max_new_tokens(max_context_len, effective_pos) is intended to ensure start_pos + max_new_tokens <= max_context_len. With effective_pos=0, a nonzero pos_ can lead to generating past max_context_len (e.g., RoPE tables typically guard input_pos + seq_len <= max_context_len), causing runtime failures. Consider keeping pos_ (or a wrapped position if the model truly wraps positions) when resolving max_new_tokens, and only relax KV-cache-capacity checks separately from the RoPE/context-length limit.

Suggested change

	// For sliding window models, the ring buffer recycles space — pos_ doesn't
	// represent consumed capacity, so pass 0 to get the full budget.
	int64_t effective_pos = (max_seq_len < max_context_len) ? 0 : pos_;
	int max_new_tokens = config.resolve_max_new_tokens(max_context_len, effective_pos);
	// Resolve generation budget from the actual starting position so
	// start_pos + max_new_tokens stays within max_context_len. Sliding-window
	// KV-cache reuse must not bypass positional/context-length limits.
	int64_t effective_pos = pos_;
	int max_new_tokens =
	config.resolve_max_new_tokens(max_context_len, effective_pos);

[Copilot]

CustomKVCacheWithAttentionSink.update() computes ring-buffer write indices but does not validate that the number of window tokens written in a single update fits within the ring portion (i.e., avoids duplicate indices due to modulo wrap). KVCacheWithAttentionSink.update() has an explicit num_window_tokens <= ring_size guard to prevent non-deterministic index_copy_ behavior; the same kind of guard is needed here as well (especially since update_cache_with_indices is also a scatter-style update).

[kirklandsign]

@claude review the epatch

[claude]

Claude finished @kirklandsign's task in 1m 51s —— View job

---

Review of PR #18860: Integrate attention sink into ET LLM export and runner

• Read the full diff and changed files
• Review custom_kv_cache.py changes
• Review export_llama_lib.py changes
• Review text_llm_runner.cpp changes
• Review test_attention_sink.py changes
• Post review feedback

---

Overall this is a well-structured PR that correctly integrates the attention sink ring buffer into the export pipeline and C++ runner. The transform ordering fix is sound, the new CustomKVCacheWithAttentionSink class mirrors existing patterns, and the tests cover the key scenarios. A few items worth addressing:

---

Issue 1: Missing num_window_tokens <= ring_size guard in CustomKVCacheWithAttentionSink.update()

KVCacheWithAttentionSink.update() (attention_sink.py:240-246) has an explicit guard that the number of window tokens doesn't exceed the ring buffer capacity — this prevents duplicate indices in the scatter write, which would cause non-deterministic behavior:

# attention_sink.py:239-245
start_pos = input_pos[0].item()
num_sink_tokens = max(0, min(seq_len, self.sink_size - start_pos))
num_window_tokens = seq_len - num_sink_tokens
assert num_window_tokens <= self.cache_positions_manager.ring_size, (
    ...
)

CustomKVCacheWithAttentionSink.update() (custom_kv_cache.py:551-561) uses update_cache_with_indices (also a scatter-style operation) but lacks this guard. While the existing seq_len <= cache_size check at line 553 prevents writing more tokens than total cache slots, it doesn't prevent the ring portion from wrapping with duplicate indices when num_window_tokens > ring_size.

Consider adding the same guard before computing indices.

Fix this →

---

Issue 2: effective_pos = 0 may over-allocate generation budget

In text_llm_runner.cpp:188:

int64_t effective_pos = (max_seq_len < max_context_len) ? 0 : pos_;
int max_new_tokens = config.resolve_max_new_tokens(max_context_len, effective_pos);

For sliding window models, passing effective_pos = 0 means resolve_max_new_tokens returns max_context_len (or close to it) as the generation budget. This is correct in the sense that the ring buffer can recycle space indefinitely, so the KV cache won't overflow. However, this effectively means generation can proceed until pos_ exceeds max_context_len, which is the RoPE frequency table size. If the model's RoPE table was exported with a specific max_context_len bound, generating past it would cause an out-of-bounds access on the frequency table.

In practice this may be fine if the caller sets max_new_tokens explicitly, but without that, the runner silently allows generating up to max_context_len tokens regardless of current position. Consider whether this should be documented more explicitly or whether the RoPE table size should be the real cap (it already is, but the connection is implicit). This is a minor concern since the RoPE module itself will guard against exceeding its table size.

---

Issue 3: _replace_kv_cache_with_custom_kv_cache doesn't recurse for matched cases

In custom_kv_cache.py:379-424, when a child matches KVCacheWithAttentionSink, RingKVCache, or KVCache, the function replaces it but does not recurse into the replacement or the child's own children. The else branch at line 423 is only reached if none of the isinstance checks match.

This is fine for the current model structure where KV caches are direct children of attention modules, but it means nested structures won't be traversed if a match is found at a parent level. The existing code before this PR had the same pattern for KVCache, so this is consistent — just worth noting that it assumes flat module hierarchies for cache modules.

---

Issue 4: Storing a module-level function reference as an instance attribute

custom_kv_cache.py:536-538:

self._create_causal_mask_for_attention_sink = (
    _create_causal_mask_for_attention_sink
)

This stores a module-level function as an instance attribute. While functional, it's unusual and could cause issues with torch.export / torch.jit tracing since it's not a submodule or buffer. CustomRingKVCache (line 594) avoids this by calling _create_causal_mask_for_ring_buffer directly from module scope. Consider doing the same for _create_causal_mask_for_attention_sink — importing at the top of the class method rather than storing as an attribute.

---

Minor observations

• Transform ordering fix (export_llama_lib.py): The reordering from KV-cache-then-SDPA to SDPA-then-KV-cache is correct. The comment at lines 1784-1787 explains the rationale clearly.

• Test coverage (test_attention_sink.py:401-444): The two new tests (test_beyond_context_window_custom_sdpa and test_sink_zero_custom_sdpa) cover the main paths well. The test at line 348 correctly calls replace_sdpa_with_custom_op before replace_kv_cache_with_custom_kv_cache, matching the new transform order.

• Error message improvement (text_llm_runner.cpp:143-148): The new prefill check against max_seq_len gives a better error message than the old combined check. The separation of prefill-chunk-size vs. context-capacity checks is an improvement.

---

[kirklandsign]

Claude finished @kirklandsign's task in 1m 51s —— View job

Review of PR #18860: Integrate attention sink into ET LLM export and runner

• Read the full diff and changed files
• Review custom_kv_cache.py changes
• Review export_llama_lib.py changes
• Review text_llm_runner.cpp changes
• Review test_attention_sink.py changes
• Post review feedback

Overall this is a well-structured PR that correctly integrates the attention sink ring buffer into the export pipeline and C++ runner. The transform ordering fix is sound, the new CustomKVCacheWithAttentionSink class mirrors existing patterns, and the tests cover the key scenarios. A few items worth addressing:

Issue 1: Missing num_window_tokens <= ring_size guard in CustomKVCacheWithAttentionSink.update()

KVCacheWithAttentionSink.update() (attention_sink.py:240-246) has an explicit guard that the number of window tokens doesn't exceed the ring buffer capacity — this prevents duplicate indices in the scatter write, which would cause non-deterministic behavior:

# attention_sink.py:239-245
start_pos = input_pos[0].item()
num_sink_tokens = max(0, min(seq_len, self.sink_size - start_pos))
num_window_tokens = seq_len - num_sink_tokens
assert num_window_tokens <= self.cache_positions_manager.ring_size, (
    ...
)

CustomKVCacheWithAttentionSink.update() (custom_kv_cache.py:551-561) uses update_cache_with_indices (also a scatter-style operation) but lacks this guard. While the existing seq_len <= cache_size check at line 553 prevents writing more tokens than total cache slots, it doesn't prevent the ring portion from wrapping with duplicate indices when num_window_tokens > ring_size.

Consider adding the same guard before computing indices.

Fix this →

Issue 2: effective_pos = 0 may over-allocate generation budget

In text_llm_runner.cpp:188:

int64_t effective_pos = (max_seq_len < max_context_len) ? 0 : pos_;
int max_new_tokens = config.resolve_max_new_tokens(max_context_len, effective_pos);

For sliding window models, passing effective_pos = 0 means resolve_max_new_tokens returns max_context_len (or close to it) as the generation budget. This is correct in the sense that the ring buffer can recycle space indefinitely, so the KV cache won't overflow. However, this effectively means generation can proceed until pos_ exceeds max_context_len, which is the RoPE frequency table size. If the model's RoPE table was exported with a specific max_context_len bound, generating past it would cause an out-of-bounds access on the frequency table.

In practice this may be fine if the caller sets max_new_tokens explicitly, but without that, the runner silently allows generating up to max_context_len tokens regardless of current position. Consider whether this should be documented more explicitly or whether the RoPE table size should be the real cap (it already is, but the connection is implicit). This is a minor concern since the RoPE module itself will guard against exceeding its table size.

Issue 3: _replace_kv_cache_with_custom_kv_cache doesn't recurse for matched cases

In custom_kv_cache.py:379-424, when a child matches KVCacheWithAttentionSink, RingKVCache, or KVCache, the function replaces it but does not recurse into the replacement or the child's own children. The else branch at line 423 is only reached if none of the isinstance checks match.

This is fine for the current model structure where KV caches are direct children of attention modules, but it means nested structures won't be traversed if a match is found at a parent level. The existing code before this PR had the same pattern for KVCache, so this is consistent — just worth noting that it assumes flat module hierarchies for cache modules.

Issue 4: Storing a module-level function reference as an instance attribute

custom_kv_cache.py:536-538:

self._create_causal_mask_for_attention_sink = (
    _create_causal_mask_for_attention_sink
)

This stores a module-level function as an instance attribute. While functional, it's unusual and could cause issues with torch.export / torch.jit tracing since it's not a submodule or buffer. CustomRingKVCache (line 594) avoids this by calling _create_causal_mask_for_ring_buffer directly from module scope. Consider doing the same for _create_causal_mask_for_attention_sink — importing at the top of the class method rather than storing as an attribute.

Minor observations

• Transform ordering fix (export_llama_lib.py): The reordering from KV-cache-then-SDPA to SDPA-then-KV-cache is correct. The comment at lines 1784-1787 explains the rationale clearly.
• Test coverage (test_attention_sink.py:401-444): The two new tests (test_beyond_context_window_custom_sdpa and test_sink_zero_custom_sdpa) cover the main paths well. The test at line 348 correctly calls replace_sdpa_with_custom_op before replace_kv_cache_with_custom_kv_cache, matching the new transform order.
• Error message improvement (text_llm_runner.cpp:143-148): The new prefill check against max_seq_len gives a better error message than the old combined check. The separation of prefill-chunk-size vs. context-capacity checks is an improvement.

Issue 1 and 4 will be addressed in a follow-up patch. 2 and 3 is not a blocker

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

---

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

[Copilot]

The new num_prompt_tokens <= max_seq_len check rejects prompts longer than max_seq_len, but TextPrefiller::prefill() already chunks prompts larger than max_seq_len_ into multiple prefill_chunk() calls. This is a behavior regression (previously long prompts could succeed via chunking) and isn’t necessary to enforce the per-chunk limit. Consider removing this guard and relying on TextPrefiller chunking, or (if needed) enforce the max chunk size inside the prefiller rather than rejecting the whole prompt here.