[MLX] Add off-graph KV cache export mode for HF models

backends/mlx/examples/llm/export_llm_hf.py

@@ -312,6 +312,147 @@ def _export_with_custom_components(
     _save_program(executorch_program, output_path)
 
 
+def _export_with_offgraph_cache(
+    model_id: str,
+    revision: Optional[str],
+    output_path: str,
+    max_seq_len: int,
+    dtype: str,
+    qlinear: Optional[str],
+    qembedding: Optional[str],
+    no_tie_word_embeddings: bool = False,
+    qlinear_group_size: Optional[int] = None,
+    qembedding_group_size: Optional[int] = None,
+) -> None:
+    """
+    Export using the off-graph KV cache op (kvcache::update_and_attend).
+
+    Unlike the custom-components path, the cache is not a graph tensor: the model
+    is run with use_cache=False (no StaticCache wrapper, no HFStaticCache swap),
+    so each attention layer emits an update_and_attend node fed this step's k/v.
+    The cache is created and bound at runtime via a cache_key.
+    """
+    import executorch.exir as exir
+    from executorch.backends.mlx import MLXPartitioner
+    from executorch.backends.mlx.llm.hf_attention import (
+        OffGraphExportWrapper,
+        register_mlx_offgraph_attention,
+    )
+    from executorch.backends.mlx.passes import get_default_passes
+    from executorch.exir import EdgeCompileConfig
+    from executorch.exir.capture._config import ExecutorchBackendConfig
+    from executorch.exir.passes import MemoryPlanningPass
+    from transformers import AutoModelForCausalLM
+
+    torch_dtype_map = {
+        "fp32": torch.float32,
+        "fp16": torch.float16,
+        "bf16": torch.bfloat16,
+    }
+    torch_dtype = torch_dtype_map.get(dtype, torch.bfloat16)
+
+    register_mlx_offgraph_attention()
+    logger.info("Registered MLX off-graph attention (update_and_attend)")
+
+    logger.info(f"Loading HuggingFace model: {model_id}")
+    load_kwargs = {
+        "torch_dtype": torch_dtype,
+        "low_cpu_mem_usage": True,
+        "attn_implementation": "mlx_offgraph",
+    }
+    if revision is not None:
+        load_kwargs["revision"] = revision
+    model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs)
+    model.eval()
+
+    from executorch.backends.mlx.llm.quantization import quantize_model_
+
+    quantize_model_(
+        model,
+        qlinear_config=qlinear,
+        qlinear_group_size=qlinear_group_size,
+        qembedding_config=qembedding,
+        qembedding_group_size=qembedding_group_size,
+        tie_word_embeddings=getattr(model.config, "tie_word_embeddings", False)
+        and not no_tie_word_embeddings,
+    )
+
+    exportable = OffGraphExportWrapper(model)
+
+    # The cache is built before the graph runs, so the runtime cannot learn its
+    # shape from the model. Publish it, one entry per cache (KV-sharing models
+    # have fewer caches than layers):
+    #   n_caches  how many caches to allocate
+    #   kv_heads  KV heads per cache
+    #   head_dims head dim per cache -- gemma-4 mixes 256 and 512
+    #   windows   sliding window per cache, 0 = flat
+    from executorch.backends.mlx.llm.cache import resolve_hf_cache_layout
+
+    # resolve_hf_cache_layout drops the KV-shared tail, so these are indexed by
+    # cache -- the same indexing _cache_id maps layers onto.
+    layer_types, cache_kv_heads, cache_head_dims = resolve_hf_cache_layout(model.config)
+    text_config = model.config.get_text_config()
+    sliding_window = getattr(text_config, "sliding_window", None) or 0
+    cache_windows = [
+        sliding_window if t == "sliding_attention" else 0 for t in layer_types
+    ]
+    kv_metadata = {
+        "get_n_caches": len(layer_types),
+        "get_kv_heads": torch.tensor(cache_kv_heads, dtype=torch.int32),
+        "get_head_dims": torch.tensor(cache_head_dims, dtype=torch.int32),
+        "get_windows": torch.tensor(cache_windows, dtype=torch.int32),
+    }
+    logger.info(
+        f"KV cache layout: {len(layer_types)} caches, "
+        f"{sum(1 for w in cache_windows if w)} sliding (window {sliding_window})"
+    )
+
+    logger.info("Exporting model with torch.export...")
+    seq_length = 3
+    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=max_seq_len - 1)
+    dynamic_shapes = {
+        "input_ids": {1: seq_len_dim},
+        "cache_position": {0: seq_len_dim},
+    }
+
+    with torch.no_grad():
+        exported_program = torch.export.export(
+            exportable,
+            args=(),
+            kwargs={
+                "input_ids": example_input_ids,
+                "cache_position": example_cache_position,
+            },
+            dynamic_shapes=dynamic_shapes,
+            strict=True,
+        )
+
+    logger.info("Delegating to MLX backend...")
+    edge_program = exir.to_edge_transform_and_lower(
+        {"forward": exported_program},
+        transform_passes=get_default_passes(),
+        constant_methods=kv_metadata,
+        partitioner=[MLXPartitioner()],
+        compile_config=EdgeCompileConfig(
+            _check_ir_validity=False,
+            _skip_dim_order=True,
+        ),
+    )
+
+    logger.info("Exporting to ExecuTorch...")
+    executorch_program = edge_program.to_executorch(
+        config=ExecutorchBackendConfig(
+            extract_delegate_segments=True,
+            memory_planning_pass=MemoryPlanningPass(alloc_graph_input=True),
+        )
+    )
+
+    _save_program(executorch_program, output_path)
+
+
 def _save_program(executorch_program, output_path: str) -> None:
     """Save the ExecuTorch program to disk."""
     os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
