Skip to content

[CUDA] Cholesky via cuSOLVER - #4208

Open
sashko-zakharchuk wants to merge 2 commits into
ml-explore:mainfrom
sashko-zakharchuk:cuda-cholesky
Open

[CUDA] Cholesky via cuSOLVER#4208
sashko-zakharchuk wants to merge 2 commits into
ml-explore:mainfrom
sashko-zakharchuk:cuda-cholesky

Conversation

@sashko-zakharchuk

Copy link
Copy Markdown
Contributor

Proposed changes

First op from the CUDA linalg gap discussed in #1392 (and #1026); inverse is next as a
follow-up.

  • Cholesky::eval_gpu in the CUDA backend, backed by cuSOLVER: cusolverDnXpotrf per
    matrix, switching to cusolverDnSpotrfBatched for batches of matrices up to n = 256
    (crossover measured on sm_120). Handles are cached per device the same way as the cuBLAS
    and cuDNN ones (cusolver_utils.{h,cpp}).
  • A small kernel zeroes the untouched triangle after potrf, matching the CPU op's output
    exactly.
  • info is allocated but never read back: reading it costs a sync, and the CPU op already
    ignores it for non positive definite inputs, so behavior matches.
  • linalg::cholesky now accepts a GPU stream when the CUDA backend is available. Metal
    still raises at graph construction with the same message as before.
  • Wheel packaging: nvidia-cusolver added to install_requires, the auditwheel excludes,
    and the MLX_LOAD_CUDA_LIBS_FROM_PYTHON rpaths (the cu12 cusolver wheel resolves its
    cusparse/nvJitLink deps through its own rpath, so no further pins are needed). On the
    Windows side, the CI toolkit install gains the cusolver subpackages and the delay-load
    helper gets a cusolver branch that also registers the cusparse/nvjitlink wheel dirs. I
    have no Windows machine, so the Windows side is untested beyond the build logic.
  • The cholesky tests also run on the GPU stream on CUDA builds now, including one batch
    through each code path checked against the CPU result.

float64 stays CPU-only: GPU streams reject float64 at array construction, so the GPU path
only ever sees float32. Non contiguous inputs go through the copy that already runs before
the factorization, so the kernels always get dense row major matrices.

Benchmarks

RTX 5050 (sm_120), float32, against the CPU path on the same machine (Threadripper PRO
5975WX):

           shape    cpu ms    gpu ms  speedup
           16x16     0.195     0.044     4.5x
           64x64     0.035     0.080     0.4x
         128x128     0.190     0.132     1.4x
         256x256     0.436     0.218     2.0x
         512x512     1.360     0.384     3.5x
       1024x1024     6.849     0.820     8.4x
       2048x2048    13.494     1.919     7.0x
       4096x4096    44.960     4.645     9.7x
        64x16x16     0.240     0.062     3.9x
        64x64x64     0.521     0.125     4.2x
       256x32x32     0.685     0.085     8.0x
      16x256x256     4.608     0.369    12.5x

A single 64x64 is the one shape measured where the CPU is still faster. The same sweep on an
RTX PRO 6000 (GB202) lands within noise of these numbers, and the batched/loop threshold held
on both cards.

Beyond the updated unit tests, a 60-case differential run against the CPU implementation
(sizes 1 to 257, three batch shapes, both triangles, non contiguous input, empty, non
positive definite) matches everywhere at float32 tolerances.

Two behavior notes from stress testing:

  • For positive semi definite input (an exact zero eigenvalue) the undefined region differs:
    LAPACK leaves finite garbage past the rank boundary, cuSOLVER writes NaN from that row on.
    The valid leading block agrees to about 1e-5.
  • Two python threads running cholesky on separate mx.new_stream streams intermittently
    poison stream capture (cudaStreamEndCapture ... previous error during capture, roughly
    half of runs). Serializing our captures behind a mutex does not change the rate, which
    rules out everything this diff controls; the same two-thread pattern with matmul is clean,
    and MLX_USE_CUDA_GRAPHS=0 is immune. Single threaded and threads sharing a stream are
    both fine. I can open a separate issue with the repro and what I ruled out.

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants