[CUDA] Cholesky via cuSOLVER

.github/actions/build-wheel/action.yml

@@ -89,6 +89,7 @@ runs:
               --exclude libcuda* \
               --exclude libcudnn* \
               --exclude libcufft* \
+              --exclude libcusolver* \
               --exclude libnccl* \
               --exclude libnvrtc*
           fi

.github/actions/setup/action.yml

@@ -170,9 +170,9 @@ runs:
           }
         PACKAGES: |
           {
-            "cuda-12.6": ["cudart_12.6", "nvcc_12.6", "cublas_12.6", "cublas_dev_12.6", "cufft_12.6", "cufft_dev_12.6", "nvrtc_12.6", "nvrtc_dev_12.6"],
-            "cuda-12.9": ["cudart_12.9", "nvcc_12.9", "cublas_12.9", "cublas_dev_12.9", "cufft_12.9", "cufft_dev_12.9", "nvrtc_12.9", "nvrtc_dev_12.9"],
-            "cuda-13.0": ["cudart_13.0", "nvcc_13.0", "cublas_13.0", "cublas_dev_13.0", "cufft_13.0", "cufft_dev_13.0", "nvrtc_13.0", "nvrtc_dev_13.0", "crt_13.0", "nvvm_13.0", "nvptxcompiler_13.0"],
+            "cuda-12.6": ["cudart_12.6", "nvcc_12.6", "cublas_12.6", "cublas_dev_12.6", "cufft_12.6", "cufft_dev_12.6", "cusolver_12.6", "cusolver_dev_12.6", "nvrtc_12.6", "nvrtc_dev_12.6"],
+            "cuda-12.9": ["cudart_12.9", "nvcc_12.9", "cublas_12.9", "cublas_dev_12.9", "cufft_12.9", "cufft_dev_12.9", "cusolver_12.9", "cusolver_dev_12.9", "nvrtc_12.9", "nvrtc_dev_12.9"],
+            "cuda-13.0": ["cudart_13.0", "nvcc_13.0", "cublas_13.0", "cublas_dev_13.0", "cufft_13.0", "cufft_dev_13.0", "cusolver_13.0", "cusolver_dev_13.0", "nvrtc_13.0", "nvrtc_dev_13.0", "crt_13.0", "nvvm_13.0", "nvptxcompiler_13.0"],
           }
       run: |
         $ErrorActionPreference = "Stop"

ACKNOWLEDGMENTS.md

@@ -22,6 +22,7 @@ MLX was developed with contributions from the following individuals:
 - Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function.
 - katlun-lgtm: Added `reflect` and `symmetric` padding modes.
 - Erwin Zhang: Added `searchsorted`. Fixed the ring backend hanging when a peer disconnects. Improved JACCL error reporting.
+- Oleksandr Zakharchuk: Added the cuSOLVER `cholesky` for the CUDA backend.
 
 <a href="https://github.com/ml-explore/mlx/graphs/contributors">
   <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx&anon=0&columns=20&max=100&r=true" />

mlx/backend/cuda/CMakeLists.txt

@@ -9,6 +9,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/arange.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/binary_two.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/cholesky.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/copy.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/copy/copy_contiguous.cu
@@ -20,6 +21,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/cusolver_utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
@@ -217,6 +219,7 @@ else()
                "$ORIGIN/../../nvidia/cuda_nvrtc/lib"
                "$ORIGIN/../../nvidia/cudnn/lib"
                "$ORIGIN/../../nvidia/cufft/lib"
+               "$ORIGIN/../../nvidia/cusolver/lib"
                "$ORIGIN/../../nvidia/nccl/lib")
   endif()
 endif()
@@ -262,6 +265,9 @@ target_link_libraries(mlx PRIVATE CUDA::cublasLt)
 # Use cuFFT.
 target_link_libraries(mlx PRIVATE CUDA::cufft)
 
