Skip to content

Commit d6a13e0

Browse files
metascroymetascroy
authored
MLX docs (pytorch#18845)
This PR adds MLX docs
1 parent baafd7e commit d6a13e0

8 files changed

Lines changed: 406 additions & 2 deletions

File tree

‎backends/mlx/pte_inspector.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -788,9 +788,9 @@ def main(): # noqa: C901
788788
parser.add_argument(
789789
"--delegate-index",
790790
type=int,
791-
default=None,
791+
default=0,
792792
metavar="N",
793-
help="Index of delegate to extract (0-based). If not specified, extracts first matching delegate.",
793+
help="Index of delegate to extract (0-based, default: 0).",
794794
)
795795
parser.add_argument(
796796
"--parse-mlx",

‎docs/source/backends-overview.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ Backends are the bridge between your exported model and the hardware it runs on.
2323
| [XNNPACK](backends/xnnpack/xnnpack-overview.md) | All | CPU | General-purpose, fallback |
2424
| [CUDA](backends/cuda/cuda-overview.md) | Linux/Windows | GPU | NVIDIA GPU acceleration |
2525
| [Core ML](backends/coreml/coreml-overview.md) | iOS, macOS | NPU/GPU/CPU | Apple devices, high performance |
26+
| [MLX](/backends/mlx/mlx-overview.md) | macOS | GPU | Apple Silicon GPU (MLX) |
2627
| [Metal Performance Shaders](backends/mps/mps-overview.md) | iOS, macOS | GPU | Apple GPU acceleration |
2728
| [Vulkan](backends/vulkan/vulkan-overview.md) | Android, Linux, Windows | GPU | Android devices (mature); Desktops (experimental) |
2829
| [Qualcomm](backends-qualcomm) | Android | NPU | Qualcomm SoCs |
@@ -55,6 +56,7 @@ Backends are the bridge between your exported model and the hardware it runs on.
5556
backends/xnnpack/xnnpack-overview
5657
backends/cuda/cuda-overview
5758
backends/coreml/coreml-overview
59+
backends/mlx/mlx-overview
5860
backends/mps/mps-overview
5961
backends/vulkan/vulkan-overview
6062
backends-qualcomm

‎docs/source/backends/mlx/mlx-op-support.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
# Op Support
2+
3+
The MLX backend supports ~90 ATen operators plus multi-node fused patterns and custom ops. The partitioner automatically determines which ops in your model can be delegated to MLX. Unsupported ops fall back to ExecuTorch's portable CPU runtime.
4+
5+
For the current list of supported operators and fused patterns, see the source:
6+
7+
- **[ops.py](https://github.com/pytorch/executorch/blob/main/backends/mlx/ops.py)** — Single-op handlers (ATen op → MLX IR node)
8+
- **[patterns.py](https://github.com/pytorch/executorch/blob/main/backends/mlx/patterns.py)** — Multi-node fused patterns (quantized linear, SDPA, KV cache, etc.)
9+
10+
During lowering, the MLX partitioner prints a summary of supported and unsupported ops so you can see which ones are delegated and which fall back to CPU.

‎docs/source/backends/mlx/mlx-overview.md‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
# MLX Backend
2+
3+
The MLX delegate is the ExecuTorch backend for Apple Silicon GPUs via the [MLX](https://github.com/ml-explore/mlx) framework. It compiles PyTorch models into a custom FlatBuffer bytecode format at export time and executes them using MLX GPU primitives at runtime.
4+
5+
::::{note}
6+
The MLX delegate is experimental and under active development.
7+
::::
8+
9+
## Features
10+
11+
- GPU acceleration on Apple Silicon (M1 and later) via MLX.
12+
- INT2/INT4/INT8 weight quantization via [TorchAO](https://github.com/pytorch/ao).
13+
- Dynamic shape support.
14+
- Mutable buffers for persistent state across inference calls (e.g., KV cache).
15+
- Zero-copy constant loading on unified memory.
16+
17+
## Target Requirements
18+
19+
- Apple Silicon Mac (M1 or later)
20+
- [macOS](https://developer.apple.com/macos) >= 14.0
21+
22+
## Development Requirements
23+
24+
- [macOS](https://developer.apple.com/macos) on Apple Silicon (M1 or later)
25+
- [Xcode](https://developer.apple.com/xcode/) (full installation, not just Command Line Tools — the Metal compiler is required)
26+
27+
Verify the Metal compiler is available:
28+
29+
```bash
30+
xcrun -sdk macosx --find metal
31+
```
32+
33+
If this prints a path (e.g., `/Applications/Xcode.app/.../metal`), you're set. If it errors, install Xcode from [developer.apple.com](https://developer.apple.com/xcode/), then switch the active developer directory:
34+
35+
```bash
36+
sudo xcode-select -s /Applications/Xcode.app/Contents/Developer
37+
```
38+
39+
----
40+
41+
## Using the MLX Backend
42+
43+
To target the MLX backend during export and lowering, pass an instance of `MLXPartitioner` to `to_edge_transform_and_lower`. The MLX backend also provides a set of graph optimization passes via `get_default_passes()` that should be passed as `transform_passes`. The example below demonstrates this process using MobileNet V2:
44+
45+
```python
46+
import torch
47+
import torchvision.models as models
48+
from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
49+
from executorch.backends.mlx import MLXPartitioner
50+
from executorch.backends.mlx.passes import get_default_passes
51+
from executorch.exir import to_edge_transform_and_lower
52+
53+
mobilenet_v2 = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).eval()
54+
sample_inputs = (torch.randn(1, 3, 224, 224), )
55+
56+
et_program = to_edge_transform_and_lower(
57+
torch.export.export(mobilenet_v2, sample_inputs),
58+
transform_passes=get_default_passes(),
59+
partitioner=[MLXPartitioner()],
60+
).to_executorch()
61+
62+
with open("mv2_mlx.pte", "wb") as file:
63+
et_program.write_to_file(file)
64+
```
65+
66+
`get_default_passes()` includes RMSNorm fusion, consecutive view/permute/dtype-cast collapsing, no-op removal, and common subexpression elimination. These are recommended for all models and required for optimal LLM performance.
67+
68+
::::{note}
69+
The MLX backend is primarily designed for LLM and generative AI workloads on Apple Silicon. The MobileNet V2 example above is shown for simplicity, but in practice you would use this backend for models like Llama, Whisper, and other transformer-based architectures. See [LLM example](https://github.com/pytorch/executorch/tree/main/backends/mlx/examples/llm) for a more representative use case.
70+
::::
71+
72+
See [Partitioner API](mlx-partitioner.md) for a reference on available partitioner options.
73+
74+
----
75+
76+
## Quantization
77+
78+
The MLX backend supports INT4, INT8, and NVFP4 weight quantization via TorchAO for both linear and embedding layers. This is particularly useful for LLM inference. See [MLX Quantization](mlx-quantization.md) for details.
79+
80+
----
81+
82+
## Runtime Integration
83+
84+
### Python (pybindings)
85+
86+
The simplest way to get started is to install ExecuTorch with Python bindings. From the repo root:
87+
88+
```bash
89+
python install_executorch.py
90+
```
91+
92+
On Apple Silicon, when the Metal compiler is available, the MLX backend is automatically included. You can then export models in Python using the MLX partitioner and run them via the ExecuTorch Python API.
93+
94+
### C++ (CMake preset)
95+
96+
To build the C++ runtime with the MLX delegate, use the `mlx-release` CMake workflow preset from the repo root:
97+
98+
```bash
99+
cmake --workflow --preset mlx-release
100+
```
101+
102+
This configures and builds a Release build of the ExecuTorch runtime with the MLX delegate and installs artifacts into `cmake-out/`. The preset enables the MLX delegate along with commonly needed extensions (module, data loader, flat tensor, LLM runner, etc.).
103+
104+
Downstream C++ apps can then `find_package(executorch)` and link against `mlxdelegate` and `mlx`. The `executorch_target_link_options_shared_lib` utility handles whole-archive linkage (required for static initializer registration) cross-platform, and `executorch_target_copy_mlx_metallib` copies the Metal kernel library next to the binary so MLX can find it at runtime:
105+
106+
```cmake
107+
# CMakeLists.txt
108+
find_package(executorch REQUIRED)
109+
110+
# Link MLX delegate (with whole-archive for static initializer registration)
111+
target_link_libraries(my_target PRIVATE mlxdelegate mlx)
112+
executorch_target_link_options_shared_lib(mlxdelegate)
113+
114+
# Copy mlx.metallib next to the binary for runtime
115+
executorch_target_copy_mlx_metallib(my_target)
116+
```
117+
118+
No additional steps are necessary to use the backend beyond linking the target. An MLX-delegated `.pte` file will automatically run on the registered backend.
119+
120+
There is also an `mlx-debug` preset useful during development:
121+
122+
```bash
123+
cmake --workflow --preset mlx-debug
124+
```
125+
126+
## Reference
127+
128+
**→{doc}`/backends/mlx/mlx-troubleshooting` — Debug common issues.**
129+
130+
**→{doc}`/backends/mlx/mlx-partitioner` — Partitioner options.**
131+
132+
**→{doc}`/backends/mlx/mlx-quantization` — Supported quantization schemes.**
133+
134+
**→{doc}`/backends/mlx/mlx-op-support` — Supported operators.**
135+
136+
```{toctree}
137+
:maxdepth: 2
138+
:hidden:
139+
:caption: MLX Backend
140+
mlx-troubleshooting
141+
mlx-partitioner
142+
mlx-quantization
143+
mlx-op-support
144+
```

‎docs/source/backends/mlx/mlx-partitioner.md‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Partitioner API
2+
3+
The MLX partitioner API allows for configuration of model delegation to the MLX backend. Passing an `MLXPartitioner` instance with no additional parameters will run as much of the model as possible on the MLX backend with default settings. This is the most common use case.
4+
5+
## Usage
6+
7+
```python
8+
import torch
9+
from executorch.backends.mlx import MLXPartitioner
10+
from executorch.exir import to_edge_transform_and_lower
11+
12+
et_program = to_edge_transform_and_lower(
13+
torch.export.export(model, example_inputs),
14+
partitioner=[MLXPartitioner()],
15+
).to_executorch()
16+
```
17+
18+
::::{important}
19+
`MLXPartitioner` must be used with `to_edge_transform_and_lower()`. The legacy `to_edge()` + `to_backend()` workflow is **not supported** because it decomposes ops that MLX has optimized implementations for.
20+
::::
21+
22+
## Unsupported Op Logging
23+
24+
During partitioning, the partitioner logs a summary of any unsupported ops. This is useful for understanding what will fall back to CPU:
25+
26+
```
27+
================================================================================
28+
MLX Partitioner: UNSUPPORTED OPS SUMMARY
29+
================================================================================
30+
[UNSUPPORTED x2] aten.some_op.default
31+
Reason: No handler registered
32+
================================================================================
33+
```
34+
35+
If all ops are supported, you'll see:
36+
37+
```
38+
(All call_function nodes are supported!)
39+
```
40+
41+
Set `ET_MLX_DEBUG=1` to see detailed per-node support decisions during partitioning.

‎docs/source/backends/mlx/mlx-quantization.md‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
# Quantization
2+
3+
The MLX backend supports weight-only quantization via [TorchAO](https://github.com/pytorch/ao) for reducing model size and improving inference performance, particularly for LLMs on Apple Silicon. Quantization is applied to the eager model in-place **before** `torch.export()`.
4+
5+
## `quantize_`
6+
7+
The MLX backend uses TorchAO's [`quantize_`](https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_api.py) API under the hood. You can call it directly for full control over quantization configs and granularity. The key TorchAO configs are:
8+
9+
- [`IntxWeightOnlyConfig`](https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_api.py) — for INT2/INT4/INT8 weight-only quantization with per-group granularity (group sizes 32, 64, 128)
10+
- [`ExportableNVFP4Config`](https://github.com/pytorch/executorch/blob/main/extension/llm/export/nvfp4.py) — for NVFP4 weight-only quantization
11+
12+
```python
13+
import torch
14+
from torchao.quantization.quant_api import quantize_, IntxWeightOnlyConfig
15+
from torchao.quantization.granularity import PerGroup
16+
17+
# INT4 weight-only quantization for linear layers (group_size=32)
18+
quantize_(
19+
model,
20+
IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)),
21+
filter_fn=lambda m, fqn: isinstance(m, torch.nn.Linear),
22+
)
23+
24+
# INT8 weight-only quantization for embedding layers (group_size=128)
25+
quantize_(
26+
model,
27+
IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerGroup(128)),
28+
filter_fn=lambda m, fqn: isinstance(m, torch.nn.Embedding),
29+
)
30+
```
31+
32+
## `quantize_model_`
33+
34+
For convenience, ExecuTorch provides `quantize_model_` which wraps `quantize_` with sensible defaults for common LLM quantization configurations:
35+
36+
```python
37+
from executorch.extension.llm.export.quantize import quantize_model_
38+
39+
# Quantize linear layers with INT4, embedding layers with INT8
40+
# Note: 8w defaults to per-axis grouping, which MLX does not support.
41+
# Always pass an explicit group size when using 8w with MLX.
42+
quantize_model_(model, qlinear_config="4w", qembedding_config="8w", qembedding_group_size=128)
43+
```
44+
45+
### Supported configs
46+
47+
| Config | Description |
48+
|--------|-------------|
49+
| `4w` | INT4 weight-only quantization (per-group) |
50+
| `8w` | INT8 weight-only quantization (per-group) |
51+
| `nvfp4` | NVIDIA FP4 weight-only quantization |
52+
53+
These can be applied independently to linear layers and embedding layers.
54+
55+
### Using the LLM Export Script
56+
57+
The simplest way to export a quantized model is via the `export_llm_hf` script, which calls `quantize_model_` internally:
58+
59+
```bash
60+
# INT4 quantization for both linear and embedding layers
61+
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
62+
--model-id "unsloth/Llama-3.2-1B-Instruct" \
63+
--output llama_int4.pte \
64+
--use-custom-sdpa \
65+
--use-custom-kv-cache \
66+
--qlinear 4w \
67+
--qembedding 4w
68+
69+
# INT8 quantization for linear layers only
70+
# Note: --qlinear-group-size is required for 8w (default is per-axis, which MLX does not support)
71+
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
72+
--model-id "unsloth/Llama-3.2-1B-Instruct" \
73+
--output llama_int8.pte \
74+
--use-custom-sdpa \
75+
--use-custom-kv-cache \
76+
--qlinear 8w \
77+
--qlinear-group-size 128
78+
```
79+
80+
### CLI Quantization Options
81+
82+
| Option | Default | Description |
83+
|--------|---------|-------------|
84+
| `--qlinear` | None | Quantization for linear layers (`4w`, `8w`, `nvfp4`) |
85+
| `--qembedding` | None | Quantization for embedding layers (`4w`, `8w`, `nvfp4`) |
86+
| `--qlinear-group-size` | Depends on config | Group size for linear layer quantization (32, 64, or 128). Defaults to 32 for `4w`, 16 for `nvfp4`. **Required for `8w`** (default is per-axis, which MLX does not support). |
87+
| `--qembedding-group-size` | Depends on config | Group size for embedding layer quantization (32, 64, or 128). Defaults to 32 for `4w`, 16 for `nvfp4`. **Required for `8w`** (default is per-axis, which MLX does not support). |
88+
| `--no-tie-word-embeddings` | False | Disable re-tying lm_head to embedding after quantization |

0 commit comments

Comments
 (0)