From d2f540f4e00ba60f800ba9828938992eb1f10e4b Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Mon, 25 Aug 2025 14:17:25 -0700
Subject: [PATCH] Use nccl header only when nccl is not present (#2539)

* use nccl header only when nccl is not present

* larger machine for cuda build
---
 .circleci/config.yml                |  2 +-
 mlx/distributed/nccl/CMakeLists.txt | 18 +++++++++++++++---
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 502640d3e..14bac4437 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -394,7 +394,7 @@ jobs:
         default: ""
     machine:
       image: ubuntu-2204:current
-      resource_class: large
+      resource_class: xlarge
     steps:
       - checkout
       - run:
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()