+# Use cuSOLVER.
+target_link_libraries(mlx PRIVATE CUDA::cusolver)
+
 # Use NVRTC and driver APIs.
 target_link_libraries(mlx PRIVATE CUDA::nvrtc CUDA::cuda_driver)
 

mlx/backend/cuda/cholesky.cu

@@ -0,0 +1,159 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/cusolver_utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+// potrf writes only one triangle; zero the other one to match the CPU op.
+template <typename T, typename IdxT>
+__global__ void zero_triangle(T* out, IdxT size, int32_t n, bool zero_below) {
+  IdxT index = cg::this_grid().thread_rank();
+  if (index >= size) {
+    return;
+  }
+  int32_t row = (index / n) % n;
+  int32_t col = index % n;
+  if (zero_below ? (col < row) : (col > row)) {
+    out[index] = T(0);
+  }
+}
+
+// potrfBatched wants a device array of per matrix pointers.
+__global__ void
+fill_matrix_pointers(void** ptrs, char* base, size_t stride, int32_t count) {
+  int32_t i = cg::this_grid().thread_rank();
+  if (i < count) {
+    ptrs[i] = base + i * stride;
+  }
+}
+
+} // namespace cu
+
+void Cholesky::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Cholesky::eval_gpu");
+  auto& s = stream();
+  auto& a = inputs[0];
+
+  // Copy the input to the output; the factorization runs in place.
+  copy_gpu(
+      a,
+      out,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      s);
+
+  int64_t n = a.shape(-1);
+  if (a.size() == 0) {
+    return;
+  }
+  int64_t num_matrices = a.size() / (n * n);
+
+  auto& encoder = cu::get_command_encoder(s);
+  encoder.set_output_array(out);
+
+  auto handle = get_cusolver_handle(encoder.device());
+  CHECK_CUSOLVER_ERROR(cusolverDnSetStream(handle, encoder.stream()));
+
+  // cuSOLVER is column major. The input is symmetric and a column major
+  // lower triangle is a row major upper one, so pass the opposite of
+  // upper_, same as the CPU op.
+  cublasFillMode_t uplo =
+      upper_ ? CUBLAS_FILL_MODE_LOWER : CUBLAS_FILL_MODE_UPPER;
+  // Only float32 reaches eval_gpu: float64 is rejected on GPU streams at
+  // array construction.
+  auto type = CUDA_R_32F;
+  auto* out_ptr = gpu_ptr<char>(out);
+  size_t matrix_bytes = n * n * out.itemsize();
+
+  // Matching the CPU op, info is never read back: a non positive definite
+  // input gives an undefined factor rather than an error.
+  auto* info = static_cast<int*>(
+      allocate_workspace(encoder, num_matrices * sizeof(int)));
+
+  // potrfBatched parallelizes across the batch, so it wins until the batch is
+  // too small to keep the device busy at that size. A loop of single
+  // factorizations only catches up for large matrices in small batches. No
+  // constant fits every measured shape, so 1024 is a compromise.
+  if (num_matrices > 1 && num_matrices <= INT32_MAX &&
+      num_matrices * 1024 > n) {
+    auto** ptrs = static_cast<void**>(
+        allocate_workspace(encoder, num_matrices * sizeof(void*)));
+    auto capture = encoder.capture_context();
+    int32_t count = num_matrices;
+    cu::fill_matrix_pointers<<<(count + 255) / 256, 256, 0, encoder.stream()>>>(
+        ptrs, out_ptr, matrix_bytes, count);
+    CHECK_CUSOLVER_ERROR(cusolverDnSpotrfBatched(
+        handle,
+        uplo,
+        n,
+        reinterpret_cast<float**>(ptrs),
+        /* lda */ n,
+        info,
+        count));
+  } else {
+    size_t device_bytes = 0;
+    size_t host_bytes = 0;
+    CHECK_CUSOLVER_ERROR(cusolverDnXpotrf_bufferSize(
+        handle,
+        /* params */ nullptr,
+        uplo,
+        n,
+        type,
+        out_ptr,
+        /* lda */ n,
+        type,
+        &device_bytes,
+        &host_bytes));
+
+    auto* device_ws = allocate_workspace(encoder, device_bytes);
+    auto host_ws = std::make_shared<std::vector<char>>(host_bytes);
+    if (host_bytes > 0) {
+      encoder.add_completed_handler([host_ws]() {});
+    }
+
+    auto capture = encoder.capture_context();
+    for (int64_t i = 0; i < num_matrices; ++i) {
+      CHECK_CUSOLVER_ERROR(cusolverDnXpotrf(
+          handle,
+          /* params */ nullptr,
+          uplo,
+          n,
+          type,
+          out_ptr + i * matrix_bytes,
+          /* lda */ n,
+          type,
+          device_ws,
+          device_bytes,
+          host_ws->data(),
+          host_bytes,
+          info + i));
+    }
+  }
+
+  encoder.set_output_array(out);
+  dispatch_bool(out.size() > INT32_MAX, [&](auto large) {
+    using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+    auto [num_blocks, block_dims] = get_launch_args(out, large());
+    encoder.add_kernel_node(
+        cu::zero_triangle<float, IdxT>,
+        num_blocks,
+        block_dims,
+        gpu_ptr<float>(out),
+        out.size(),
+        n,
+        upper_);
+  });
+}
+
+} // namespace mlx::core

