Expose MLX JIT build flag (EMILY_MLX_JIT) — closes #45 - #47
Merged
Conversation
Pipe EMILY_MLX_JIT=1 through to -DMLX_METAL_JIT=ON in the MLX CMake
configure so users can opt in to MLX's runtime JIT-compilation of
Metal kernels. Default (unset / 0) preserves the existing AOT build.
The flag is incorporated into the MLX install-dir cache key
("mlx-<hash>" vs "mlx-<hash>-jit") so toggling does not reuse a stale
artefact. All 425 tests pass under both modes.
JIT on trades a 150 MB AOT metallib for a 3.5 MB stub plus lazy
per-kernel compile at first use; priv/ footprint drops from ~175 MB
to ~25 MB. README "How to build" documents the trade-off.
Closes #45.
Add a matrix dimension over EMILY_MLX_JIT ∈ {0, 1} so both MLX build
flavours are exercised on every push / PR. _build and MLX caches are
partitioned by the flag so the two matrix entries don't clobber each
other's artefacts.
The MLX JIT branch in backend/metal/CMakeLists.txt unconditionally runs make_jit_source on NAX headers, which #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h> — only present in the macOS 26.2+ SDK. The AOT branch in kernels/CMakeLists.txt already gates NAX with that SDK check; the JIT branch needs the same gate so JIT=1 builds work on older SDKs (e.g. the macos-14 GitHub runner). Rather than fork the submodule, ship a patch file under patches/ and apply it idempotently from mix.exs before cmake configure. When the SDK gate fails, MLX_METAL_NO_NAX is defined and is_nax_available() returns false at runtime, matching the existing AOT behaviour.
CI JIT=1 on macos-14 got past cmake configure but hit a linker
error:
symbol not found in flat namespace '__ZN3mlx4core5metal13quantized_naxEv'
jit_kernels.cpp references mlx::core::metal::{gemm_nax,
quantized_nax, fp_quantized_nax, steel_gemm_fused_nax, …} as strong
external symbols. Those symbols are normally supplied by the
auto-generated jit/<name>.cpp files that make_jit_source() creates;
when we gate those out for old SDKs, nothing provides them and the
NIF fails to link.
Extend the patch to emit an nax_stubs.cpp with empty-string returns
under the gated-out branch. The stubs are never exercised because
is_nax_available() returns false when MLX_METAL_NO_NAX is defined.
Owner
Author
|
A PR that applies the patch can be found at ml-explore/mlx#3426 |
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.
Summary
EMILY_MLX_JIT=1at build time passes-DMLX_METAL_JIT=ONto MLX's CMake, selecting runtime JIT compilation of Metal kernels. Default (unset /0) preserves the existing AOT path.mlx-<hash>vsmlx-<hash>-jit), so flipping it doesn't reuse a stale libmlx.a / metallib.Measured artefact sizes (M-series Mac, Release)
libemily.somlx.metallibpriv/totalEMILY_MLX_JIT=1)The ~150 MB delta on disk comes from
mlx.metallibshrinking to a stub when kernels are JIT-compiled lazily at first use instead of shipped AOT-compiled. Issue #45's framing of JIT as the "bloating" mode is inverted — AOT is the bloated mode. Default is unchanged for now (per the deferred-default point in #45).Test plan
EMILY_MLX_JIT=1 EMILY_CACHE=/tmp/… mix compile.emily_mlx --forcesucceeds; MLX install dir gets-jitsuffix.mix testpasses 425/425 under JIT off (default).mix testpasses 425/425 under JIT on (EMILY_MLX_JIT=1 EMILY_CACHE=… mix test).mix precommitclean on default config.ls $EMILY_CACHEshowsmlx-<hash>/andmlx-<hash>-jit/side by side.Notes
During investigation I initially observed 322/425 failures under JIT on — this turned out to be stale cached artefacts from iterative testing, not a real bug. Clean rebuild produced the passing run. No tests needed to be skipped or gated for JIT-only code paths.