This example demonstrates how to export and run LLMs using the MLX delegate for Apple Silicon.
- Export: Convert HuggingFace LLMs to ExecuTorch format with MLX delegate
- Quantization: Optional INT4/INT8 weight quantization via TorchAO
- KV Cache: Efficient KV cache implementation for autoregressive generation
- Off-graph KV Cache: Optional runtime-owned cache, sized and configured per run instead of at export
- Custom Ops: Uses
mlx::custom_sdpaandmlx::kv_cache_updatefor optimal execution on MLX - Pybindings: Run inference using ExecuTorch Python bindings
- Gemma 4: Text-only export and run flow supports processor-backed checkpoints such as
google/gemma-4-E2B-it
pip install transformers optimum-executorch| Script | Description |
|---|---|
export_llm_hf |
Export LLMs using optimum-executorch pipeline, with optional custom MLX SDPA/KV cache |
run_llm_hf |
Run exported models with token-by-token generation |
run_llm_hf.cpp |
C++ runner; the run path for off-graph-cache exports, and runs in-graph ones too |
run_llm_batched.cpp |
C++ runner for independent off-graph sessions, with optional shared prefix caching |
For exporting via the ExecuTorch LLM pipeline (e.g. examples/models/llama), use --mlx to enable the MLX delegate.
Uses optimum-executorch's CausalLMExportableModule by default. Optional flags enable custom MLX-optimized components (custom SDPA and/or KV cache).
# Baseline export using optimum-executorch
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
--model-id "unsloth/Llama-3.2-1B-Instruct" \
--output llama_hf.pte
# With custom MLX components
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
--model-id "unsloth/Llama-3.2-1B-Instruct" \
--output llama_hf_mlx.pte \
--use-custom-sdpa \
--use-custom-kv-cache
# With 4-bit quantization
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
--model-id "unsloth/Llama-3.2-1B-Instruct" \
--output llama_hf_int4.pte \
--use-custom-sdpa \
--use-custom-kv-cache \
--qlinear 4w \
--qembedding 4w
# Off-graph KV cache: the cache is owned by the runtime, not baked into the .pte
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
--model-id "unsloth/gemma-3-1b-it" \
--output gemma3_offgraph.pte \
--use-offgraph-cache \
--max-ctx-len 1024
# Gemma 4 text-only export
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
--model-id "google/gemma-4-E2B-it" \
--output gemma4_hf_int4.pte \
--use-custom-sdpa \
--use-custom-kv-cache \
--qlinear 4wGemma 4 support is currently validated for the text-only path using
--use-custom-sdpa --use-custom-kv-cache --qlinear 4w.
Validated with transformers commit
61461a7bcb458db7cf6eeea49678b9ab776a7821:
pip install -U "transformers @ git+https://github.com/huggingface/transformers.git@61461a7bcb458db7cf6eeea49678b9ab776a7821"| Option | Default | Description |
|---|---|---|
--model-id |
unsloth/Llama-3.2-1B-Instruct |
HuggingFace model ID |
--output |
(required) | Output .pte file path |
--max-ctx-len |
1024 |
Maximum context length / KV cache capacity |
--dtype |
bf16 |
Model dtype (fp32, fp16, bf16) |
--qlinear |
None | Quantization for linear layers (4w, 8w, nvfp4) |
--qembedding |
None | Quantization for embedding layers (4w, 8w, nvfp4) |
--no-tie-word-embeddings |
False |
Disable re-tying lm_head to embedding after quantization |
--use-custom-sdpa |
False |
Use MLX custom SDPA (mlx::custom_sdpa) |
--use-custom-kv-cache |
False |
Use MLX custom KV cache (mlx::kv_cache_update) |
--use-offgraph-cache |
False |
Use the off-graph KV cache (kvcache::update_and_attend); replaces the two flags above |
--prefill-chunk-size |
512 |
Max tokens per forward step. Bounds the traced seq_len dimension and is published as get_max_seq_len for the runner. It is also the largest single cache write, so a ring layer is sized window + chunk - 1; it may not exceed the sliding window or the context length. Ignored on the optimum-executorch path, which owns its own seq_len bound |
Off-graph exports keep no cache in the .pte, so the pybindings run_llm_hf
cannot run them — use mlx_run_llm_hf below, which builds
the cache and binds it at load time.
Run models exported with export_llm_hf. Supports both full-prompt prefill (dynamic seq len exports) and token-by-token prefill (fixed seq len exports).
python -m executorch.backends.mlx.examples.llm.run_llm_hf \
--pte llama_hf.pte \
--model-id unsloth/Llama-3.2-1B-Instruct \
--prompt "Explain quantum computing in simple terms"Gemma 4 checkpoints may use AutoProcessor instead of AutoTokenizer; run_llm_hf now supports both paths automatically for text-only prompts.
Validated Gemma 4 run command:
python -m executorch.backends.mlx.examples.llm.run_llm_hf \
--pte gemma4_hf_int4.pte \
--model-id google/gemma-4-E2B-it \
--prompt "What is the capital of France?" \
--max-new-tokens 50| Option | Default | Description |
|---|---|---|
--pte |
llama_hf.pte |
Path to .pte file |
--model-id |
unsloth/Llama-3.2-1B-Instruct |
HuggingFace model ID (for tokenizer or processor) |
--prompt |
The quick brown fox |
Input prompt |
--max-new-tokens |
50 |
Maximum tokens to generate |
Native runner for models exported with --use-offgraph-cache. That cache lives
outside the graph and is handed to the backend by key at load time, which the
pybindings run_llm_hf above cannot do. It also runs in-graph .pte files
unchanged — omit --kv-max-capacity — so the same binary compares both cache
paths. Greedy decode.
A standalone find_package(executorch) project, so ExecuTorch must be installed
first:
cmake --preset mlx-release
cmake --build cmake-out --target install -j$(( $(sysctl -n hw.ncpu) - 1 ))
cmake -S backends/mlx/examples/llm -B cmake-out/backends/mlx/examples/llm \
-DCMAKE_BUILD_TYPE=Release
cmake --build cmake-out/backends/mlx/examples/llm -j$(( $(sysctl -n hw.ncpu) - 1 ))The binary lands at cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf.
The same build also produces mlx_run_llm_batched in that directory.
mlx_run_llm_batched uses the backend-independent prefix cache in
extension/llm.
Enable it with --prefix_cache_entries; zero, the default, disables it.
Cold/warm logit and token parity are not guaranteed on MLX. Prefix reuse
changes prefill shapes and can change numerical results. Keep reuse opt-in and
restrict experiments to greedy generation until non-greedy sampling is
validated; greedy sampling itself does not guarantee matching tokens.
The application owns the shared batching::PrefixCache and captures snapshots
with capture_prompt(session, full_prompt_tokens). Each capture wraps one
generation callback and is collected on the caller thread after
generation.wait(). Capture binds to the original session identity, so the
source Session and capture handle can move. Keep the source's current owner
alive until generation completes; capture does not delay source closure.
Capture is best effort and requests its clone before delivering the
first tokens; user callback exceptions retain the Runner's normal failure
behavior. ModuleExecutor has no prefix-policy hooks.
--prompt_rounds 2 repeats the prompt on a fresh session while retaining the
cache. Use a model exported with --use-offgraph-cache, full or selected
logits, and a matching tokenizer and chat template.
After the Gemma 3 off-graph export below, run from the ExecuTorch root, replacing the tokenizer path with the tokenizer used for that export:
ET_PREFIX_TOKENIZER=/path/to/gemma3/tokenizer.json
ET_PREFIX_RESULTS=$(mktemp -d)
for entries in 0 2; do
cmake-out/backends/mlx/examples/llm/mlx_run_llm_batched \
--pte gemma3_offgraph.pte --tokenizer "$ET_PREFIX_TOKENIZER" \
--chat gemma --cache_kind batched-sequence \
--max_session_tokens 1024 --max_new_tokens 64 \
--max_decode_sequences 1 --temperature 0 --seed 42 \
--prefix_cache_entries "$entries" --prompt_rounds 2 \
--out_prefix "$ET_PREFIX_RESULTS/cache-$entries" \
"What is the capital of France?" \
> "$ET_PREFIX_RESULTS/cache-$entries.log" 2>&1 || exit 1
done
cmp "$ET_PREFIX_RESULTS/cache-0_0.txt" "$ET_PREFIX_RESULTS/cache-0_1.txt"
cmp "$ET_PREFIX_RESULTS/cache-0_0.txt" "$ET_PREFIX_RESULTS/cache-2_0.txt"
cmp "$ET_PREFIX_RESULTS/cache-0_0.txt" "$ET_PREFIX_RESULTS/cache-2_1.txt"
cat "$ET_PREFIX_RESULTS/cache-2.log"The generated text should agree across all four runs. In the cached run's
request report, generation 0 should have zero cached prompt tokens and generation 1
should reuse the prompt except for its final token, which is replayed for
fresh logits. Repeat with --cache_kind batched-cell to check the alternative
storage implementation. These short prompts check integration; the cache
tests additionally cover sliding-window wrap, divergent continuations, and
unavailable history. Model validation should also include prompts longer than
the model's sliding window.
python -m executorch.backends.mlx.examples.llm.export_llm_hf \
--model-id unsloth/gemma-3-1b-it \
--output gemma3_offgraph.pte \
--use-offgraph-cache \
--max-ctx-len 1024
cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf \
--pte gemma3_offgraph.pte \
--tokenizer ~/.cache/huggingface/hub/models--unsloth--gemma-3-1b-it/snapshots/*/tokenizer.json \
--chat gemma \
--kv-max-capacity 1024 \
--prompt "What is the capital of France?" \
--max-new-tokens 50--chat selects the instruct template — llama3, gemma, gemma4, or 0 for
raw text. It matters: raw text confuses an instruct model into emitting turn
markers, and using the wrong template invalidates a comparison between two
.pte files.
Only the cache geometry is fixed at export — how many caches, their KV heads, head dims and windows, which the export publishes as constant methods the runner reads before building the cache. Everything else is chosen per run, with no re-export. The runner reports what it built:
[cache] off-graph seq | capacity=1024 initial=512 kv_dtype=BFloat16
26 layers: 4 flat + 22 ring(window 512)
--kv-windows overrides the attention pattern, repeating a comma-separated list
over the caches (0 = flat). The geometry each cache declared is untouched, so
this cannot desync from the graph. For gemma-3-1b, whose 26 layers are 22
sliding at 512 and 4 full:
# (omitted) 26 layers: 4 flat + 22 ring(window 512)
--kv-windows 512,512,512,512,512,0 # the same, spelling out gemma-3's 5:1 period
--kv-windows 0 # 26 layers: 26 flat
--kv-windows 512 # 26 layers: 0 flat + 26 ring(window 512)
--kv-windows 512,256 # 26 layers: 0 flat + 13 ring(256) + 13 ring(512)A ring layer allocates its whole window + chunk - 1 slots up front, while a
flat layer starts at --kv-initial-capacity and doubles as the sequence grows,
so an all-flat cache can look smaller than a sliding one early in a run and
larger later. Prefill runs in steps of the chunk size the export published, so
the ring is bounded by the chunk rather than by the prompt: prefilling 24k
tokens in one step would need 541 MiB of ring here, against 24 MiB at the 512
default.
--help lists every flag with its default.
| Option | Default | Description |
|---|---|---|
--pte |
(required) | Path to .pte file |
--tokenizer |
(required) | Path to tokenizer.json |
--prompt |
The quick brown fox |
Input prompt |
--max-new-tokens |
50 |
Tokens to generate, excluding the prompt |
--temperature |
0 |
Sampling temperature; 0 is greedy argmax, which is what makes two .pte files comparable |
--chat |
llama3 |
Chat template: llama3, gemma, gemma4, or 0 to disable |
--kv-max-capacity |
0 |
Off-graph: history the cache may hold. Setting it selects the off-graph path |
--kv-storage-dtype |
PTE activation dtype | Off-graph: optional KV storage override (bf16, fp16, fp32); defaults to the PTE's get_activation_dtype, which is required, so a .pte exported before this metadata must be re-exported |
--kv-initial-capacity |
-1 |
Off-graph: starting pool size; grows by doubling up to capacity |
--kv-windows |
(model's own) | Off-graph: attention pattern override, e.g. 512 |
--interactive |
false |
Multi-turn chat on stdin; off-graph only |
--warmup |
false |
Run once before measuring, to absorb JIT and pool growth |
The export_llm_hf script uses optimum-executorch's CausalLMExportableModule by default. When custom flags are enabled, it uses TorchExportableModuleWithStaticCache from HuggingFace transformers, with optional MLX-specific replacements:
--use-custom-sdpa: Registersmlx::custom_sdpaattention implementation--use-custom-kv-cache: Replaces HF'sStaticCachewithHFStaticCacheusingmlx::kv_cache_update