[Executorch] Add MLX recipe

@@ -76,12 +76,30 @@ pip install transformers "torch==2.12.0" "onnx==1.21.0" "onnxscript==0.7.0" onnx ```  </hfoption>-<hfoption id="ExecuTorch">+<hfoption id="ExecuTorch XNNPACK">  ```bash pip install transformers "torch==2.12.0" "executorch==1.3.1" ``` +</hfoption>+<hfoption id="ExecuTorch MLX">++Requires macOS 14 or later on Apple Silicon. MLX is included in the macOS ARM64 ExecuTorch+nightly wheels; it is not a separate package or extra. The following nightly pair was validated+with the MLX exporter tests. Install it from the standard ExecuTorch nightly registry:++```bash+pip install transformers+pip install \+  "executorch==1.6.0.dev20260924" \+  "torch==2.15.0.dev20260924" \+  --extra-index-url https://download.pytorch.org/whl/nightly/cpu+```++Install `torch` explicitly because nightly ExecuTorch wheels do not declare it as a dependency.+For source builds, see the [MLX installation guide](https://docs.pytorch.org/executorch/main/backends/mlx/mlx-overview.html).+ </hfoption> <hfoption id="OpenVINO"> @@ -145,7 +163,10 @@ outputs = session.run(None, ort_inputs) </hfoption> <hfoption id="ExecuTorch"> -[`~exporters.ExecutorchConfig#backend`] defaults to `xnnpack` which targets the CPU and works on CPU-only installations. `cuda` targets the GPU and requires a CUDA-enabled environment. Requesting it without CUDA raises a `RuntimeError`.+[`~exporters.ExecutorchConfig#backend`] selects the target: `xnnpack` (the default) for CPU,+`mlx` for Apple Silicon GPU, or `cuda` for a CUDA-enabled GPU. Choose the corresponding+[installation option](#installation). MLX execution requires the MLX delegate and its Metal libraries;+requesting CUDA without a CUDA-enabled environment raises a `RuntimeError`.  ```python from transformers import AutoModelForCausalLM, AutoTokenizer@@ -156,7 +177,7 @@ tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") inputs = tokenizer("Hello, world!", return_tensors="pt")  exporter = ExecutorchExporter()-config = ExecutorchConfig(backend="xnnpack", dynamic=True)+config = ExecutorchConfig(backend="xnnpack", dynamic=True)  # Use "mlx" on Apple Silicon. et_program = exporter.export(model, inputs, config=config)  # save for on-device deployment@@ -277,7 +298,7 @@ seq = torch.export.Dim("seq", min=1, max=2048)  exporter = ExecutorchExporter() config = ExecutorchConfig(-    backend="xnnpack",+    backend="xnnpack",  # Use "mlx" on Apple Silicon.     dynamic_shapes={"input_ids": {0: batch, 1: seq}, "attention_mask": {0: batch, 1: seq}},     # Emit data-dependent shape guards as runtime asserts instead of failing the export when a     # guard wouldn't hold across the explicit symbolic range. Most LLMs need this under fine-grained@@ -388,7 +409,7 @@ text = processor.apply_chat_template(messages, add_generation_prompt=True, token inputs = processor(text=text, images=messages[0]["content"][0]["url"], return_tensors="pt").to(model.device)  exporter = ExecutorchExporter()-config = ExecutorchConfig(backend="xnnpack", dynamic=True)+config = ExecutorchConfig(backend="xnnpack", dynamic=True)  # Use "mlx" on Apple Silicon. components = exporter.export_for_generation(model, inputs, config=config) # components = {"image_encoder": ExecutorchProgramManager, "language_model": ..., "lm_head": ..., "decode": ...} ```@@ -490,7 +511,7 @@ components = exporter.export_for_generation(model, inputs, config=config, multi_ from transformers.exporters import ExecutorchExporter, ExecutorchConfig  exporter = ExecutorchExporter()-config = ExecutorchConfig(backend="xnnpack", dynamic=True)+config = ExecutorchConfig(backend="xnnpack", dynamic=True)  # Use "mlx" on Apple Silicon. components = exporter.export_for_generation(model, inputs, config=config, multi_token_decode=True) # components["decode"] now accepts a variable number of query tokens ```@@ -505,6 +526,10 @@ cache handles where they land internally.  ### Static KV cache +> [!NOTE]+> The ExecuTorch examples in this section, including zero-copy updates and the decode loop,+> are **XNNPACK-only**. The MLX exporter rejects `StaticCache`; use `DynamicCache` with MLX.+ `generate()` grows a `DynamicCache` by default, reallocating as the sequence extends — a moving target for an exported graph. A **static** cache is a fixed-size buffer, allocated once and written in place at the current position each step. Combined with a [multi-token decode](#multi-token-decode) it collapses@@ -542,7 +567,7 @@ components = exporter.export_for_generation( ```  </hfoption>-<hfoption id="ExecuTorch">+<hfoption id="ExecuTorch XNNPACK">  ```python from transformers import GenerationConfig@@ -577,7 +602,7 @@ own arena. What that takes is the only per-backend part left:   each pair to one device buffer, so the cache is read and updated in place across the loop with no host   round-trips. -- **ExecuTorch** — turn off the memory-planning allocations on [`ExecutorchConfig`] so the in-place+- **ExecuTorch (XNNPACK)** — turn off the memory-planning allocations on [`ExecutorchConfig`] so the in-place   write can land in the caller's own tensor (see the reference for what each flag does):    ```python@@ -718,7 +743,7 @@ for position in range(prompt_len, max_cache_len): ```  </hfoption>-<hfoption id="ExecuTorch">+<hfoption id="ExecuTorch XNNPACK">  ExecuTorch's on-device runtime is C++, and the in-place cache update relies on `Method::set_output_data_ptr` — **not surfaced by the Python runtime** (`executorch.runtime.Method`@@ -807,8 +832,8 @@ visible from the public `export` API, but the most common things to know: setting and `eager` also works (slower). Set one of them on the model before calling `export` if it's using something else. - `grouped_mm` traces fine through `DynamoExporter` and is auto-translated for `OnnxExporter`.-For `ExecutorchExporter` with the XNNPACK backend, the exporter swaps MoE experts to-`batched_mm` because XNNPACK has no `_grouped_mm.out` kernel.+For `ExecutorchExporter` with either XNNPACK or MLX, the exporter swaps MoE experts to+`batched_mm` before export.  ## Next steps 

