[MLX] Reduce physical footprint memory in RingBufferKVCache for chunked prefill (pytorch#20341)

When doing chunked prefill, the RingBufferKVCache does not need 2x
window size, but instead window_size + max_write_length - 1 (prefill
chunk size). This PR exposes that knob and wires it to gemma4 31b MLX
export, which uses chunk_size 256, smaller than gemma4's window size
(1024).

Reduces phys_footprint on a 4K export by around −0.68 GiB (from 13.84
GiB to 13.16 GiB).

---------

Co-authored-by: uddeshsingh <uddeshsingh@gmail.com>

Author: metascroy

backends/mlx/llm/cache.py

@@ -228,7 +228,12 @@ class RingBufferKVCache(nn.Module):
     the attention function creates the mask lazily by accessing the cache
     via a closure. This avoids tracing issues with torch.export.
 
-    Layout: BHSD [batch_size, num_heads, window_size, head_dim]
+    Layout: BHSD [batch_size, num_heads, window_size + max_write_len - 1, head_dim]
+
+    ``max_write_len`` bounds the largest single ``update`` write and sets the
+    buffer slack beyond the window (buffer_size = window_size + max_write_len -
+    1). It defaults to the full context length, i.e. ``2 * window_size - 1``; pass the chunk size when doing
+    chunked prefill to shrink it further.
 
     Example:
         >>> cache = RingBufferKVCache(
@@ -250,6 +255,7 @@ def __init__(
         n_heads: int,
         head_dim: int,
         dtype: torch.dtype = torch.float32,
+        max_write_len: int | None = None,
     ):
         super().__init__()
         assert (
@@ -258,13 +264,23 @@ def __init__(
         self.max_batch_size = max_batch_size
         self.max_context_length = max_context_length
         self.window_size = max_context_length
-        self.buffer_size = 2 * max_context_length
+        # Slack beyond the window so a multi-token write never overwrites data
+        # that earlier queries in the same batch still need. After a write of
+        # seq_len tokens, the in-window positions that must coexist span a
+        # contiguous range of seq_len + window_size - 1 positions, so the ring
+        # needs buffer_size >= window_size + max_write_len - 1. Defaults to the
+        # full window; with chunked prefill, pass the chunk size to shrink it.
+        self.max_write_len = (
+            max_write_len if max_write_len is not None else max_context_length
+        )
+        assert (
+            self.max_write_len <= self.window_size
+        ), f"max_write_len={self.max_write_len} must be <= window_size={self.window_size}"
+        self.buffer_size = self.window_size + self.max_write_len - 1
         self.n_heads = n_heads
         self.head_dim = head_dim
 
-        # Cache buffers [B, H, 2*window_size, D]
-        # 2× buffer ensures multi-token writes never overwrite data that
-        # earlier queries in the same batch still need (matches ET behavior).
+        # Cache buffers [B, H, window_size + max_write_len - 1, D]
         cache_shape = (max_batch_size, n_heads, self.buffer_size, head_dim)
         self.register_buffer(
             "k_cache", torch.zeros(cache_shape, dtype=dtype, device="cpu")
@@ -292,7 +308,7 @@ def update(
             seq_len = k_val.size(2)
             torch._check(seq_len == v_val.size(2))
             torch._check(start_pos >= 0)
-            torch._check(seq_len <= self.window_size)
+            torch._check(seq_len <= self.max_write_len)
         else:
             start_pos = input_pos
 

backends/mlx/test/test_ops.py

@@ -2704,6 +2704,7 @@ def __init__(
         max_context_length: int,
         n_heads: int,
         head_dim: int,
+        max_write_len: Optional[int] = None,
     ):
         super().__init__()
         from executorch.backends.mlx.llm.cache import RingBufferKVCache
@@ -2713,6 +2714,7 @@ def __init__(
             max_context_length=max_context_length,
             n_heads=n_heads,
             head_dim=head_dim,
+            max_write_len=max_write_len,
         )
 
     def forward(
@@ -2749,18 +2751,31 @@ def __init__(
         head_dim: int = 64,
         max_context_length: int = 64,
         seq_step: int = 4,
+        max_write_len: Optional[int] = None,
     ):
         self.max_batch_size = max_batch_size
         self.n_heads = n_heads
         self.head_dim = head_dim
         self.max_context_length = max_context_length
         self.seq_step = seq_step
+        self.max_write_len = max_write_len
 
     @classmethod
     def get_test_configs(cls) -> List["RingBufferKVCacheTest"]:
         return [
             cls(),
             cls(n_heads=8, head_dim=32, max_context_length=32, seq_step=2),
+            # Reduced buffer (max_write_len < window) with a wrapping write:
+            # window=8, max_write_len=4 -> buffer_size=11. The test write of 4
+            # tokens at start_pos=8 wraps the ring (slots 8,9,10 then 0),
+            # exercising the tighter buffer
+            cls(
+                n_heads=8,
+                head_dim=32,
+                max_context_length=8,
+                seq_step=2,
+                max_write_len=4,
+            ),
         ]
 
     def create_model(self) -> nn.Module:
@@ -2769,6 +2784,7 @@ def create_model(self) -> nn.Module:
             max_context_length=self.max_context_length,
             n_heads=self.n_heads,
             head_dim=self.head_dim,
+            max_write_len=self.max_write_len,
         )
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
@@ -2793,7 +2809,13 @@ def create_test_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (input_pos, k_val, v_val)
 
     def get_dynamic_shapes(self) -> Optional[Dict[str, any]]:
-        seq_dim = Dim("seq_step", min=1, max=self.max_context_length)
+        # A single write can't exceed max_write_len (defaults to the window).
+        seq_max = (
+            self.max_write_len
+            if self.max_write_len is not None
+            else self.max_context_length
+        )
+        seq_dim = Dim("seq_step", min=1, max=seq_max)
         return {
             "input_pos": None,
             "k_val": {2: seq_dim},

examples/models/gemma4_31b/export.py

@@ -328,13 +328,17 @@ def _export_mlx(
     from executorch.exir.passes import MemoryPlanningPass
     from torch.export import Dim, export
 
+    max_prefill = 256
+
     mlx_source_transformations(
-        model, dtype=torch.bfloat16, use_turboquant=use_turboquant
+        model,
+        dtype=torch.bfloat16,
+        use_turboquant=use_turboquant,
+        max_write_len=max_prefill,
     )
 
     materialize_runtime_buffers(model, dtype=torch.bfloat16)
 
-    max_prefill = 256
     seq_dim = Dim("seq_len", min=1, max=max_prefill)
 
     print(f"Exporting (T in [1, {max_prefill}])...")

examples/models/gemma4_31b/mlx_source_transformations.py

@@ -158,6 +158,7 @@ def mlx_source_transformations(
     model: nn.Module,
     dtype: torch.dtype = torch.bfloat16,
     use_turboquant: bool = False,
+    max_write_len: int | None = None,
 ) -> None:
     """Apply MLX source transformations to a Gemma 4 31B model in-place.
 
@@ -177,19 +178,33 @@ def mlx_source_transformations(
         use_turboquant: If True, swap full-attention layers' KV caches
             for ``MLXTurboQuantKVCache`` (~3.8× cache memory savings).
             Sliding-window layers are unaffected.
+        max_write_len: Largest single prefill chunk written to the sliding
+            ring caches (the export-time max prefill / dynamic seq_len bound),
+            clamped to ``sliding_window``. Sizes the ring buffer to
+            ``sliding_window + max_write_len - 1``. If None, defaults to the
+            full window, i.e. ``2 * sliding_window - 1``.
     """
     config = model.config
 
     for layer in model.layers:
         attn = layer.self_attn
 
         if attn.is_sliding:
+            # A single write into a sliding ring buffer can never exceed the
+            # window, so clamp the chunk size to the window. Tiny test configs
+            # can have sliding_window < the export max prefill.
+            sliding_write_len = (
+                min(max_write_len, config.sliding_window)
+                if max_write_len is not None
+                else None
+            )
             attn.kv_cache = MLXRingKVCache(
                 max_batch_size=1,
                 max_context_length=config.sliding_window,
                 n_heads=attn.n_kv_heads,
                 head_dim=attn.head_dim,
                 dtype=dtype,
+                max_write_len=sliding_write_len,
             )
             attn.is_turboquant = False
         elif use_turboquant: