From fd6d304b3a3f0e756eb960b727c5bce5ff8aff93 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Tue, 10 Mar 2026 20:58:04 -0700 Subject: [PATCH] win: fix cuda build (#3204) --- mlx/backend/cuda/CMakeLists.txt | 6 ++++-- mlx/backend/cuda/device.cpp | 6 ++++-- mlx/backend/cuda/eval.cpp | 4 ++++ mlx/backend/cuda/quantized/qmm/CMakeLists.txt | 2 +- mlx/backend/cuda/quantized/qmm/{qmm.cpp => qmm.cu} | 0 mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh | 4 +++- mlx/backend/cuda/scaled_dot_product_attention.cpp | 8 ++++++++ mlx/distributed/nccl/CMakeLists.txt | 2 +- 8 files changed, 25 insertions(+), 7 deletions(-) rename mlx/backend/cuda/quantized/qmm/{qmm.cpp => qmm.cu} (100%) diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index e968b1c1..5cac9e59 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -118,8 +118,10 @@ target_compile_options( mlx PRIVATE "$<$:--expt-relaxed-constexpr>") # Required for generating optimized CUTLASS code. -target_compile_options( - mlx PRIVATE "$<$:-Xcompiler=-fno-strict-aliasing>") +if(NOT MSVC) + target_compile_options( + mlx PRIVATE "$<$:-Xcompiler=-fno-strict-aliasing>") +endif() # Suppress nvcc warnings on C++ headers. target_compile_options( diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index ac46a0df..e4a11b07 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -66,8 +66,10 @@ Device::~Device() { void Device::make_current() { // We need to set/get current CUDA device very frequently, cache it to reduce - // actual calls of CUDA APIs. - static thread_local int current = 0; + // actual calls of CUDA APIs. Use -1 as sentinel so the first call on each + // new thread always calls cudaSetDevice (which establishes the CUDA primary + // context). Without this, device 0 would never get set on a new thread. + static thread_local int current = -1; if (current != device_) { CHECK_CUDA_ERROR(cudaSetDevice(device_)); current = device_; diff --git a/mlx/backend/cuda/eval.cpp b/mlx/backend/cuda/eval.cpp index 1d04c8f3..c8da4258 100644 --- a/mlx/backend/cuda/eval.cpp +++ b/mlx/backend/cuda/eval.cpp @@ -21,6 +21,10 @@ void new_stream(Stream s) { void eval(array& arr) { nvtx3::scoped_range r("gpu::eval"); + // Ensure CUDA context is active on this thread. Required when MLX is called + // from threads that have not yet established a CUDA context (e.g. thread + // pools, language runtimes that migrate work across OS threads). + cu::device(arr.primitive().stream().device).make_current(); auto outputs = arr.outputs(); { // If the array is a tracer hold a reference diff --git a/mlx/backend/cuda/quantized/qmm/CMakeLists.txt b/mlx/backend/cuda/quantized/qmm/CMakeLists.txt index 74a1c8a6..35162a6b 100644 --- a/mlx/backend/cuda/quantized/qmm/CMakeLists.txt +++ b/mlx/backend/cuda/quantized/qmm/CMakeLists.txt @@ -1,6 +1,6 @@ target_sources( mlx - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cpp + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n16_m1.cu diff --git a/mlx/backend/cuda/quantized/qmm/qmm.cpp b/mlx/backend/cuda/quantized/qmm/qmm.cu similarity index 100% rename from mlx/backend/cuda/quantized/qmm/qmm.cpp rename to mlx/backend/cuda/quantized/qmm/qmm.cu diff --git a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh index b8dc63ef..0007b9db 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh +++ b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh @@ -111,7 +111,9 @@ void qmm_sm90( reinterpret_cast(kernel), gemm.get_grid_shape(gemm.params()), GemmKernel::get_block_shape(), - {get<0>(cluster), get<1>(cluster), get<2>(cluster)}, + {static_cast(get<0>(cluster)), + static_cast(get<1>(cluster)), + static_cast(get<2>(cluster))}, GemmKernel::SharedStorageSize, kernel_params); } diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index 7f0e1f70..8b5a71ff 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -318,7 +318,15 @@ bool supports_sdpa_cudnn( bool has_arr_mask, bool do_causal, Stream s) { +#ifdef _WIN32 + // On Windows (WDDM), cuDNN SDPA has severe performance issues due to + // high per-kernel-launch overhead in the WDDM driver model. cuDNN's + // multi-kernel SDPA amplifies this, making it much slower than the + // single-kernel sdpa_vector path for both prefill and generation. + static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SDPA", 0); +#else static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SDPA", 1); +#endif if (!enabled) { return false; } diff --git a/mlx/distributed/nccl/CMakeLists.txt b/mlx/distributed/nccl/CMakeLists.txt index 644d060d..44182c39 100644 --- a/mlx/distributed/nccl/CMakeLists.txt +++ b/mlx/distributed/nccl/CMakeLists.txt @@ -1,4 +1,4 @@ -if(MLX_BUILD_CUDA) +if(MLX_BUILD_CUDA AND NOT WIN32) target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nccl.cpp) find_package(NCCL) if(NCCL_FOUND)