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> 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()