Optimize CPU SDPA for ring attention (pytorch#22358)

## Summary

- skip fully masked KV blocks and trim masked boundaries in CPU custom
SDPA
- split wrapped ring windows around large masked gaps while preserving
additive-mask semantics
- size local KV caches as sliding window plus the maximum in-flight
prefill chunk, capped by full context
- support regular, custom-op, and quantized ring caches

Author: JacobSzwejbka

examples/models/llama/attention.py

@@ -262,6 +262,23 @@ def calculate_positions_and_update_indices(self, input_pos: torch.Tensor, seq_le
         return indices
 
 
+def _get_ring_cache_size(
+    max_context_length: int,
+    window_size: int,
+    max_seq_len: Optional[int] = None,
+) -> int:
+    """Size an SWA cache for one retained window plus one in-flight chunk."""
+    assert window_size > 0, "Sliding-window size must be positive"
+    if max_seq_len is None:
+        max_seq_len = max_context_length
+    assert max_seq_len > 0, "Maximum sequence length must be positive"
+    assert window_size <= max_context_length, (
+        f"Sliding-window size ({window_size}) cannot exceed the full context "
+        f"length ({max_context_length})"
+    )
+    return min(max_context_length, window_size + max_seq_len)
+
+
 class RingKVCache(KVCache):
     def __init__(
         self,
@@ -271,10 +288,17 @@ def __init__(
         head_dim: int,
         enable_dynamic_shape: bool,
         dtype=torch.float32,
+        *,
+        window_size: int,
+        max_seq_len: Optional[int] = None,
     ):
-        self.window_size = max_context_length
         """
-        Reason why we want the kv cache size to be twice the context length:
+        The cache needs room for the retained sliding window and the current
+        prefill chunk. Its size is window_size + max_seq_len, capped by the
+        full-context cache. If max_seq_len is omitted, the full context length
+        is used.
+
+        Reason why a cache larger than the sliding window is needed:
         Sliding window attention without ringbuffer
         pos   0  1  2  3  4  5  6  7  8  9  10
         0     x  0  0  0  0  0  0  0  0  0  0
@@ -307,20 +331,28 @@ def __init__(
         So not having kept 2, 3 and 4 in cache means we will have divergent behavior.
         Worst case of this would have been when update it equal to the length of
         the cache. like in our case pos = 5 seq len = 4.
-        Thus we need to have a cache that is larger. How much larger, as much as
-        the sliding window size. So twice the max_context_length.
+        Thus we need a cache larger than the sliding window by enough space for
+        the largest in-flight input chunk.
         How would that have helped. Lets see. At pos = 5 our cache would have
         [0, 1, 2, 3, 4, NA, NA, NA] After cache update we would have
         [8, 1, 2, 3, 4, 5, 6, 7]. We kicked out token at pos = 0. However, the
         current step still has access to [pos - sliding_window_size, pos] tokens.
-        
+
         To make sure we dont over attend, i.e. we dont have pos = 5
         to attend to pos = 1, mask calculaton has to account for the sliding window
         size.
         """
+        self.window_size = window_size
+        self.full_context_length = max_context_length
+        self.max_seq_len = (
+            max_context_length if max_seq_len is None else int(max_seq_len)
+        )
+        ring_cache_size = _get_ring_cache_size(
+            max_context_length, window_size, self.max_seq_len
+        )
         super().__init__(
             max_batch_size,
-            max_context_length * 2,
+            ring_cache_size,
             n_heads,
             head_dim,
             enable_dynamic_shape,

examples/models/llama/export_llama_lib.py

@@ -969,6 +969,7 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager:
             preq_group_size=llm_config.base.preq_group_size,
             preq_embedding_quantize=llm_config.base.preq_embedding_quantize,
             local_global_attention=llm_config.model.local_global_attention,
+            max_seq_len=llm_config.export.max_seq_length,
             use_torchao_kernels_linear=llm_config.backend.torchao.use_torchao_kernels_linear,
             use_torchao_kernels_tied_embedding=llm_config.backend.torchao.use_torchao_kernels_tied_embedding,
             quantize_with_hqq=llm_config.quantization.use_hqq,
@@ -2064,6 +2065,7 @@ def _get_source_transforms(  # noqa
     preq_group_size: Optional[int] = None,
     preq_embedding_quantize: Optional[str] = None,
     local_global_attention: Optional[List[int]] = None,
+    max_seq_len: Optional[int] = None,
     use_torchao_kernels_linear: bool = False,
     use_torchao_kernels_tied_embedding: bool = False,
     quantize_with_hqq: bool = True,
@@ -2104,6 +2106,8 @@ def _get_source_transforms(  # noqa
         preq_mode: Pre-quantization mode.
         preq_group_size: Pre-quantization group size.
         preq_embedding_quantize: Pre-quantization embedding quantize.
+        max_seq_len: Largest input chunk accepted by local-attention layers. If
+            omitted, each layer's full context length is used.
 
     Returns:
         A list of transformation functions.
@@ -2280,6 +2284,7 @@ def _get_source_transforms(  # noqa
             partial(
                 replace_kv_cache_with_ring_kv_cache,
                 layer_sizes=local_global_attention,
+                max_seq_len=max_seq_len,
             )
         )
 

examples/models/llama/source_transformation/attention_sink.py

@@ -18,13 +18,36 @@
     _create_causal_mask_for_ring_buffer,
     AttentionMHA,
     KVCache,
-    RingKVCache,
 )
 from executorch.examples.models.llama.model_args import ModelArgs
 from executorch.examples.models.llama.rope import Rope
 from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter
 
 
+def _get_attention_sink_cache_size(
+    max_context_length: int,
+    window_size: int,
+    sink_size: int,
+    max_seq_len: Optional[int] = None,
+) -> int:
+    """Size a sink cache for fixed sinks, one window, and one input chunk."""
+    assert sink_size >= 0, "Attention sink size must be non-negative"
+    assert sink_size < max_context_length, (
+        f"Attention sink size ({sink_size}) must be smaller than the full "
+        f"context length ({max_context_length})"
+    )
+    assert window_size > 0, "Sliding-window size must be positive"
+    assert sink_size + window_size <= max_context_length, (
+        f"Attention sink size ({sink_size}) plus sliding-window size "
+        f"({window_size}) cannot exceed the full context length "
+        f"({max_context_length})"
+    )
+    if max_seq_len is None:
+        max_seq_len = max_context_length
+    assert max_seq_len > 0, "Maximum sequence length must be positive"
+    return sink_size + window_size + max_seq_len
+
+
 class RopeWithAttentionSink(Rope):
     """
     Rope subclass for Attention Sink models.
@@ -37,10 +60,11 @@ 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 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.
+    The ring buffer holds one retained window plus one maximum-size input
+    chunk. It is larger than the live window 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__(
@@ -52,7 +76,26 @@ def __init__(
         super().__init__(params)
         self.window_size = window_size
         self.sink_size = sink_size
-        self.ring_size = window_size * 2
+        max_seq_len = (
+            params.max_context_len
+            if getattr(params, "max_seq_len", None) is None
+            else int(params.max_seq_len)
+        )
+        cache_size = _get_attention_sink_cache_size(
+            params.max_context_len,
+            window_size,
+            sink_size,
+            max_seq_len,
+        )
+        self.ring_size = cache_size - sink_size
+        if self.freqs_cos.size(0) < cache_size:
+            freqs_cos, freqs_sin = self.precompute_freqs_cis(
+                self.params.head_dim,
+                cache_size,
+                self.params.rope_freq_base,
+            )
+            self.freqs_cos = freqs_cos
+            self.freqs_sin = freqs_sin
 
     def _remap_input_pos(self, input_pos: torch.Tensor) -> torch.Tensor:
         """Remap positions: sink tokens stay, window tokens wrap in ring buffer."""
@@ -133,7 +176,7 @@ class CachePositionsManagerWithSink(nn.Module):
     For sink_size=0: behaves exactly like original CachePositionsManager.
     For sink_size>0: sink tokens go to fixed positions, rest uses ring buffer.
 
-    IMPORTANT: cache_size should be the actual cache dimension size (2x window for ring buffer).
+    IMPORTANT: cache_size is the actual cache dimension, including sink slots.
     """
 
     def __init__(self, cache_size: int, sink_size: int = 0):
@@ -191,7 +234,7 @@ class KVCacheWithAttentionSink(KVCache):
     Uses a ring buffer approach for the sliding window portion while keeping
     the first sink_size tokens fixed. This avoids dynamic shape operations.
 
-    Cache layout: [sink: 0 to sink_size-1] [ring_buffer: sink_size to sink_size + window_size*2 - 1]
+    Cache layout: [fixed sink tokens] [ring buffer for one window + one chunk]
     """
 
     def __init__(
@@ -202,16 +245,27 @@ def __init__(
         rope: RopeWithAttentionSink,
         window_size: int,
         sink_size: int,
+        max_context_length: int,
+        max_seq_len: Optional[int] = None,
         max_batch_size: int = 1,
         dtype=torch.float32,
     ):
-        # Total cache size is sink_size + window_size * 2.
-        # The ring buffer needs 2x the window size because at the moment a new
-        # token is written, the previous window_size tokens must still be readable
-        # (they haven't been overwritten yet). With only 1x, writing a new entry
-        # would immediately evict the oldest visible token, leaving fewer than
-        # window_size tokens available for attention.
-        total_cache_size = sink_size + window_size * 2
+        self.full_context_length = max_context_length
+        self.max_seq_len = (
+            max_context_length if max_seq_len is None else int(max_seq_len)
+        )
+        # Keep the fixed sinks separate from the ring space needed for the
+        # retained window and the largest in-flight input chunk.
+        total_cache_size = _get_attention_sink_cache_size(
+            max_context_length,
+            window_size,
+            sink_size,
+            self.max_seq_len,
+        )
+        assert rope.ring_size == total_cache_size - sink_size, (
+            f"RoPE ring size ({rope.ring_size}) must match the KV-cache ring "
+            f"size ({total_cache_size - sink_size})"
+        )
         super().__init__(
             max_batch_size=max_batch_size,
             max_context_length=total_cache_size,
@@ -310,6 +364,7 @@ def _replace_attention(
     rope_with_attention_sink: RopeWithAttentionSink,
     sink_size: int,
     window_size: int,
+    max_seq_len: Optional[int],
 ):
     for _, child_module in module._modules.items():
         if len(list(child_module.children())) > 0:  # pyre-ignore [16]
@@ -318,32 +373,23 @@ def _replace_attention(
                 rope_with_attention_sink=rope_with_attention_sink,
                 sink_size=sink_size,
                 window_size=window_size,
+                max_seq_len=max_seq_len,
             )
 
         if isinstance(child_module, AttentionMHA):
             kv_cache = child_module.kv_cache
-            if sink_size == 0:
-                # No sink tokens needed — use standard RingKVCache directly
-                child_module.kv_cache = RingKVCache(
-                    kv_cache.max_batch_size,
-                    window_size,  # RingKVCache expects user-provided window size
-                    kv_cache.n_heads,
-                    kv_cache.head_dim,
-                    kv_cache.enable_dynamic_shape,
-                    kv_cache.k_cache.dtype,
-                )
-            else:
-                kv_cache_with_attention_sink = KVCacheWithAttentionSink(
-                    n_heads=kv_cache.n_heads,
-                    head_dim=kv_cache.head_dim,
-                    enable_dynamic_shape=kv_cache.enable_dynamic_shape,
-                    rope=rope_with_attention_sink,
-                    max_batch_size=kv_cache.max_batch_size,
-                    window_size=window_size,
-                    sink_size=sink_size,
-                    dtype=kv_cache.k_cache.dtype,
-                )
-                child_module.kv_cache = kv_cache_with_attention_sink
+            child_module.kv_cache = KVCacheWithAttentionSink(
+                n_heads=kv_cache.n_heads,
+                head_dim=kv_cache.head_dim,
+                enable_dynamic_shape=kv_cache.enable_dynamic_shape,
+                rope=rope_with_attention_sink,
+                max_batch_size=kv_cache.max_batch_size,
+                window_size=window_size,
+                sink_size=sink_size,
+                max_context_length=kv_cache.max_context_length,
+                max_seq_len=max_seq_len,
+                dtype=kv_cache.k_cache.dtype,
+            )
             # Don't replace forward - let the original AttentionMHA.forward handle it
             # since our KVCache has is_ring_buffer=True, it will use the ring buffer mask
 
@@ -365,11 +411,13 @@ def enable_attention_sink(
         window_size=window_size,
         sink_size=sink_size,
     )
+    max_seq_len = getattr(params, "max_seq_len", None)
     _replace_rope(module, rope_with_attention_sink)
     _replace_attention(
         module=module,
         rope_with_attention_sink=rope_with_attention_sink,
         sink_size=sink_size,
         window_size=window_size,
+        max_seq_len=max_seq_len,
     )
     return module