[pssrawat]

Summary:
Enabling the attention sink on a model that also quantizes its KV cache
silently produced a model with no attention sink, and the model then failed
once a conversation ran past the cache.

The export applies replace_kv_cache_with_custom_kv_cache and then
replace_kv_cache_with_quantized_kv_cache. The second one dispatches on
isinstance(child, CustomKVCache), and CustomKVCacheWithAttentionSink and
CustomRingKVCache subclass CustomKVCache, so both were rebuilt as a plain
full-context QuantizedKVCache with no ring, no pinned prefix and no mask,
and replace_sdpa_with_quantized_sdpa built QuantizedSDPA without
use_attention_mask. Once the position passes the cache's slot count,
update_cache rejects the write, so the session fails instead of evicting.

Adds QuantizedKVCacheWithAttentionSink, the quantized counterpart of
CustomKVCacheWithAttentionSink. It takes its slot count from the source
cache, so the layout carries over exactly; the write slots come from
CachePositionsManagerWithSink and the mask from
_create_causal_mask_for_attention_sink, as on the float path. Entries are
int8, quantized per token and head (affine asymmetric) with float32 scales and
int8 zero points, and are written through update_cache_with_indices; the
constructor defaults to that and rejects the configurations it cannot run.

CustomRingKVCache and RingKVCache (attention sink with sink_size=0, and
native sliding-window layers) now become the existing QuantizedRingKVCache
instead of being flattened. Any other ring cache with no quantized
counterpart, such as a sink cache when the custom KV cache transform was
skipped, raises instead of being flattened. A cache that is already quantized
is left alone, so running the transform twice is a no-op.

QuantizedSDPA now uses the mask whenever its cache is a ring or sink cache,
whatever the float SDPA was built with; those caches build an additive float
mask. A full-context cache keeps plain causal attention, so a model built with
use_custom_sdpa_with_attention_mask and no ring is unchanged.

Authored with Claude.

Differential Revision: D126929701

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit d8d4226 with merge base 77f5323 ():

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

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

• pull / nxp-mcuxpresso-test / linux-job (gh) (trunk failure)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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

[meta-codesync]

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

[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.