From ed9ffdf4c142292693c529a99685803b08bcceac Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 18 Mar 2026 10:33:04 +0100 Subject: [PATCH 1/4] Forbid building with CUDA 13.1 --- CMakeLists.txt | 4 ++++ mlx/backend/cuda/CMakeLists.txt | 8 ++++++-- mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh | 2 +- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 041a476cf8..e315c160a8 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 08d1ba1d4c..1cee777bbe 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -118,8 +118,12 @@ target_compile_options(mlx target_compile_options( mlx PRIVATE "$<$:--expt-relaxed-constexpr>") -# Required for generating optimized CUTLASS code. -if(NOT MSVC) +if(MSVC) + # Ignore warnings from CUTLASS. + target_compile_options( + mlx PRIVATE $<$:-Xcudafe="--diag_suppress=2908">) +else() + # Required for generating optimized CUTLASS code. target_compile_options( mlx PRIVATE "$<$:-Xcompiler=-fno-strict-aliasing>") endif() diff --git a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh index 46bfe80667..a202683c31 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh +++ b/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); } From 0f757ada01685de1539411bfb820bbff5ccd3b8f Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 18 Mar 2026 19:12:56 +0900 Subject: [PATCH 2/4] Remove nccl_stubs --- mlx/distributed/nccl/CMakeLists.txt | 18 +------ mlx/distributed/nccl/nccl_stub/CMakeLists.txt | 1 - mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp | 54 ------------------- 3 files changed, 2 insertions(+), 71 deletions(-) delete mode 100644 mlx/distributed/nccl/nccl_stub/CMakeLists.txt delete mode 100644 mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp diff --git a/mlx/distributed/nccl/CMakeLists.txt b/mlx/distributed/nccl/CMakeLists.txt index 44182c3924..d67f18e482 100644 --- a/mlx/distributed/nccl/CMakeLists.txt +++ b/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 $) + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp) 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 deleted file mode 100644 index 8b13789179..0000000000 --- a/mlx/distributed/nccl/nccl_stub/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp b/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp deleted file mode 100644 index cd1b522297..0000000000 --- a/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp +++ /dev/null @@ -1,54 +0,0 @@ -#if !defined(_WIN32) -#include -#endif -#include -#include - -ncclResult_t ncclGetUniqueId(ncclUniqueId*) { - return ncclSuccess; -} - -const char* ncclGetErrorString(ncclResult_t result) { - return nullptr; -} - -ncclResult_t -ncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank) { - return ncclSuccess; -} - -ncclResult_t ncclCommDestroy(ncclComm_t comm) { - return ncclSuccess; -} - -ncclResult_t ncclAllGather( - const void* sendbuff, - void* recvbuff, - size_t sendcount, - ncclDataType_t datatype, - ncclComm_t comm, - cudaStream_t stream) { - return ncclSuccess; -} - -ncclResult_t ncclAllReduce( - const void* sendbuff, - void* recvbuff, - size_t count, - ncclDataType_t datatype, - ncclRedOp_t op, - ncclComm_t comm, - cudaStream_t stream) { - return ncclSuccess; -} - -ncclResult_t ncclReduceScatter( - const void* sendbuff, - void* recvbuff, - size_t recvcount, - ncclDataType_t datatype, - ncclRedOp_t op, - ncclComm_t comm, - cudaStream_t stream) { - return ncclSuccess; -} From 85229640a11d98a997c98e45ef323163a53b888f Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 18 Mar 2026 20:09:34 +0900 Subject: [PATCH 3/4] Add __launch_bounds__ for binary ops --- mlx/backend/cuda/binary/binary.cuh | 33 +++++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/mlx/backend/cuda/binary/binary.cuh b/mlx/backend/cuda/binary/binary.cuh index 5f09572fae..bbbb5ffaf2 100644 --- a/mlx/backend/cuda/binary/binary.cuh +++ b/mlx/backend/cuda/binary/binary.cuh @@ -16,8 +16,14 @@ namespace cu { namespace cg = cooperative_groups; +constexpr int BINARY_MAX_BLOCK_DIM = 1024; + template -__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 -__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 -__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 -__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; } 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, From ce04042294320d7d375688fc796f77ca8263d0a7 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 18 Mar 2026 12:44:52 +0100 Subject: [PATCH 4/4] Fix test on windows --- python/tests/test_array.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 4d673d5b4a..4555990e07 100644 --- a/python/tests/test_array.py +++ b/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())