Skip to content

Commit b4ffdbe

Browse files
metascroymetascroy
authored
Adjust how nax are excluded on < macOS 26.0 (#21554)
As titled
1 parent e54d663 commit b4ffdbe

3 files changed

Lines changed: 113 additions & 45 deletions

File tree

‎backends/mlx/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ message(
185185
# mismatch after an MLX bump fails loudly rather than silently no-op'ing. See
186186
# each patch file under patches/ for its rationale.
187187
set(_mlx_patches
188-
${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_has_include.patch
188+
${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_jit_sdk_gate.patch
189189
${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch
190190
${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_gather_mm_rhs_lda.patch
191191
)

‎backends/mlx/patches/mlx_nax_has_include.patch‎

Lines changed: 0 additions & 44 deletions
This file was deleted.

‎backends/mlx/patches/mlx_nax_jit_sdk_gate.patch‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
Gate the NAX JIT kernel sources behind the SDK requirement.
2+
3+
MLX's NAX kernels (GEMM and attention) include
4+
<MetalPerformancePrimitives/MetalPerformancePrimitives.h>, a framework that only
5+
ships in the macOS 26 / Xcode 26 SDK. With MLX_METAL_JIT=ON (which ExecuTorch
6+
uses) on an older SDK, the JIT preamble generator (make_compiled_preamble.sh)
7+
runs `metal -E` over these headers, which fatals on the missing include and
8+
fails the build. MLX already gates NAX on the non-JIT metallib path
9+
(kernels/CMakeLists.txt), but the JIT path was ungated.
10+
11+
Instead of guarding the includes with __has_include, gate the NAX
12+
make_jit_source() calls behind the same
13+
MLX_METAL_VERSION/MACOS_SDK_VERSION/CMAKE_OSX_DEPLOYMENT_TARGET check the
14+
metallib path uses, and define MLX_METAL_NO_NAX when the requirement is unmet.
15+
On those SDKs NAX is already runtime-gated via is_nax_available()
16+
(device.cpp), so the get_*_nax_kernel entry points in jit_kernels.cpp are
17+
unreachable; guarded empty preamble getters keep that translation unit linking.
18+
macOS 26+ SDKs still build NAX.
19+
20+
Upstream candidate; carried locally until an MLX release gates the JIT path.
21+
22+
diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt
23+
--- a/mlx/backend/metal/CMakeLists.txt
24+
+++ b/mlx/backend/metal/CMakeLists.txt
25+
@@ -84,19 +84,32 @@ if(MLX_METAL_JIT)
26+
27+
make_jit_source(steel/attn/kernels/steel_attention)
28+
29+
- make_jit_source(
30+
- steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h
31+
- kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h)
32+
- make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax)
33+
- make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax)
34+
- make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax)
35+
- make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax)
36+
+ if(MLX_METAL_VERSION GREATER_EQUAL 400
37+
+ AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2
38+
+ AND CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER_EQUAL 26.2)
39+
40+
- make_jit_source(quantized_nax kernels/quantized_utils.h)
41+
- make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h
42+
- kernels/fp4.h)
43+
+ make_jit_source(
44+
+ steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h
45+
+ kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h)
46+
+ make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax)
47+
+ make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax)
48+
+ make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax)
49+
+ make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax)
50+
+
51+
+ make_jit_source(quantized_nax kernels/quantized_utils.h)
52+
+ make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h
53+
+ kernels/fp4.h)
54+
+
55+
+ make_jit_source(steel/attn/kernels/steel_attention_nax)
56+
57+
- make_jit_source(steel/attn/kernels/steel_attention_nax)
58+
+ else()
59+
+ message(
60+
+ WARNING "NAX kernels require Metal 4, macOS SDK >= 26.2, and "
61+
+ "MACOSX_DEPLOYMENT_TARGET >= 26.2 (SDK ${MACOS_SDK_VERSION}, "
62+
+ "CMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}). "
63+
+ "Building without NAX kernels.")
64+
+ target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX)
65+
+ endif()
66+
67+
else()
68+
target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nojit_kernels.cpp)
69+
diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp
70+
--- a/mlx/backend/metal/jit_kernels.cpp
71+
+++ b/mlx/backend/metal/jit_kernels.cpp
72+
@@ -8,6 +8,40 @@ using namespace fmt::literals;
73+
74+
namespace mlx::core {
75+
76+
+#ifdef MLX_METAL_NO_NAX
77+
+// NAX JIT preambles are only generated (via make_jit_source) when the SDK
78+
+// requirement is met. On older SDKs they are skipped and MLX_METAL_NO_NAX is
79+
+// defined, so is_nax_available() returns false and the get_*_nax_kernel entry
80+
+// points below are never reached. These empty definitions only exist to satisfy
81+
+// the linker for this translation unit.
82+
+namespace metal {
83+
+const char* gemm_nax() {
84+
+ return "";
85+
+}
86+
+const char* steel_gemm_fused_nax() {
87+
+ return "";
88+
+}
89+
+const char* steel_gemm_gather_nax() {
90+
+ return "";
91+
+}
92+
+const char* steel_gemm_splitk_nax() {
93+
+ return "";
94+
+}
95+
+const char* steel_gemm_segmented_nax() {
96+
+ return "";
97+
+}
98+
+const char* quantized_nax() {
99+
+ return "";
100+
+}
101+
+const char* fp_quantized_nax() {
102+
+ return "";
103+
+}
104+
+const char* steel_attention_nax() {
105+
+ return "";
106+
+}
107+
+} // namespace metal
108+
+#endif // MLX_METAL_NO_NAX
109+
+
110+
MTL::ComputePipelineState* get_arange_kernel(
111+
metal::Device& d,
112+
const std::string& kernel_name,

0 commit comments

Comments
 (0)