diff --git a/CMakeLists.txt b/CMakeLists.txt index 6e1301cf..2e2c05da 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -155,6 +155,7 @@ target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS}) if(MLX_BUILD_CUDA) enable_language(CUDA) + find_package(CUDAToolkit REQUIRED) endif() if(MLX_BUILD_METAL) diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index abaca309..dde2507b 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -202,7 +202,6 @@ FetchContent_MakeAvailable(nvtx3) target_link_libraries(mlx PUBLIC $) # Make cuda runtime APIs available in non-cuda files. -find_package(CUDAToolkit REQUIRED) target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) # Use cublasLt. diff --git a/mlx/distributed/nccl/CMakeLists.txt b/mlx/distributed/nccl/CMakeLists.txt index 4c59c56e..644d060d 100644 --- a/mlx/distributed/nccl/CMakeLists.txt +++ b/mlx/distributed/nccl/CMakeLists.txt @@ -9,17 +9,17 @@ if(MLX_BUILD_CUDA) STATUS "NCCL not found, using stubs. To run distributed with NCCL backend, install NCCL." ) - - include(ExternalProject) - ExternalProject_Add( + 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_stub OBJECT + ${CMAKE_CURRENT_SOURCE_DIR}/nccl_stub/nccl_stubs.cpp) + target_include_directories( nccl_stub - 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}) + PRIVATE ${CUDAToolkit_INCLUDE_DIRS} + PUBLIC ${CMAKE_CURRENT_BINARY_DIR}) + target_link_libraries(mlx PRIVATE $) endif() else() target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp) diff --git a/mlx/distributed/nccl/nccl_stub/CMakeLists.txt b/mlx/distributed/nccl/nccl_stub/CMakeLists.txt index a93ba3c0..8b137891 100644 --- a/mlx/distributed/nccl/nccl_stub/CMakeLists.txt +++ b/mlx/distributed/nccl/nccl_stub/CMakeLists.txt @@ -1,14 +1 @@ -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) -set_target_properties(nccl PROPERTIES SOVERSION 2) -find_package(CUDAToolkit REQUIRED) -target_include_directories(nccl PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) -target_include_directories(nccl PRIVATE ${CMAKE_CURRENT_BINARY_DIR})