zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-28 12:53:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

use nccl header only when nccl is not present

Browse Source
This commit is contained in:
Awni Hannun 2025-08-25 09:14:43 -07:00
parent 333ffea273
commit 6e2c392815
1 changed files with 15 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
mlx/distributed/nccl/CMakeLists.txt
View File

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