L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix building with CUDA toolkit 13.2 (#3273)

Browse Source
This commit is contained in:
Cheng
2026-03-19 08:31:44 +09:00
committed by GitHub
parent b41b349b67
commit 75f74ea9bc
8 changed files with 42 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+4
View File
@@ -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
+6 -2
View File
@@ -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
+28 -5
View File
@@ -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
+1 -1
View File
@@ -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
+2 -16
View File
@@ -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)
mlx/distributed/nccl/nccl_stub/CMakeLists.txt
-1
View File
@@ -1 +0,0 @@
mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp
-54
View File
@@ -1,54 +0,0 @@
#if !defined(_WIN32)
#include <arpa/inet.h>
#endif
#include <cuda_runtime.h>
#include <nccl.h>
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;
}
python/tests/test_array.py
+1
View File
@@ -2181,6 +2181,7 @@ class TestArray(mlx_tests.MLXTestCase):
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())