ExecuTorch MLX delegate by metascroy · Pull Request #16718 · pytorch/executorch

metascroy · 2026-01-20T22:10:33Z

Summary

This PR adds an MLX backend for ExecuTorch, enabling Metal-accelerated inference on Apple Silicon. It runs Llama, Qwen, Gemma, Whisper, Voxtral, and Parakeet models end-to-end, with 637 passing op tests and multithreaded execution support. For many models, it offers best performance among all ExecuTorch backends on Apple Silicon, offering 2-6x speedups over what was previously possible with ExecuTorch, and up to 30% smaller model sizes compared to XNNPACK due to BF16 support and tied quantized embedding support.

The PR is large due to extensive op coverage, testing, and documentation, but almost all changes are confined to backends/mlx/. The design is described in backends/mlx/README.md.

Suggested review approach:

Review files outside backends/mlx/ carefully — these integrate with ExecuTorch's build system and are the most likely to need changes.
For backends/mlx/, focus on structural design (see README) and test coverage (CI job is .github/workflows/mlx.yml)

Prerequisite PRs

These fixes were developed alongside the MLX backend. Once merged, this PR can be rebased to remove the duplicated changes:

#17257 — Improve lowering time with NamedDataMap
#17679 — Allow transform passes in etLLM
#17678 — Fix dynamic shape bug in remove_noop_pass
#17378 — Fix pocketfft intermittent bus errors on macOS (upstream fix)

Tests

CI is defined in .github/workflows/mlx.yml:

test_ops.py: 637 passing op tests
Multithreading: launches models on 50 threads, verifies correctness
GenAI E2E: parakeet, voxtral, etLLM (stories110m), HF LLMs (llama, qwen, gemma)
backend-tester: 380 passed, 1 failed, 86 skipped operator tests; 34 passed, 3 failed, 5 skipped model tests

The 1 failing operator test is a test-side issue being fixed in #17539. The 3 failing model tests will be addressed in follow-ups — they are not an initial focus compared to the GenAI models above.

pytorch-bot · 2026-01-20T22:10:38Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16718

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 2 Cancelled Jobs, 1 Unrelated Failure

As of commit d8ee9d2 with merge base 25f2a3f ():

NEW FAILURES - The following jobs have failed:

MLX / test-mlx-voxtral-realtime / test-mlx-voxtral-realtime (gh)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
pull / test-samsung-models-linux / linux-job (gh)
test_edsr_fp16
Test CUDA Windows Export and E2E / test-model-cuda-windows-e2e (mistralai, Voxtral-Mini-3B-2507, non-quantized) / windows-job (gh)
Process completed with exit code 1.

CANCELLED JOBS - The following jobs were cancelled. Please retry:

Test ARM Backend / test-arm / test-backend-linux (arm_tosa_fp, models) / linux-job (gh)
##[error]The operation was canceled.
Test ARM Backend / test-arm / test-backend-linux (arm_vgf_fp, models) / linux-job (gh)
##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

Test CUDA Windows Export and E2E / test-model-cuda-windows-e2e (mistralai, Voxtral-Mini-3B-2507, quantized-int4-weight-only) / windows-job (gh) (trunk failure)
Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

github-actions · 2026-01-20T22:11:20Z

This PR needs a `release notes:` label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

backends/mlx/examples/llm/export_llm_hf.py

metascroy · 2026-02-24T20:23:05Z

extension/llm/export/builder.py