@@ -332,6 +473,7 @@ def export_llama_hf(
     qembedding: Optional[str] = None,
     use_custom_sdpa: bool = False,
     use_custom_kv_cache: bool = False,
+    use_offgraph_cache: bool = False,
     no_tie_word_embeddings: bool = False,
     qlinear_group_size: Optional[int] = None,
     qembedding_group_size: Optional[int] = None,
@@ -349,7 +491,26 @@ def export_llama_hf(
         use_custom_sdpa: Use MLX custom SDPA (mlx::custom_sdpa)
         use_custom_kv_cache: Use MLX custom KV cache (mlx::kv_cache_update)
     """
-    if use_custom_sdpa or use_custom_kv_cache:
+    if use_offgraph_cache:
+        if use_custom_sdpa or use_custom_kv_cache:
+            raise ValueError(
+                "--use-offgraph-cache is exclusive with --use-custom-sdpa / "
+                "--use-custom-kv-cache (it replaces both)"
+            )
+        logger.info("Using off-graph KV cache (update_and_attend)")
+        _export_with_offgraph_cache(
+            model_id=model_id,
+            revision=revision,
+            output_path=output_path,
+            max_seq_len=max_seq_len,
+            dtype=dtype,
+            qlinear=qlinear,
+            qembedding=qembedding,
+            no_tie_word_embeddings=no_tie_word_embeddings,
+            qlinear_group_size=qlinear_group_size,
+            qembedding_group_size=qembedding_group_size,
+        )
+    elif use_custom_sdpa or use_custom_kv_cache:
         logger.info(
             f"Using custom components: sdpa={use_custom_sdpa}, "
             f"kv_cache={use_custom_kv_cache}"
@@ -434,6 +595,13 @@ def main():
         default=False,
         help="Use MLX custom KV cache (mlx::kv_cache_update)",
     )
+    parser.add_argument(
+        "--use-offgraph-cache",
+        action="store_true",
+        default=False,
+        help="Use the off-graph KV cache op (kvcache::update_and_attend); "
+        "replaces --use-custom-sdpa/--use-custom-kv-cache",
+    )
 
     args = parser.parse_args()
 
@@ -447,6 +615,7 @@ def main():
         qembedding=args.qembedding,
         use_custom_sdpa=args.use_custom_sdpa,
         use_custom_kv_cache=args.use_custom_kv_cache,
+        use_offgraph_cache=args.use_offgraph_cache,
         no_tie_word_embeddings=args.no_tie_word_embeddings,
         qlinear_group_size=args.qlinear_group_size,
         qembedding_group_size=args.qembedding_group_size,

backends/mlx/llm/cache.py

@@ -525,7 +525,9 @@ def update(
             # Current HF ExecuTorch wrappers copy the requested cache position
             # into each StaticCache layer's cumulative_length before forward().
             if hasattr(self.layers[layer_idx], "cumulative_length"):
-                cache_position = self.layers[layer_idx].cumulative_length
+                # cumulative_length is a scalar; KVCache.update indexes [0], so
+                # give it the 1-D shape a cache_kwargs caller would have passed.
+                cache_position = self.layers[layer_idx].cumulative_length.reshape(1)
             else:
                 raise RuntimeError(
                     "cache_position was not provided and the pinned "