Delay load CUDA libs and resolve DLL paths at runtime (#3061)

mlx/backend/cuda/CMakeLists.txt

@@ -151,16 +151,46 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
 set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                      "${MLX_CUDA_ARCHITECTURES}")
 
-if(MLX_BUILD_PYTHON_BINDINGS)
-  set_property(
-    TARGET mlx
-    APPEND
-    PROPERTY INSTALL_RPATH
-             # The paths here should match the install_requires in setup.py.
-             "$ORIGIN/../../nvidia/cublas/lib"
-             "$ORIGIN/../../nvidia/cuda_nvrtc/lib"
-             "$ORIGIN/../../nvidia/cudnn/lib"
-             "$ORIGIN/../../nvidia/nccl/lib")
+# Search CUDA libs from installed python packages.
+if(WIN32)
+  # Resolve paths of unfound DLL at runtime.
+  if(BUILD_SHARED_LIBS)
+    target_link_libraries(mlx PRIVATE "delayimp.lib")
+    target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp)
+  else()
+    # For static library the delayload must be compiled into final executables.
+    target_link_libraries(mlx PUBLIC "delayimp.lib")
+    target_sources(
+      mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp>)
+  endif()
+  # Get all the CUDA DLLs we could link with.
+  file(
+    GLOB CUDA_DLL_NAMES
+    RELATIVE "${CUDAToolkit_BIN_DIR}/x64"
+    "${CUDAToolkit_BIN_DIR}/x64/*.dll")
+  # Delay load CUDA and cuDNN libs.
+  foreach(CUDA_DLL ${CUDA_DLL_NAMES} ${CUDNN_DLL_NAMES})
+    target_link_options(mlx PUBLIC "/DELAYLOAD:${CUDA_DLL}")
+  endforeach()
+  # Pass the locations where CUDA DLLs are placed.
+  if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON)
+    target_compile_definitions(
+      mlx PUBLIC MLX_CUDA_BIN_DIR="${CUDAToolkit_BIN_DIR}/x64"
+                 MLX_CUDNN_BIN_DIR="${CUDNN_BIN_DIR}")
+  endif()
+else()
+  # For POSIX we rely on RPATH to search for CUDA libs.
+  if(MLX_LOAD_CUDA_LIBS_FROM_PYTHON)
+    set_property(
+      TARGET mlx
+      APPEND
+      PROPERTY INSTALL_RPATH
+               # The paths here should match the install_requires in setup.py.
+               "$ORIGIN/../../nvidia/cublas/lib"
+               "$ORIGIN/../../nvidia/cuda_nvrtc/lib"
+               "$ORIGIN/../../nvidia/cudnn/lib"
+               "$ORIGIN/../../nvidia/nccl/lib")
+  endif()
 endif()
 
 # ------------------------ Dependencies ------------------------

mlx/backend/cuda/delayload.cpp

@@ -0,0 +1,80 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+
+// clang-format off
+#include <windows.h> // must be included first
+#include <delayimp.h>
+// clang-format on
+
+namespace mlx::core {
+
+namespace fs = std::filesystem;
+
+inline fs::path relative_to_current_binary(const char* relative) {
+  return fs::absolute(current_binary_dir() / relative);
+}
+
+inline fs::path cublas_bin_dir() {
+#if defined(MLX_CUDA_BIN_DIR)
+  return MLX_CUDA_BIN_DIR;
+#else
+  return relative_to_current_binary("../nvidia/cublas/bin");
+#endif
+}
+
+fs::path load_nvrtc() {
+#if defined(MLX_CUDA_BIN_DIR)
+  fs::path nvrtc_bin_dir = MLX_CUDA_BIN_DIR;
+#else
+  fs::path nvrtc_bin_dir =
+      relative_to_current_binary("../nvidia/cuda_nvrtc/bin");
+#endif
+  // Internally nvrtc loads some libs dynamically, add to search dirs.
+  ::AddDllDirectory(nvrtc_bin_dir.c_str());
+  return nvrtc_bin_dir;
+}
+
+fs::path load_cudnn() {
+#if defined(MLX_CUDNN_BIN_DIR)
+  fs::path cudnn_bin_dir = MLX_CUDNN_BIN_DIR;
+#else
+  fs::path cudnn_bin_dir = relative_to_current_binary("../nvidia/cudnn/bin");
+#endif
+  // Must load cudnn_graph64_9.dll before locating symbols, otherwise We would
+  // get errors like "Invalid handle. Cannot load symbol cudnnCreate".
+  for (const auto& dll : fs::directory_iterator(cudnn_bin_dir)) {
+    if (dll.path().filename().string().starts_with("cudnn_graph") &&
+        dll.path().extension() == ".dll") {
+      ::LoadLibraryW(dll.path().c_str());
+      break;
+    }
+  }
+  // Internally cuDNN loads some libs dynamically, add to search dirs.
+  load_nvrtc();
+  ::AddDllDirectory(cudnn_bin_dir.c_str());
+  ::AddDllDirectory(cublas_bin_dir().c_str());
+  return cudnn_bin_dir;
+}
+
+// Called by system when failed to locate a lazy-loaded DLL.
+FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) {
+  HMODULE mod = NULL;
+  if (dliNotify == dliNotePreLoadLibrary) {
+    std::string dll = pdli->szDll;
+    if (dll.starts_with("cudnn")) {
+      static auto cudnn_bin_dir = load_cudnn();
+      mod = ::LoadLibraryW((cudnn_bin_dir / dll).c_str());
+    } else if (dll.starts_with("cublas")) {
+      mod = ::LoadLibraryW((cublas_bin_dir() / dll).c_str());
+    } else if (dll.starts_with("nvrtc")) {
+      static auto nvrtc_bin_dir = load_nvrtc();
+      mod = ::LoadLibraryW((nvrtc_bin_dir / dll).c_str());
+    }
+  }
+  return reinterpret_cast<FARPROC>(mod);
+}
+
+} // namespace mlx::core
+
+extern "C" const PfnDliHook __pfnDliNotifyHook2 = mlx::core::delayload_helper;

mlx/backend/cuda/jit_module.cpp

@@ -31,7 +31,10 @@ const std::vector<std::string>& include_path_args() {
   static std::vector<std::string> cached_args = []() {
     std::vector<std::string> args;
     // Add path to bundled CCCL headers.
-    auto root_dir = current_binary_dir().parent_path();
+    auto root_dir = current_binary_dir();
+#if !defined(_WIN32)
+    root_dir = root_dir.parent_path();
+#endif
     auto path = root_dir / "include" / "cccl";
 #if defined(MLX_CCCL_DIR)
     if (!std::filesystem::exists(path)) {

mlx/backend/cuda/sort.cu

@@ -73,8 +73,8 @@ struct LessThan {
 
   __device__ __forceinline__ bool operator()(T a, T b) const {
     if constexpr (std::is_floating_point_v<T>) {
-      bool an = std::isnan(a);
-      bool bn = std::isnan(b);
+      bool an = cuda::std::isnan(a);
+      bool bn = cuda::std::isnan(b);
       if (an | bn) {
         return (!an) & bn;
       }

setup.py

@@ -110,6 +110,8 @@ class CMakeBuild(build_ext):
                 )
             )
             cmake_args += [f"-DMLX_CUDA_ARCHITECTURES={cuda_archs}"]
+            # Search CUDA libs from python packages.
+            cmake_args += ["-DMLX_LOAD_CUDA_LIBS_FROM_PYTHON=ON"]
 
         # Some generators require explcitly passing config when building.
         build_args = ["--config", cfg]