mlx/backend/cuda/cusolver_utils.cpp

@@ -0,0 +1,47 @@
+// Copyright © 2026 Apple Inc.
+#include "mlx/backend/cuda/cusolver_utils.h"
+#include "mlx/backend/gpu/device_info.h"
+
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+namespace {
+
+auto& cusolver_handles_cache() {
+  struct CusolverHandle {
+    ~CusolverHandle() {
+      if (handle) {
+        // Not checked: runs at thread exit where a throw would terminate.
+        cusolverDnDestroy(handle);
+      }
+    }
+    cusolverDnHandle_t handle{nullptr};
+  };
+  static thread_local std::vector<CusolverHandle> cache(gpu::device_count());
+  return cache;
+}
+
+} // namespace
+
+cusolverDnHandle_t get_cusolver_handle(cu::Device& device) {
+  auto& storage = cusolver_handles_cache().at(device.cuda_device());
+  if (!storage.handle) {
+    device.make_current();
+    CHECK_CUSOLVER_ERROR(cusolverDnCreate(&storage.handle));
+  }
+  return storage.handle;
+}
+
+void init_cusolver_handles_cache() {
+  cusolver_handles_cache();
+}
+
+void check_cusolver_error(const char* name, cusolverStatus_t err) {
+  if (err != CUSOLVER_STATUS_SUCCESS) {
+    throw std::runtime_error(
+        fmt::format("{} failed with code: {}.", name, static_cast<int>(err)));
+  }
+}
+
+} // namespace mlx::core

mlx/backend/cuda/cusolver_utils.h

@@ -0,0 +1,18 @@
+// Copyright © 2026 Apple Inc.
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+
+#include <cusolverDn.h>
+
+namespace mlx::core {
+
+void check_cusolver_error(const char* name, cusolverStatus_t err);
+
+#define CHECK_CUSOLVER_ERROR(cmd) check_cusolver_error(#cmd, (cmd))
+
+void init_cusolver_handles_cache();
+
+cusolverDnHandle_t get_cusolver_handle(cu::Device& device);
+
+} // namespace mlx::core

mlx/backend/cuda/delayload.cpp

@@ -24,6 +24,24 @@ inline fs::path cublas_dir() {
                         : relative_to_current_binary("../nvidia/cublas/bin");
 }
 
