Skip to content

[CUDA] Cholesky via cuSOLVER - #4262

Closed
HaoXuAI wants to merge 1 commit into
ml-explore:mainfrom
HaoXuAI:cuda-linalg-cholesky
Closed

[CUDA] Cholesky via cuSOLVER#4262
HaoXuAI wants to merge 1 commit into
ml-explore:mainfrom
HaoXuAI:cuda-linalg-cholesky

Conversation

@HaoXuAI

@HaoXuAI HaoXuAI commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

Cholesky is a NO_GPU stub on CUDA, so mx.linalg.cholesky requires a CPU
stream. This backs it with cuSOLVER potrf / potrfBatched, and adds the
cusolver_utils handle and error plumbing that getrf, geqrf, gesvd and
syevd can reuse.

linalg::cholesky no longer forces a CPU stream when CUDA is available.
Metal behaviour is unchanged: cu::is_available() is false there, so the
existing error is still raised.

Testing

Verified on an L40S (sm_89) with CUDA 12.6, against the CPU/LAPACK path:
test_linalg.py passes 19 tests / 266 subtests, including under DEVICE=gpu.
The new test covers batched, non-contiguous and empty inputs, and asserts the
unrequested triangle is exactly zero.

Benchmarks

benchmarks/python/cholesky_bench.py, fp32, MLX_ENABLE_TF32=0, vs
torch.linalg.cholesky on the same GPU (MLX ms / PyTorch ms):

shape MLX PyTorch ratio
512² 0.435 0.186 2.34x
2048² 1.891 0.951 1.99x
4096² 4.536 2.366 1.92x
4096² upper 3.830 4.332 0.88x
256 x 32² 0.089 0.058 1.53x
64 x 128² 0.223 0.134 1.66x
16 x 512² upper 0.666 0.664 1.00x

Two thirds of the single-matrix gap is potrf itself — even in its faster fill
mode 4096² takes 3.82 ms against PyTorch's 2.40 ms. The rest is that
upper=False maps to CUBLAS_FILL_MODE_UPPER, consistently slower than
FILL_MODE_LOWER. Computing the fast mode and transposing is equivalent and
measures +13% at n=4096 but -8% at n=512, so I left it out.

Note

check_cpu_stream_unless_cuda is new cross-cutting surface that each further
linalg op would touch. Happy to reshape it — a per-primitive capability query
may be preferable to a CUDA-specific branch.

Backs Cholesky::eval_gpu with potrf and potrfBatched, and adds the shared
cusolver_utils handle and error plumbing for the remaining linalg ops.
@zcbenz

zcbenz commented Aug 15, 2026

Copy link
Copy Markdown
Member

Thanks for the PR! There is actually already a almost identical PR #4208 so I'm closing this as a duplicate.

@zcbenz zcbenz closed this Aug 15, 2026
@HaoXuAI

HaoXuAI commented Aug 15, 2026

Copy link
Copy Markdown
Contributor Author

thanks, when I started the PR was not there

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants