lablup / lablup/mlxcel

epic: transferable low-precision performance for the OpenXLA backend (#449)

Open
#513 1 comment 0 reactions 0 assignees View on GitHub
area:inference status:in-progress type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.