+inline fs::path cusolver_dir() {
+  return cuda_bin_dir() ? fs::path(cuda_bin_dir())
+                        : relative_to_current_binary("../nvidia/cusolver/bin");
+}
+
+fs::path load_cusolver() {
+  fs::path dir = cusolver_dir();
+  // cusolver resolves cusparse, nvjitlink and cublas at load time, add their
+  // wheel dirs to the search path.
+  ::AddDllDirectory(dir.c_str());
+  ::AddDllDirectory(cublas_dir().c_str());
+  ::AddDllDirectory(
+      relative_to_current_binary("../nvidia/cusparse/bin").c_str());
+  ::AddDllDirectory(
+      relative_to_current_binary("../nvidia/nvjitlink/bin").c_str());
+  return dir;
+}
+
 fs::path load_nvrtc() {
   fs::path nvrtc_dir = cuda_bin_dir()
       ? fs::path(cuda_bin_dir())
@@ -63,6 +81,9 @@ FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) {
       mod = ::LoadLibraryW((cudnn_dir / dll).c_str());
     } else if (dll.starts_with("cublas")) {
       mod = ::LoadLibraryW((cublas_dir() / dll).c_str());
+    } else if (dll.starts_with("cusolver")) {
+      static auto cusolver_dir = load_cusolver();
+      mod = ::LoadLibraryW((cusolver_dir / dll).c_str());
     } else if (dll.starts_with("nvrtc")) {
       static auto nvrtc_dir = load_nvrtc();
       mod = ::LoadLibraryW((nvrtc_dir / dll).c_str());

mlx/backend/cuda/eval.cpp

@@ -4,6 +4,7 @@
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/cublas_utils.h"
 #include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/cusolver_utils.h"
 #include "mlx/backend/cuda/event.h"
 #include "mlx/primitives.h"
 #include "mlx/scheduler.h"
@@ -20,6 +21,7 @@ void init() {
   mlx::core::cu::CudaEvent::init_pool();
   init_cublas_handles_cache();
   init_cudnn_handles_cache();
+  init_cusolver_handles_cache();
   init_cudnn_conv_cache();
   init_cudnn_sdpa_cache();
 }

mlx/backend/cuda/primitives.cpp

@@ -28,7 +28,6 @@ NO_GPU_MULTI(LUF)
 NO_GPU_MULTI(QRF)
 NO_GPU_MULTI(SVD)
 NO_GPU(Inverse)
-NO_GPU(Cholesky)
 NO_GPU_MULTI(Eig)
 NO_GPU_MULTI(Eigh)
 

mlx/linalg.cpp

@@ -4,6 +4,7 @@
 #include <ostream>
 #include <vector>
 
+#include "mlx/backend/cuda/cuda.h"
 #include "mlx/linalg.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
@@ -18,6 +19,18 @@ void check_cpu_stream(const StreamOrDevice& s, const std::string& prefix) {
         "Explicitly pass a CPU stream to run it.");
   }
 }
+
+// For ops that have a CUDA implementation but no Metal one yet.
+void check_cpu_or_cuda_stream(
+    const StreamOrDevice& s,
+    const std::string& prefix) {
+  if (to_stream(s).device == Device::gpu && !cu::is_available()) {
+    throw std::invalid_argument(
+        prefix +
+        " This op is not yet supported on the GPU. "
+        "Explicitly pass a CPU stream to run it.");
+  }
+}
 void check_float(Dtype dtype, const std::string& prefix) {
   if (dtype != float32 && dtype != float64) {
     std::ostringstream msg;
@@ -336,7 +349,7 @@ array cholesky(
     const array& a,
     bool upper /* = false */,
     StreamOrDevice s /* = {} */) {
-  check_cpu_stream(s, "[linalg::cholesky]");
+  check_cpu_or_cuda_stream(s, "[linalg::cholesky]");
   check_float(a.dtype(), "[linalg::cholesky]");
   if (a.ndim() < 2) {
     std::ostringstream msg;