epic: transferable low-precision performance for the OpenXLA backend (#449)
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## 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
- [x] #514
### Phase 2
- [ ] #515 (depends on #514)
### Phase 3
- [ ] #516 (depends on #514, #515)
### Phase 4
- [ ] #517 (depends on #514, #515, #516)
Contributor guide
Research direction
Start with the OpenXLA emitter, loader, and normal load/compile/serve path, then review the dependent sub-issues #514–#517. Done means low-precision graphs are integrated by default, the test suite has a precision accuracy gate, speedups are measured on Metal and CPU, and the int8 path and documentation are complete.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100