Context
The OpenXLA / StableHLO (xla-iree) backend (#449) runs the same StableHLO graph on any IREE target (CPU, CUDA, Metal, and future NPUs via IREE HAL drivers). It is a portability / parity path, not a hand-tuned engine. Today its bundled graphs and resident weights are f32, which leaves large, transferable performance on the table.
Measured on an M1 Ultra (Llama-3.2-1B-Instruct, greedy, one decode step via iree-benchmark-module, pure runtime so no host glue):
| variant |
Metal |
CPU (13 threads) |
| f32 (today) |
~600 ms |
~233 ms |
| f16 matmul inputs |
~291 ms (~2.06x) |
~187 ms (~1.25x) |
End-to-end on real weights (48 tokens, Metal): f32 1.57 tok/s vs f16 3.33 tok/s (~2.1x), byte-identical output (the f16 matmul inputs with f32 accumulation preserved the greedy token stream). For reference MLX runs the same model at ~186 tok/s.
Thesis: invest in the graph, not in Metal kernels
The performance levers split in two:
- Graph-level (precision, quantization, op selection) is authored once in the portable StableHLO graph and transfers to every IREE target, demonstrated above (one f16 change sped up both Metal and CPU with zero per-backend code). For NPUs this is not a 2x optimization, it is the entry ticket: NPUs are int8 / fp16 native, so a low-precision / quantized graph is what unlocks their systolic kernels. This is the investment worth making.
- Per-backend kernel codegen (tiling, simdgroup / MPS matmul, fusion) is where the remaining ~50x to MLX lives. That is upstream IREE's job, is Metal-specific (does not transfer to non-SPIR-V NPUs), and MLX already owns Apple Silicon performance. Not in scope.
Metal is therefore used here only as a convenient local proxy to validate correctness and relative speedup; absolute Metal tok/s is a pessimistic stand-in for an NPU with its own optimized kernels.
Scope
Land the transferable low-precision arc on the OpenXLA path: f16/bf16 first (proven, token-exact, ~2x), then int8 weight quantization (the NPU lever), each guarded by a precision accuracy gate. The OpenXLA emitter is already config-driven (Llama any size, Qwen2, Gemma2) and the loader already dequantizes MLX 4/8-bit checkpoints to f32, so the work extends existing paths rather than adding a new one.
Non-goals
- Not chasing MLX on Apple Silicon. MLX stays the default and primary Apple-Silicon backend; this is the portability path.
- No hand-written Metal kernels and no IREE Metal-codegen tuning (upstream / out of scope).
- No new model families (use the existing config-driven emitter coverage).
Epic acceptance (functionally integrated deliverable)
By epic close the OpenXLA backend, on a GPU target, runs low-precision graphs through its normal load / compile / serve path by default, with: a precision accuracy gate in the test suite; measured speedups across Metal and CPU; a working int8-quantized path for MLX-quantized checkpoints; and docs recording the perf table and the transferable-precision decision. Not standalone modules: each piece must be wired into the real generation path and the default behavior.
Sub-issues
Phase 1
Phase 2
Phase 3
Phase 4
Context
The OpenXLA / StableHLO (
xla-iree) backend (#449) runs the same StableHLO graph on any IREE target (CPU, CUDA, Metal, and future NPUs via IREE HAL drivers). It is a portability / parity path, not a hand-tuned engine. Today its bundled graphs and resident weights are f32, which leaves large, transferable performance on the table.Measured on an M1 Ultra (Llama-3.2-1B-Instruct, greedy, one decode step via
iree-benchmark-module, pure runtime so no host glue):End-to-end on real weights (48 tokens, Metal): f32 1.57 tok/s vs f16 3.33 tok/s (~2.1x), byte-identical output (the f16 matmul inputs with f32 accumulation preserved the greedy token stream). For reference MLX runs the same model at ~186 tok/s.
Thesis: invest in the graph, not in Metal kernels
The performance levers split in two:
Metal is therefore used here only as a convenient local proxy to validate correctness and relative speedup; absolute Metal tok/s is a pessimistic stand-in for an NPU with its own optimized kernels.
Scope
Land the transferable low-precision arc on the OpenXLA path: f16/bf16 first (proven, token-exact, ~2x), then int8 weight quantization (the NPU lever), each guarded by a precision accuracy gate. The OpenXLA emitter is already config-driven (Llama any size, Qwen2, Gemma2) and the loader already dequantizes MLX 4/8-bit checkpoints to f32, so the work extends existing paths rather than adding a new one.
Non-goals
Epic acceptance (functionally integrated deliverable)
By epic close the OpenXLA backend, on a GPU target, runs low-precision graphs through its normal load / compile / serve path by default, with: a precision accuracy gate in the test suite; measured speedups across Metal and CPU; a working int8-quantized path for MLX-quantized checkpoints; and docs recording the perf table and the transferable-precision decision. Not standalone modules: each piece must be wired into the real generation path and the default behavior.
Sub-issues
Phase 1
Phase 2
Phase 3
Phase 4