Fix JIT build with old macOS SDK

[metascroy]

Summary

The NAX kernel 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 build fails because the JIT preamble generator preprocesses these headers and the include cannot be resolved.

Details

make_jit_source → make_compiled_preamble.sh runs xcrun -sdk macosx metal -x metal -E -H over each NAX header to enumerate/inline its dependencies. On an SDK without MetalPerformancePrimitives, that step fatals:

nax.h:12:10: fatal error: 'MetalPerformancePrimitives/MetalPerformancePrimitives.h' file not found
Error: Metal compiler header resolution failed for .../steel/gemm/gemm_nax.h

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 NAX make_jit_source entries unconditionally, so a JIT build on a pre-26 SDK always tries to preprocess nax.h and fails. make_compiled_preamble.sh already 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(...):

#if __has_include(<MetalPerformancePrimitives/MetalPerformancePrimitives.h>)
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
#endif

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=ON build with Xcode 16.4 / macOS 15.5 SDK: fails before, succeeds after.
• NAX-capable hardware + macOS 26 SDK: NAX still compiled and used (include present).

[jagrit06]

Thanks for the fix, I will find a pre-26 system as soon as I can and test it. Can you share what the jit pre-amble on the machine you tried in the meantime ?

[metascroy]

Thanks for the fix, I will find a pre-26 system as soon as I can and test it. Can you share what the jit pre-amble on the machine you tried in the meantime ?

Thanks @jagrit06! I'm also on macOS 26 locally, so the preamble just includes the NAX/MPP code as before for me.

But I did apply this exact patch in ExecuTorch because our test-mlx-backend-operators job runs on Xcode 15.4 / MacOSX14.5 SDK:

• Both nax.h patches apply cleanly
• The JIT preamble generation over the NAX headers runs and compiles without error (see jit/steel_gemm_fused_nax.cpp, steel_attention_nax.cpp in logs); previously this step failed on MLX v32.0
• The operator suite passes (390 passed, 77 skipped, 0 failed)

Full log: https://ossci-raw-job-status.s3.amazonaws.com/log/pytorch/executorch/89566514106

[zcbenz]

The correct way to fix this is to only make_jit_source the nax headers when the SDK requirement is met.

[metascroy]

The correct way to fix this is to only make_jit_source the nax headers when the SDK requirement is met.

Thanks @zcbenz, I have updated the approach as suggested!

[zcbenz]

Looks good to me, thanks!