From 8654b8281df54bb5f7a5933a34fa3c2b7f334552 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Tue, 13 Jan 2026 15:47:45 -0800
Subject: [PATCH] Don't try to use NAX at run-time if kernels aren't there
 (#2982)

---
 mlx/backend/metal/device.h               | 4 ++++
 mlx/backend/metal/kernels/CMakeLists.txt | 2 ++
 2 files changed, 6 insertions(+)

diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h
index 564d15a9..a433b804 100644
--- a/mlx/backend/metal/device.h
+++ b/mlx/backend/metal/device.h
@@ -266,6 +266,9 @@ Device& device(mlx::core::Device);
 std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool();
 
 inline bool is_nax_available() {
+#ifdef MLX_METAL_NO_NAX
+  return false;
+#else
   auto _check_nax = []() {
     bool can_use_nax = false;
     if (__builtin_available(
@@ -278,6 +281,7 @@ inline bool is_nax_available() {
   };
   static bool is_nax_available_ = _check_nax();
   return is_nax_available_;
+#endif
 }
 
 } // namespace mlx::core::metal
diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt
index 262b0495..eb4a7b91 100644
--- a/mlx/backend/metal/kernels/CMakeLists.txt
+++ b/mlx/backend/metal/kernels/CMakeLists.txt
@@ -164,6 +164,8 @@ if(NOT MLX_METAL_JIT)
     build_kernel(steel/attn/kernels/steel_attention_nax
                  ${STEEL_NAX_ATTN_HEADERS})
 
+  else()
+    target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX)
   endif()
 
 endif()