feat: f16/bf16 precision mode for the OpenXLA emitter (#449)#553
Merged
Conversation
Add a Precision to the StableHLO emitter's Builder that demotes the f32 inputs of every dot_general (matmul) to f16/bf16 while keeping the f32 accumulate and output, so only the matmuls change and norm/softmax/RoPE stay f32. Authored in the graph so it is portable to every IREE target (CPU, CUDA, Metal, future NPUs), matching --iree-global-opt-demote-contraction-inputs-type=f16. Selected by MLXCEL_XLA_PRECISION (f16|bf16, default f32); the emitted MLIR differs, so the vmfb content-hash cache keys each precision separately. A blanket program-wide f32 to f16 is deliberately not done: it regressed norm/softmax/accumulation and was slower. Weights stay f32-resident and are demoted in the graph (no FFI change); keeping the resident weights in the narrow type (a memory-bandwidth win that needs the f16 weight FFI) lands with the quantized-weight path, where it matters more. Validated on an M1 Ultra (Llama-3.2-1B-Instruct, greedy): f16 is ~1.9x the f32 tok/s (2.89 vs 1.53) and token-exact with it; the same graph change also speeds up the CPU path. The F32 default is byte-identical to before (the emitter byte-exact regression test still passes). New builder unit tests cover the f16/bf16 demotion.
inureyes
added
type:enhancement
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #514. Part of #513.
What
A
Precisionon the OpenXLA StableHLO emitter'sBuilderthat demotes the f32 inputs of everydot_general(matmul) to f16/bf16 while keeping the f32 accumulate and output. Only the matmuls change; the sensitive elementwise ops (norm, softmax, RoPE) stay f32. Authored in the graph, so it is portable to every IREE target (CPU, CUDA, Metal, future NPUs), matching--iree-global-opt-demote-contraction-inputs-type=f16.Selected by
MLXCEL_XLA_PRECISION(f16|bf16, defaultf32). The emitted MLIR differs, socompile_one's content-hash vmfb cache keys each precision separately with no extra plumbing.Changes
emitter/builder.rs:Precisionenum +precision_from_env();Buildercarries it (defaultF32,with_precision());dot_generalconverts f32 operands to the narrow type (f32 output) when set. Unit tests for the f16/bf16 demotion.emitter/model.rs: the fouremit_*entry points readprecision_from_env().mlxcel-xla/README.md: a Precision section.Design notes
--iree-input-demote-f32-to-f16was measured slower (f16 norm/softmax/accumulation backfires), so only the contraction inputs are demoted.Validation (M1 Ultra, Llama-3.2-1B-Instruct, greedy)
MLXCEL_XLA_PRECISION=f16on Metal: 2.89 tok/s vs 1.53 (f32) = ~1.9x, output token-exact with f32.local-task) path too.from_json_reproduces_bundled_assets_byte_for_byte) still passes. 46/46mlxcel-xlaunit tests green;cargo fmt --checkclean; emitter clippy clean.xla-backend-only crate; the F32 path is a no-op passthrough).Note on performance
This is a transferable, correctness-first lever, not a path to MLX parity. XLA-on-Metal remains a dev/parity path (the remaining gap is IREE metal-spirv kernel codegen, out of scope). The value is that this one graph change benefits every IREE target, which is the entry ticket for low-precision-native NPUs.