link with nccl (#2546)

mlx/distributed/nccl/CMakeLists.txt

@@ -9,11 +9,17 @@ if(MLX_BUILD_CUDA)
       STATUS
         "NCCL not found, using stubs. To run distributed with NCCL backend, install NCCL."
     )
-    file(
+
-      DOWNLOAD
+    include(ExternalProject)
-      "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in"
+    ExternalProject_Add(
-      "${CMAKE_CURRENT_BINARY_DIR}/nccl.h")
+      nccl_stub
-    target_include_directories(mlx PRIVATE "${CMAKE_CURRENT_BINARY_DIR}")
+      SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/nccl_stub"
+      BUILD_COMMAND ${CMAKE_COMMAND} --build .
+      INSTALL_COMMAND "")
+    set(NCCL_PATH
+        "${CMAKE_CURRENT_BINARY_DIR}/nccl_stub-prefix/src/nccl_stub-build/")
+    target_link_libraries(mlx PRIVATE ${NCCL_PATH}/libnccl.so)
+    target_include_directories(mlx PRIVATE ${NCCL_PATH})
   endif()
 else()
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp)

mlx/distributed/nccl/nccl_stub/CMakeLists.txt

@@ -0,0 +1,13 @@
+cmake_minimum_required(VERSION 3.25)
+
+project(nccl LANGUAGES C CXX)
+
+file(
+  DOWNLOAD
+  "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in"
+  "${CMAKE_CURRENT_BINARY_DIR}/nccl.h")
+
+add_library(nccl SHARED nccl_stubs.cpp)
+find_package(CUDAToolkit REQUIRED)
+target_include_directories(nccl PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+target_include_directories(nccl PRIVATE ${CMAKE_CURRENT_BINARY_DIR})

mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp

@@ -0,0 +1,7 @@
+#include <arpa/inet.h>
+#include <cuda_runtime.h>
+#include <nccl.h>
+
+ncclResult_t ncclGetUniqueId(ncclUniqueId*) {
+  return ncclSuccess;
+}