@@ -262,23 +262,6 @@ 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,
@@ -288,17 +271,10 @@ 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
         """
-        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:
+        Reason why we want the kv cache size to be twice the context length:
         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
@@ -331,28 +307,20 @@ 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 a cache larger than the sliding window by enough space for
-        the largest in-flight input chunk.
+        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.
         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,
-            ring_cache_size,
+            max_context_length * 2,
             n_heads,
             head_dim,
             enable_dynamic_shape,
 
@@ -969,7 +969,6 @@ 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,
@@ -2065,7 +2064,6 @@ 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,
@@ -2106,8 +2104,6 @@ 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.
@@ -2284,7 +2280,6 @@ def _get_source_transforms(  # noqa
             partial(
                 replace_kv_cache_with_ring_kv_cache,
                 layer_sizes=local_global_attention,
-                max_seq_len=max_seq_len,
             )
         )
 
 
@@ -18,36 +18,13 @@
     _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.
@@ -60,11 +37,10 @@ class RopeWithAttentionSink(Rope):
       - Window tokens (pos >= sink_size): wrapped into ring buffer range
         [sink_size, sink_size + ring_size) via modulo
 
-    The ring buffer 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.
+    The ring buffer is 2x window_size for write-ahead headroom, not to keep the
+    live window contiguous -- it can span a wrap. Across a wrap two positions
+    remap to a difference that is not their true distance, so RoPE preserves
+    relative distance only within a wrap.
     """
 
     def __init__(
@@ -76,26 +52,7 @@ def __init__(
         super().__init__(params)
         self.window_size = window_size
         self.sink_size = sink_size
-        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
+        self.ring_size = window_size * 2
 
     def _remap_input_pos(self, input_pos: torch.Tensor) -> torch.Tensor:
         """Remap positions: sink tokens stay, window tokens wrap in ring buffer."""
@@ -176,7 +133,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 is the actual cache dimension, including sink slots.
+    IMPORTANT: cache_size should be the actual cache dimension size (2x window for ring buffer).
     """
 
     def __init__(self, cache_size: int, sink_size: int = 0):
@@ -234,7 +191,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: [fixed sink tokens] [ring buffer for one window + one chunk]
+    Cache layout: [sink: 0 to sink_size-1] [ring_buffer: sink_size to sink_size + window_size*2 - 1]
     """
 
     def __init__(
@@ -245,27 +202,16 @@ 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,
     ):
-        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})"
-        )
+        # 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
         super().__init__(
             max_batch_size=max_batch_size,
             max_context_length=total_cache_size,
@@ -364,7 +310,6 @@ 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]
@@ -373,23 +318,32 @@ 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
-            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,
-            )
+            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
             # 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
 
@@ -411,13 +365,11 @@ 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