Add Neural Accelerator Support (#2772)

Author: jagrit06

mlx/backend/metal/CMakeLists.txt

@@ -121,6 +121,14 @@ if(NOT MLX_METAL_PATH)
   set(MLX_METAL_PATH ${CMAKE_CURRENT_BINARY_DIR}/kernels/)
 endif()
 
+if((MLX_METAL_VERSION GREATER_EQUAL 400) AND (MACOS_SDK_VERSION GREATER_EQUAL
+                                              26.2))
+  set(MLX_ENABLE_NAX TRUE)
+  target_compile_definitions(mlx PRIVATE MLX_ENABLE_NAX)
+else()
+  set(MLX_ENABLE_NAX FALSE)
+endif()
+
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels)
 
 target_compile_definitions(mlx

mlx/backend/metal/device.h

@@ -265,4 +265,14 @@ Device& device(mlx::core::Device);
 
 std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool();
 
+#ifdef MLX_ENABLE_NAX
+
+inline bool is_nax_available() {
+  static bool is_nax_available_ =
+      metal::device(mlx::core::Device::gpu).get_architecture_gen() >= 17;
+  return is_nax_available_;
+}
+
+#endif // MLX_ENABLE_NAX
+
 } // namespace mlx::core::metal

mlx/backend/metal/kernels/CMakeLists.txt

@@ -9,10 +9,13 @@ set(BASE_HEADERS
     utils.h)
 
 function(build_kernel_base TARGET SRCFILE DEPS)
-  set(METAL_FLAGS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
+  set(METAL_FLAGS -x metal -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
   if(MLX_METAL_DEBUG)
     set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources)
   endif()
+  if(MLX_ENABLE_NAX)
+    set(METAL_FLAGS ${METAL_FLAGS} -Wno-c++20-extensions -std=metal4.0)
+  endif()
   if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
     set(METAL_FLAGS ${METAL_FLAGS}
                     "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
@@ -120,6 +123,30 @@ if(NOT MLX_METAL_JIT)
   build_kernel(gemv_masked steel/utils.h)
 endif()
 
+if(MLX_ENABLE_NAX)
+
+  set(STEEL_NAX_HEADERS
+      steel/defines.h
+      steel/utils.h
+      steel/gemm/transforms.h
+      steel/gemm/nax.h
+      steel/gemm/gemm_nax.h
+      steel/utils/type_traits.h
+      steel/utils/integral_constant.h)
+
+  build_kernel(steel/gemm/kernels/steel_gemm_fused_nax ${STEEL_NAX_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_gather_nax ${STEEL_NAX_HEADERS})
+
+  build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS})
+  build_kernel(fp_quantized_nax fp_quantized_nax.h ${STEEL_NAX_HEADERS})
+
+  set(STEEL_NAX_ATTN_HEADERS
+      steel/defines.h steel/utils.h steel/attn/nax.h steel/utils/type_traits.h
+      steel/utils/integral_constant.h)
+
+  build_kernel(steel/attn/kernels/steel_attention_nax ${STEEL_NAX_ATTN_HEADERS})
+endif()
+
 add_custom_command(
   OUTPUT ${MLX_METAL_PATH}/mlx.metallib
   COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o