Guard NAX MetalPerformancePrimitives include behind __has_include#3853
Open
metascroy wants to merge 1 commit into
Open
Guard NAX MetalPerformancePrimitives include behind __has_include#3853metascroy wants to merge 1 commit into
metascroy wants to merge 1 commit into
Conversation
The NAX GEMM and attention headers unconditionally include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>, which only ships in the macOS 26 / Xcode 26 SDK. With MLX_METAL_JIT=ON on an older SDK, the JIT preamble generator (make_compiled_preamble.sh) runs `metal -E` over these headers and fatals on the missing include, breaking the build. The metallib path already gates NAX on the SDK/Metal version; the JIT path did not. Wrap the include in `#if __has_include(...)` so preprocessing succeeds on pre-26 SDKs (NAX kernels are runtime-gated via is_nax_available() and never compiled there), while newer SDKs still get NAX.
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
The NAX kernel headers unconditionally include
<MetalPerformancePrimitives/MetalPerformancePrimitives.h>, which only ships in the macOS 26 / Xcode 26 SDK. WithMLX_METAL_JIT=ONon an older SDK, the build fails because the JIT preamble generator preprocesses these headers and the include cannot be resolved.Details
make_jit_source→make_compiled_preamble.shrunsxcrun -sdk macosx metal -x metal -E -Hover each NAX header to enumerate/inline its dependencies. On an SDK withoutMetalPerformancePrimitives, that step fatals:The metallib path already gates NAX on
MLX_METAL_VERSION >= 400 && MACOS_SDK_VERSION >= 26.2 && CMAKE_OSX_DEPLOYMENT_TARGET >= 26.2(kernels/CMakeLists.txt), but the JIT path adds the NAXmake_jit_sourceentries unconditionally, so a JIT build on a pre-26 SDK always tries to preprocessnax.hand fails.make_compiled_preamble.shalready intends to tolerate the framework (grep -v "Xcode"strips it from the inlined preamble); this only breaks when the header is absent from the SDK entirely.Fix
Wrap the framework include in both NAX headers with
#if __has_include(...):On pre-26 SDKs, preprocessing now succeeds and NAX kernels are simply never JIT-compiled (they are runtime-gated by
is_nax_available(), which is already false on non-NAX hardware). On the macOS 26 SDK the include is kept and NAX is unchanged.Testing
MLX_METAL_JIT=ONbuild with Xcode 16.4 / macOS 15.5 SDK: fails before, succeeds after.