Fix building with CUDA toolkit 13.2

CMakeLists.txt

@@ -156,6 +156,10 @@ if(MLX_BUILD_CUDA)
   enable_language(CUDA)
   find_package(CUDAToolkit REQUIRED)
   find_package(CUDNN REQUIRED)
+  if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.1" AND CUDAToolkit_VERSION
+                                                          VERSION_LESS "13.2")
+    message(FATAL_ERROR "CUDA Toolkit 13.1 is not supported.")
+  endif()
 endif()
 
 if(MLX_BUILD_METAL)

mlx/backend/cuda/CMakeLists.txt

@@ -118,8 +118,12 @@ target_compile_options(mlx
 target_compile_options(
   mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")
 
-# Required for generating optimized CUTLASS code.
-if(NOT MSVC)
+if(MSVC)
+  # Ignore warnings from CUTLASS.
+  target_compile_options(
+    mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=2908">)
+else()
+  # Required for generating optimized CUTLASS code.
   target_compile_options(
     mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-strict-aliasing>")
 endif()

mlx/backend/cuda/binary/binary.cuh

@@ -16,8 +16,14 @@ namespace cu {
 
 namespace cg = cooperative_groups;
 
+constexpr int BINARY_MAX_BLOCK_DIM = 1024;
+
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_ss(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size) {
   IdxT index = cg::this_grid().thread_rank();
 
   if ((index + 1) * N_READS > size) {
@@ -36,7 +42,11 @@ __global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_sv(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size) {
   IdxT index = cg::this_grid().thread_rank();
 
   if ((index + 1) * N_READS > size) {
@@ -57,7 +67,11 @@ __global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vs(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size) {
   IdxT index = cg::this_grid().thread_rank();
 
   if ((index + 1) * N_READS > size) {
@@ -78,7 +92,11 @@ __global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
 }
 
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_vv(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vv(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size) {
   IdxT index = cg::this_grid().thread_rank();
 
   if ((index + 1) * N_READS > size) {
@@ -331,7 +349,12 @@ void binary_op_gpu_inplace(
               kernel = cu::binary_vv<Op, InType, OutType, IdxT, N_READS>;
             }
             auto [num_blocks, block_dims] = get_launch_args(
-                out.data_size(), out.shape(), out.strides(), large(), N_READS);
+                out.data_size(),
+                out.shape(),
+                out.strides(),
+                large(),
+                N_READS,
+                cu::BINARY_MAX_BLOCK_DIM);
             encoder.add_kernel_node(
                 kernel,
                 num_blocks,

mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh

@@ -205,8 +205,8 @@ __global__ void qmm_sm80_kernel(
   auto K_BLOCK_MAX = size<2>(tCrA);
 
   // Prefetch beginning tiles.
-  CUTE_UNROLL
   int tile_pipe = 0;
+  CUTE_UNROLL
   for (; tile_pipe < K_PIPE_MAX - 1; ++tile_pipe) {
     fetch_gmem(tile_pipe);
   }

mlx/distributed/nccl/CMakeLists.txt

@@ -1,25 +1,11 @@
 if(MLX_BUILD_CUDA AND NOT WIN32)
-  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nccl.cpp)
   find_package(NCCL)
   if(NCCL_FOUND)
     target_link_libraries(mlx PRIVATE ${NCCL_LIBRARIES})
     target_include_directories(mlx PRIVATE ${NCCL_INCLUDE_DIRS})
+    target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nccl.cpp)
   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")
-    add_library(nccl_stub OBJECT
-                ${CMAKE_CURRENT_SOURCE_DIR}/nccl_stub/nccl_stubs.cpp)
-    target_include_directories(
-      nccl_stub
-      PRIVATE ${CUDAToolkit_INCLUDE_DIRS}
-      PUBLIC ${CMAKE_CURRENT_BINARY_DIR})
-    target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:nccl_stub>)
+    target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp)
   endif()
 else()
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp)

python/tests/test_array.py

@@ -2081,6 +2081,7 @@ def test_deep_graphs(self):
             x = mx.sin(x)
         mx.eval(x)
 
+    @unittest.skipIf(platform.system() == "Windows", "Memory info not accurate")
     def test_siblings_without_eval(self):
         def get_mem():
             process = psutil.Process(os.getpid())