[CUDA] Cholesky via cuSOLVER - #4208
Open
sashko-zakharchuk wants to merge 2 commits into
Open
Conversation
sashko-zakharchuk
force-pushed
the
cuda-cholesky
branch
from
August 13, 2026 14:17
4ac8ddc to
397e2cd
Compare
Contributor
Author
|
Force-pushed: rebased on main, and fixed the cuda-12.6 failure. The GPU assertions I had added to The existing assertions are back to untouched upstream code on the CPU stream, and the GPU |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed changes
First op from the CUDA linalg gap discussed in #1392 (and #1026); inverse would follow.
Cholesky::eval_gpuin the CUDA backend, backed by cuSOLVER:cusolverDnXpotrfpermatrix, switching to
cusolverDnSpotrfBatchedfor 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}).potrf, matching the CPU op's outputexactly.
infois allocated but never read back: reading it costs a sync, and the CPU op alsoignores a positive
info, so neither path reports a non positive definite input.linalg::choleskynow accepts a GPU stream when the CUDA backend is available. Metalstill raises at graph construction with the same message as before.
nvidia-cusolveradded toinstall_requires, the auditwheel excludes,and the
MLX_LOAD_CUDA_LIBS_FROM_PYTHONrpaths (the cu12 cusolver wheel resolves itscusparse/nvJitLink deps through its own rpath, so no further pins are needed).
learns to resolve cusolver, registering the cusparse/nvjitlink wheel dirs alongside it. I
have no Windows machine, so that path is only compile tested.
inputs: a single 3x3, 16 8x8 through the batched path, two 512x512 through the loop, plus
empty and non contiguous inputs.
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):
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:
LAPACK leaves finite garbage past the rank boundary, cuSOLVER usually writes NaN from that
row on, and whether it does varies by version. The valid leading block agrees to about
1e-5. Worth knowing because
test_cholesky's matrix is singular (sqrtAthere has rank2): on it cuSOLVER 12.6 writes NaN into the upper factor while 12.9 and 13.0 do not, so
the new GPU checks use positive definite inputs instead.
mx.new_streamstreams intermittentlypoison stream capture (
cudaStreamEndCapture ... previous error during capture, roughlyhalf of runs). Serializing our captures behind a mutex does not change the rate, and the
same two-thread pattern with matmul does not fail at all, so I do not think it is the
cholesky call itself. It does not happen single threaded, with threads sharing a stream,
or with
MLX_USE_CUDA_GRAPHS=0. I can open a separate issue with the repro.Checklist
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes