[MLX] Fix HF sliding-window KV cache sizing and enable ring buffers (#21635)

**Summary**
export_llm_hf sized every layer's KV cache to the model's
sliding_window, which truncates the full-attention layers of a hybrid
model. Sliding attention was also never actually enabled on this path:
replace_hf_cache_with_mlx gave all layers plain linear caches and plain
causal SDPA.

Hybrid models now install ring buffers for sliding layers and
full-length linear caches for full-attention layers.
replace_hf_cache_with_mlx_ring_buffer takes max_cache_len separately
from window_size, because it previously applied one size to both. The
dynamic sequence dimension is bounded by the window rather than the
cache length, since a ring cannot absorb a single step longer than its
window.

**Files**
* source_transformation.py — replace_hf_cache_with_mlx_ring_buffer gains
a max_cache_len parameter so full-attention layers are sized to the
whole context while sliding layers stay bounded by the window, instead
of applying window_size to both.
* export_llm_hf.py — stops capping the cache length to the window,
routes models with a sliding_window to the ring-buffer cache and
sliding-window SDPA, and bounds the dynamic sequence dimension by the
window rather than the cache length.

**Test**

Exported:
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
        --model-id unsloth/gemma-3-1b-it --output gemma3_1b.pte \
--use-custom-sdpa --use-custom-kv-cache --max-seq-len 1024 --dtype fp32

Ran the .pte step by step through pybindings and compared each step's
top-1 token against HuggingFace eager, teacher-forced on the same token
sequence, over a full 1024-token context split at the sliding-window
boundary. Exported and run in fp32, the two agree on every step.
Repeating in bf16 gives a small number of disagreements on both sides of
the boundary (2.8% below, 1.4% above).

Author: kiymetakdemir

backends/mlx/examples/llm/export_llm_hf.py

@@ -184,13 +184,15 @@ def _export_with_custom_components(
     # Gemma 4 keep transformer attributes under text_config.
     text_config = model.config.get_text_config()
     sliding_window = getattr(text_config, "sliding_window", None)
+    # Full-attention layers retain all history, so they cannot be bounded by the
+    # sliding window. Sizing every layer to the window truncates them and the
+    # model loses its context past `sliding_window` tokens.
+    effective_cache_len = max_seq_len
     if sliding_window is not None:
-        logger.info(f"Model has sliding_window={sliding_window}")
-        # Cap max_seq_len to sliding window size for cache allocation
-        effective_cache_len = min(max_seq_len, sliding_window)
-        logger.info(f"  Capping cache length to sliding window: {effective_cache_len}")
-    else:
-        effective_cache_len = max_seq_len
+        logger.info(
+            f"Model has sliding_window={sliding_window}; "
+            f"cache length {effective_cache_len}"
+        )
 
     # The HF ExecuTorch cache wrappers validate both generation_config.use_cache
     # and the text config's use_cache flag before constructing static caches.
@@ -228,26 +230,52 @@ def _export_with_custom_components(
         )
 
     if use_custom_kv_cache:
-        from executorch.backends.mlx.llm.source_transformation import (
-            replace_hf_cache_with_mlx,
-        )
-
         if sliding_window is not None:
+            from executorch.backends.mlx.llm.source_transformation import (
+                replace_hf_cache_with_mlx_ring_buffer,
+            )
+
             logger.info(
-                "Replacing HuggingFace StaticCache with HFStaticCache "
-                f"(capped to sliding window: {effective_cache_len})..."
+                "Replacing HuggingFace HybridCache with MLX ring buffers "
+                f"(window {sliding_window}, cache length {effective_cache_len})..."
+            )
+            replace_hf_cache_with_mlx_ring_buffer(
+                exportable,
+                model.config,
+                max_batch_size=1,
+                window_size=sliding_window,
+                max_cache_len=effective_cache_len,
+                dtype=torch_dtype,
             )
         else:
-            logger.info("Replacing HuggingFace StaticCache with HFStaticCache...")
+            from executorch.backends.mlx.llm.source_transformation import (
+                replace_hf_cache_with_mlx,
+            )
 
-        replace_hf_cache_with_mlx(
-            exportable,
-            model.config,
-            max_batch_size=1,
-            max_cache_len=effective_cache_len,
-            dtype=torch_dtype,
-        )
-        logger.info("  HFStaticCache installed successfully")
+            logger.info(
+                "Replacing HuggingFace StaticCache with HFStaticCache "
+                f"(cache length {effective_cache_len})..."
+            )
+            replace_hf_cache_with_mlx(
+                exportable,
+                model.config,
+                max_batch_size=1,
+                max_cache_len=effective_cache_len,
+                dtype=torch_dtype,
+            )
+        logger.info("  MLX cache installed successfully")
+
+        if use_custom_sdpa and sliding_window is not None:
+            from executorch.backends.mlx.llm.hf_attention import (
+                register_mlx_sliding_window_attention,
+            )
+
+            # Registered after the wrapper exists: the SDPA closure captures it
+            # to reach the ring buffers when building masks.
+            register_mlx_sliding_window_attention(exportable)
+            model.config._attn_implementation = "mlx_sliding_window"
+            text_config._attn_implementation = "mlx_sliding_window"
+            logger.info("Registered MLX sliding-window SDPA")
 
     from executorch.backends.mlx.llm.quantization import quantize_model_
 
@@ -266,7 +294,13 @@ def _export_with_custom_components(
     example_input_ids = torch.zeros((1, seq_length), dtype=torch.long)
     example_cache_position = torch.arange(seq_length, dtype=torch.long)
 
-    seq_len_dim = torch.export.Dim("seq_length_dim", max=effective_cache_len - 1)
+    # A sliding layer holds `sliding_window` cells, so it cannot absorb a single
+    # step longer than that -- a separate limit from how much history the caches
+    # retain.
+    max_step_len = (
+        min(max_seq_len, sliding_window) if sliding_window is not None else max_seq_len
+    )
+    seq_len_dim = torch.export.Dim("seq_length_dim", max=max_step_len - 1)
     dynamic_shapes = {
         "input_ids": {1: seq_len_dim},
         "cache_position": {0: seq_len_dim},

backends/mlx/llm/source_transformation.py

@@ -163,6 +163,7 @@ def replace_hf_cache_with_mlx_ring_buffer(
     config,
     max_batch_size: int = 1,
     window_size: int = 512,
+    max_cache_len: int | None = None,
     dtype: torch.dtype = torch.float32,
 ) -> nn.Module:
     """
@@ -176,19 +177,26 @@ def replace_hf_cache_with_mlx_ring_buffer(
         module: HF exportable module with static_cache or cache attribute
         config: HF model config
         max_batch_size: Maximum batch size (default: 1)
-        window_size: Sliding window size (cache capacity per layer)
+        window_size: Sliding window size (capacity of the sliding-layer rings)
+        max_cache_len: Capacity of the full-attention layers; defaults to
+            ``window_size``, which is only correct for models with no
+            full-attention layers
         dtype: Cache tensor dtype
 
     Raises:
         ValueError: If module has no recognized cache attribute
     """
     from transformers.cache_utils import StaticCache
 
+    # Full-attention layers retain every position, so they are sized to the whole
+    # context; only the sliding layers are bounded by the window.
+    full_cache_len = max_cache_len if max_cache_len is not None else window_size
+
     # Create HFStaticCache with ring buffer layers
     mlx_cache = HFStaticCache(
         config=config,
         max_batch_size=max_batch_size,
-        max_cache_len=window_size,
+        max_cache_len=full_cache_len,
         dtype=dtype,
     )
 
@@ -242,8 +250,9 @@ def _install_cache(attr_name):
         raise ValueError("Module must have 'static_cache' or 'cache' attribute")
 
     logger.info(
-        f"Installed hybrid MLX cache: {num_ring_layers} ring-buffer layers / "
-        f"{num_cache_layers} total cache layers, window_size={window_size}"
+        f"Installed hybrid MLX cache: {num_ring_layers} ring-buffer layers "
+        f"(window_size={window_size}) / {num_cache_layers} total cache layers "
+        f"(full-attention length {full_cache_len})"
     )
 
     return module