Summary
Give the OpenXLA emitter and loader a targeted low-precision mode: emit f16 (and optionally bf16) for the matmul / contraction path and keep the resident weights in that precision, instead of f32. Targeted, not blanket: a blunt global f32 to f16 demotion measured slower (norm / softmax / accumulation in f16 backfired), so only the contraction inputs are demoted while accumulation and the sensitive elementwise ops stay f32.
Measured win (M1 Ultra, decode step): ~2.06x on Metal, ~1.25x on CPU, end-to-end ~2.1x with byte-identical greedy output on real weights.
Scope
- Emitter (
src/lib/mlxcel-xla/src/emitter/): a precision parameter that emits the contraction inputs in f16/bf16 with f32 accumulation (equivalent to IREE --iree-global-opt-demote-contraction-inputs-type=f16, but authored in the graph so it is target-independent).
- Loader (
iree.rs): keep the resident weights in the chosen precision rather than widening bf16/f16 to f32, for the memory-bandwidth win.
- A selection knob (env and/or config) for the precision; default unchanged in this issue (flip happens in the gated follow-up).
Integration (required for done)
The engine compiles and runs the f16 graphs through the normal IreeLlama::load -> compile -> prefill/decode path, selectable at runtime; not a standalone emitter function. A real mlxcel generate run on Metal uses the f16 path when selected.
Validation
- Token stream on real weights compared to the f32 path (expect token-exact or within the accuracy gate from the sibling issue).
- tok/s reported on Metal and CPU.
Part of #513
Summary
Give the OpenXLA emitter and loader a targeted low-precision mode: emit f16 (and optionally bf16) for the matmul / contraction path and keep the resident weights in that precision, instead of f32. Targeted, not blanket: a blunt global f32 to f16 demotion measured slower (norm / softmax / accumulation in f16 backfired), so only the contraction inputs are demoted while accumulation and the sensitive elementwise ops stay f32.
Measured win (M1 Ultra, decode step): ~2.06x on Metal, ~1.25x on CPU, end-to-end ~2.1x with byte-identical greedy output on real weights.
Scope
src/lib/mlxcel-xla/src/emitter/): a precision parameter that emits the contraction inputs in f16/bf16 with f32 accumulation (equivalent to IREE--iree-global-opt-demote-contraction-inputs-type=f16, but authored in the graph so it is target-independent).iree.rs): keep the resident weights in the chosen precision rather than widening bf16/f16 to f32, for the memory-bandwidth win.Integration (required for done)
The engine compiles and runs the f16 graphs through the normal
IreeLlama::load-> compile -> prefill/decode path, selectable at runtime; not a standalone emitter function. A realmlxcel generaterun on Metal uses the f16 path when selected.Validation
Part of #513