From 6e2c3928156d68199bd5dd7f197bcd9f8bd91892 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Mon, 25 Aug 2025 09:14:43 -0700 Subject: [PATCH] use nccl header only when nccl is not present --- mlx/distributed/nccl/CMakeLists.txt | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/mlx/distributed/nccl/CMakeLists.txt b/mlx/distributed/nccl/CMakeLists.txt index 2f764c6ac..389a86c66 100644 --- a/mlx/distributed/nccl/CMakeLists.txt +++ b/mlx/distributed/nccl/CMakeLists.txt @@ -1,8 +1,20 @@ if(MLX_BUILD_CUDA) target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nccl.cpp) - find_package(NCCL REQUIRED) - target_link_libraries(mlx PRIVATE ${NCCL_LIBRARIES}) - target_include_directories(mlx PRIVATE ${NCCL_INCLUDE_DIRS}) + find_package(NCCL) + if(NCCL_FOUND) + target_link_libraries(mlx PRIVATE ${NCCL_LIBRARIES}) + target_include_directories(mlx PRIVATE ${NCCL_INCLUDE_DIRS}) + else() + message( + STATUS + "NCCL not found, using stubs. To run distributed with NCCL backend, install NCCL." + ) + file( + DOWNLOAD + "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in" + "${CMAKE_CURRENT_BINARY_DIR}/nccl.h") + target_include_directories(mlx PRIVATE "${CMAKE_CURRENT_BINARY_DIR}") + endif() else() target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp) endif()