From c06259ea43cdf8d9fa5d3e780d41892dd296e06e Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni.hannun@gmail.com>
Date: Fri, 9 Jan 2026 14:06:51 -0800
Subject: [PATCH] Don't try to use nax at run-time if kernels aren't there

---
 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 564d15a9b4..a433b80491 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 262b049568..eb4a7b9185 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()