@@ -451,7 +451,9 @@ def to_backend(self, partitioners: Optional[List[Partitioner]]) -> "LLMEdgeManag
        return self

    def to_edge_transform_and_lower(
-        self, partitioners: Optional[List[Partitioner]]
+        self,


Will remove once #17679 lands

metascroy · 2026-02-24T20:23:25Z

exir/passes/remove_noop_pass.py

-                # Only do this check if all the dims are static.
-                if all(isinstance(dim, int) for dim in orig_tensor.size()):
-                    if orig_tensor.shape == node.meta["val"].shape:
+                output_tensor = node.meta["val"]


Will remove once #17678 lands

Approved that PR

backends/mlx/examples/llm/export_llm_hf.py

This introduces a CSE pass to ExecuTorch, which eliminates common subexpressions that occur in exported programs. This pass was first developed as part of the MLX delegate (#16718) to optimize transformers, but I'm introducing it to ExecuTorch more generally because I believe it could benefit many other backends. Examples of common subexpressions that occur in transformers: * Repeated mask constructions per layer (only needs to be done once) * Repeated extraction of symints from 1d tensors for cache position (emits .item calls, which cause tensor materialization) This pass eliminates these inefficiencies without having to rewrite the model.

JacobSzwejbka · 2026-02-27T21:15:02Z

backends/mlx/serialization/schema.fbs

+}
+
+// For fields that can be either a literal int or a runtime Vid
+table IntOrVid {


Why not just use a union type?

FlatBuffer unions can only contain tables, not scalar primitives. Using a union here would require wrapping the int64 literal in a table (table IntLiteral { value: int64; }), adding pointer indirection and a table allocation for every shape dimension. The is_vid bool keeps the common case (literal) as a direct inline field read.

The more important reason is flatbutter doesn't allow vector of union

JacobSzwejbka · 2026-02-27T21:15:54Z

backends/mlx/serialization/schema.fbs

+    scale: float = 1.0;
+}
+
+table SdpaNode {


Whats the story for out of tree custom ops? Just not possible?

What do you mean by out of tree custom ops?

From MLX perspective, they allow custom ops using raw shader code. I guess we could add a node for that to the schema.

JacobSzwejbka · 2026-02-27T21:19:03Z

backends/mlx/serialization/schema.fbs

+    causal: bool = false;
+}
+
+table AddNode {


Similarly what happens if MLX changes the schema of any of these operators in a BC breaking way, or even adds a new optional arg with default values. Seems like there would be a lot of op versioning issues by explicitly encoding every ops schema in the serialization format. We also have a lot less influence on their op versioning story compared with ATen (where its already a big pita).

AddNode is an abstraction in the MLX delegate that we control (not part of MLX itself).

The runtime takes AddNode and translates it to an MLX call (add op) by parsing the args. So even if mlx changes the signature of their add op, it doesn't mean we break AddNode in the graph. We would just update our runtime appropriately. (With that said, MLX did say they try to maintain BC of their underlying ops).

Maybe I'm missing the concern here, though?

JacobSzwejbka · 2026-02-27T21:24:28Z

backends/mlx/ops.py

+        return out
+
+except ImportError:
+    # Edge IR ops not available (e.g., when building from ATen dialect)


When are you doing this?

I'm not. I'll remove this

JacobSzwejbka · 2026-02-27T21:28:57Z

backends/mlx/patterns.py

+
+
+@REGISTRY.register_pattern(name="ET_KV_CACHE_UPDATE")
+class ETKVCacheUpdateHandler(PatternHandler):


Why is both a custom mutator op supported, but also the slice_update ops earlier?

I guess you mean the difference between the IndexUpdateNode pattern above and this one?

They are similar, but ET_KV_CACHE_UPDATE has slightly better perf and is tied to a KVCache building block in mlx/llm/cache.py.

In terms of the difference:

INDEX_UPDATE matches a functionalized copy_ pattern (used in many HF models) and maps it onto IndexUpdateNode, which is an in-place update op in the MLX backend. IndexUpdateNode works by finding contiguous ranges in indices and repeatedly calling slice update. So in practice it will be 1 slice update call for regular kv cache, but up to 2 calls for ring buffer stuff, but it does require the indices be materialized in order to parse the index ranges (which can slow down MLX computation).

ET_KV_CACHE_UPDATE emits as one slice update call, and does not require materialization of the indices tensor.

At this point, the benefit is maybe 10-15% extra perf.

JacobSzwejbka · 2026-02-27T21:39:12Z

backends/mlx/pattern_utils.py

+
+
+@dataclass
+class PatternMatch:


XNNPack has their own pattern matching util iirc. Could we just offer this in Exir cc @GregoryComer

Happy to reuse any pre-existing utilities if they're in exir.

If the XNNPACK one is based on the torch.fx pattern finding util, I don't know if it'll work, though. I recall it being very bad at fuzzy matching of args/kwargs (like you needed exact matches, so you have to write like 3-4 pattern matchers for each op, and it's a bit error prone)

backends/mlx/examples/voxtral/export_voxtral_hf.py

manuelcandales · 2026-03-02T21:15:28Z

backends/mlx/examples/voxtral/export_voxtral_hf.py

+
+    logger.info("Exporting audio preprocessor with MLX backend...")
+
+    model = WhisperAudioProcessor(


Rename if this is used in models other than Whisper (like here in Voxtral)

metascroy · 2026-03-03T02:12:33Z

In the runtime unittest, can we compile with -Wall, -Werror, and -Wconversion, -Wsign-conversion, -Wshorten-64-to-32, asan, ubsan flag to uncover any security issues

Use overflow-safe arithmetic for all bounds checks. Use __builtin_add_overflow / __builtin_mul_overflow or c10's safe arithmetic utilities.

Check for null pointer deferences

I added a strict compile test, addressed security issues, and addressed null ptr references.

metascroy requested review from cccclai, kirklandsign, larryliu0820 and shoumikhin as code owners January 20, 2026 22:10

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jan 20, 2026

mergennachin requested review from manuelcandales and mergennachin January 20, 2026 22:37

metascroy force-pushed the mlx-delegate branch from 5e84aa5 to 1867cfc Compare January 27, 2026 01:58

metascroy requested review from JacobSzwejbka and lucylq as code owners January 30, 2026 06:33

metascroy force-pushed the mlx-delegate branch from 6541f23 to f296724 Compare February 9, 2026 04:41

metascroy force-pushed the mlx-delegate branch 7 times, most recently from e0e015c to 240b241 Compare February 25, 2026 01:06

metascroy changed the title ~~[draft] MLX delegate~~ ExecuTorch MLX delegate Feb 25, 2026

larryliu0820 reviewed Feb 26, 2026

View reviewed changes

backends/mlx/examples/llm/export_llm_hf.py Show resolved Hide resolved

metascroy mentioned this pull request Feb 26, 2026

Introduce CSE pass to ExecuTorch #17752

Merged

metascroy commented Feb 26, 2026

View reviewed changes

metascroy and others added 5 commits February 27, 2026 11:56

MLX delegate

54f1cdc

backends/apple/mlx/patches/mlx_json.patch

bff9194

Exclude auto-generated files from git

ed278de

up

585f932

more testing

1930a32

JacobSzwejbka reviewed Feb 27, 2026

View reviewed changes

metascroy added 13 commits February 28, 2026 00:36

up

f08f1dd

up

dc25245

up

0c76afa

up

f8f571f

up

1842ad3

up

10ff69f

up

339d1f0

up

2f45116

up

d50e8d6

Merge branch 'main' into mlx-delegate

2c73c84

up

868d955

up

c247996

up

89f202d

manuelcandales reviewed Mar 2, 2026

View reviewed changes

backends/mlx/examples/voxtral/export_voxtral_hf.py Show resolved Hide resolved

manuelcandales reviewed Mar 2, 2026

View reviewed changes

metascroy added 7 commits March 2, 2026 13:29

up

cc0a110

up

508498a

up

ede3a6e

up

fab5620

up

d181635

up

d4ba78f

up

d8ee9d2

metascroy mentioned this pull request Mar 3, 2026

MLX delegate (part 1) #17803

Open



		@REGISTRY.register_pattern(name="ET_KV_CACHE_UPDATE")
		class ETKVCacheUpdateHandler(PatternHandler):


		logger.info("Exporting audio preprocessor with MLX backend...")

		model = WhisperAudioProcessor(

Conversation

metascroy commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Prerequisite PRs

Tests

Uh oh!

pytorch-bot bot commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16718

❌ 3 New Failures, 2 Cancelled Jobs, 1 Unrelated Failure

Uh oh!

github-actions bot commented Jan 20, 2026

This PR needs a release notes: label

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

metascroy commented Mar 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

metascroy commented Jan 20, 2026 •

edited

Loading

pytorch-bot bot commented Jan 20, 2026 •

edited

Loading

This PR needs a `release notes:` label