@@ -163,6 +163,7 @@ class ExecutorchConfig(DynamoConfig):              - `"xnnpack"` — CPU inference via the XNNPACK library (default; runs anywhere).             - `"cuda"` — GPU inference via the ExecuTorch CUDA backend.+            - `"mlx"` — GPU inference via the ExecuTorch MLX backend on Apple Silicon.         alloc_graph_input (`bool`, *optional*, defaults to `True`):             Whether the memory-planning pass reserves arena memory for graph inputs. When `False`,             the runtime uses the caller-provided input buffers directly instead of copying into the

@@ -17,8 +17,8 @@ Extends `DynamoExporter` to produce an `ExecutorchProgramManager` for mobile and edge deployment. The export pipeline runs: -1. **Backend preparation** (`_BACKEND_PREPARE`): `prepare_for_xnnpack` / `prepare_for_cuda`-   move the model to the target device/dtype and build the partitioner list.+1. **Backend preparation** (`_BACKEND_PREPARE`): move the model to the target device/dtype+   and build the partitioner list. 2. **Torch patches** (`_PATCHES["executorch"]` via `apply_patches("executorch")`, plus the    backend-specific `_PATCHES[f"executorch.{backend}"]`): reversibly swap `torch` ops the    ExecuTorch backends can't accept (`split_copy`, `avg_pool2d`, …) with decomposed equivalents.@@ -67,6 +67,7 @@     from torch.utils._sympy.numbers import IntInfinity     from torch.utils._sympy.value_ranges import ValueRanges +    from ..cache_utils import EncoderDecoderCache, StaticCache     from ..modeling_utils import PreTrainedModel      # Runtime-assert ops dropped before lowering (see `_drop_runtime_asserts`).@@ -145,7 +146,10 @@ def export(             apply_fx_program_fixes("executorch", exported_program)             apply_fx_node_fixes("executorch", exported_program.graph_module)             edge_program_manager: EdgeProgramManager = to_edge_transform_and_lower(-                exported_program, partitioner=partitioner, compile_config=_get_edge_compile_config()+                exported_program,+                partitioner=partitioner,+                compile_config=_get_edge_compile_config(config.backend),+                transform_passes=_get_transform_passes(config.backend),             )             executorch_programs_manager: ExecutorchProgramManager = edge_program_manager.to_executorch(                 config=_get_backend_config(config)@@ -154,7 +158,16 @@ def export(         return executorch_programs_manager  -def _get_edge_compile_config() -> EdgeCompileConfig:+def _get_transform_passes(backend: str):+    """Return backend-specific graph transforms, or ``None`` for defaults."""+    if backend == "mlx":+        from executorch.backends.mlx.passes import get_default_passes++        return get_default_passes()+    return None+++def _get_edge_compile_config(backend: str) -> EdgeCompileConfig:     """Build the ``EdgeCompileConfig`` used for ``to_edge_transform_and_lower``.      Adds non-core ATen ops to ``_core_aten_ops_exception_list`` so torch.export@@ -164,6 +177,8 @@ def _get_edge_compile_config() -> EdgeCompileConfig:     but aren't in the core ATen opset. The CPU portable kernels handle them at     runtime; XNNPACK leaves them in the non-delegated CPU portion of the graph.     """+    if backend == "mlx":+        return EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True)     return EdgeCompileConfig(         _core_aten_ops_exception_list=[             torch.ops.aten._fft_c2c.default,@@ -258,9 +273,33 @@ def prepare_for_cuda(model: PreTrainedModel, sample_inputs: dict[str, Any]):     return model, _make_contiguous(sample_inputs), partitioner  +def prepare_for_mlx(model: PreTrainedModel, sample_inputs: dict[str, Any]):+    """Apple Silicon GPU inference via the ExecuTorch MLX backend."""+    for value in sample_inputs.values():+        caches = [value]+        if isinstance(value, EncoderDecoderCache):+            caches = [value.self_attention_cache, value.cross_attention_cache]+        if any(isinstance(cache, StaticCache) for cache in caches):+            raise ValueError(+                "StaticCache is not supported by the ExecuTorch MLX backend. "+                "Use DynamicCache or set cache_implementation='dynamic' in GenerationConfig."+            )++    from executorch.backends.mlx import MLXPartitioner++    model.requires_grad_(False)+    model = model.to(device="cpu")+    # MLX does not support grouped MoE kernels.+    if isinstance(model, PreTrainedModel) and model._can_set_experts_implementation():+        model.set_experts_implementation("batched_mm")+    partitioner = [MLXPartitioner()]+    return model, _make_contiguous(sample_inputs), partitioner++ _BACKEND_PREPARE = {     "xnnpack": prepare_for_xnnpack,     "cuda": prepare_for_cuda,+    "mlx": prepare_for_mlx, }  @@ -339,7 +378,8 @@ def patch(input, k, dim=None, largest=True, sorted=True):     return patch  -@register_patch("executorch", "torch.detach", "torch.Tensor.detach")+@register_patch("executorch.xnnpack", "torch.detach", "torch.Tensor.detach")+@register_patch("executorch.cuda", "torch.detach", "torch.Tensor.detach") def _patch_detach(_original):     """No-op detach.""" 

@@ -14,6 +14,7 @@  import copy import functools+import importlib.util import inspect import itertools import re@@ -47,6 +48,7 @@     slow,     torch_device, )+from transformers.utils import is_executorch_available   # ──────────────────────────── skip lists ────────────────────────────@@ -579,6 +581,38 @@ )  +_EXECUTORCH_BACKENDS = ("xnnpack",)+if is_executorch_available() and importlib.util.find_spec("executorch.backends.mlx") is not None:+    try:+        from executorch.runtime import Runtime++        if "MLXBackend" in Runtime.get().backend_registry.registered_backend_names:+            _EXECUTORCH_BACKENDS += ("mlx",)+    except ImportError:+        pass  # The Python backend can be installed without the native runtime.++EXECUTORCH_EXPORT_PARAMS = parameterized.expand(+    list(itertools.product(_EXECUTORCH_BACKENDS, _EXPORT_SHAPE_MODES)),+    name_func=lambda f, _, p: f"{f.__name__}_{'dynamic' if p.args[1] else 'static'}_{p.args[0]}",+)+EXECUTORCH_GENERATE_EXPORT_PARAMS = parameterized.expand(+    [+        (backend, dynamic, generation_config)+        for backend, dynamic, generation_config in itertools.product(+            _EXECUTORCH_BACKENDS, _EXPORT_SHAPE_MODES, _EXPORT_GENERATION_CONFIGS+        )+        if not (+            backend == "mlx" and generation_config is not None and generation_config.cache_implementation == "static"+        )+    ],+    name_func=lambda f, _, p: (+        f"{f.__name__}_{'dynamic' if p.args[1] else 'static'}"+        + (f"_{p.args[2].cache_implementation}_cache" if p.args[2] is not None else "")+        + f"_{p.args[0]}"+    ),+)++ def _needs_static_cache(generation_config) -> bool:     """True if `generation_config` requests a cache the model must explicitly support (a static impl).     Such variants only run on models that can (see the `_can_compile_fullgraph` gate in the tests)."""@@ -1081,24 +1115,25 @@ def test_openvino_export(self, dynamic, atol=1e-3, rtol=1e-3):      # ──────────────────── ExecuTorch tests ─────────────────────── -    @DYNAMIC_EXPORT_PARAMS+    @EXECUTORCH_EXPORT_PARAMS     @slow     @require_executorch     @pytest.mark.executorch_export_test     @pytest.mark.timeout(EXPORT_TEST_TIMEOUT)     @require_torch_greater_or_equal(MIN_EXPORT_TORCH_VERSION)     @disable_hub_kernels-    def test_executorch_export(self, dynamic):+    def test_executorch_export(self, backend, dynamic):         """Export each model class to ExecuTorch, run it, and verify output count matches eager."""          self._skip_if_not_exportable()         exporter = ExecutorchExporter()-        config = ExecutorchConfig(dynamic=dynamic)+        config = ExecutorchConfig(backend=backend, dynamic=dynamic)          for model_class in self.all_model_classes:-            if self._should_skip(model_class, dynamic=dynamic, backend="executorch"):+            if any(+                self._should_skip(model_class, dynamic=dynamic, backend=scope) for scope in ("executorch", backend)+            ):                 continue-             # Trace on CPU: XNNPACK targets CPU, and CPU tracing yields device-consistent graphs.             # Tracing on CUDA surfaces per-model device bugs — models create in-`forward` tensors             # (arange/zeros/sinusoids) without `device=`, which default to CPU and then mismatch a@@ -1282,26 +1317,28 @@ def test_openvino_export_generate(self, dynamic, generation_config):      # ──────────────────── ExecuTorch tests ─────────────────────── -    @GENERATE_EXPORT_PARAMS+    @EXECUTORCH_GENERATE_EXPORT_PARAMS     @slow     @require_executorch     @pytest.mark.executorch_export_test     @pytest.mark.timeout(EXPORT_TEST_TIMEOUT)     @require_torch_greater_or_equal(MIN_EXPORT_TORCH_VERSION)     @disable_hub_kernels-    def test_executorch_export_generate(self, dynamic, generation_config):+    def test_executorch_export_generate(self, backend, dynamic, generation_config):         """Export prefill and decode stages to ExecuTorch, run each, and verify output count matches eager."""          self._skip_if_not_exportable()         exporter = ExecutorchExporter()-        config = ExecutorchConfig(dynamic=dynamic)+        config = ExecutorchConfig(backend=backend, dynamic=dynamic)          for model_class in self.all_generative_model_classes:-            if self._should_skip(-                model_class, generate=True, dynamic=dynamic, backend="executorch", generation_config=generation_config+            if any(+                self._should_skip(+                    model_class, generate=True, dynamic=dynamic, backend=scope, generation_config=generation_config+                )+                for scope in ("executorch", backend)             ):                 continue-             components = self._prepare_export_generate_model_and_inputs(                 model_class,                 device="cpu",