[CUDA] Cholesky via cuSOLVER - #4208
Conversation
zcbenz
added
the
await verification
4ac8ddc to
397e2cd
Compare
|
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 |
|
Nice work — I'd independently implemented the same op before this landed (my PR was closed as a duplicate, correctly). Two things from my testing that might be useful, and one thing yours does better than mine did. The
So The fill mode is asymmetric in cuSOLVER. I also have a PyTorch comparison benchmark if that's useful — single matrices land at ~2x One note for my own benefit: launching the pointer-fill kernel inside the capture context so stream order handles the ordering is neater than what I did (allocating the pointer array as an mlx array just to get a graph dependency edge). Stealing that. |
|
Corrected the benchmark table in the description. The GPU column was measured on a different card Speedups on those rows become 7.0x, 4.4x, 4.2x and 5.7x. The other eight rows are unchanged. The I also removed the claim that a sweep on a second card landed within noise of these numbers and |
397e2cd to
c4f5c48
Compare
|
You are right. I had only measured batches at or below n = 256. Ratio of looped to batched on an RTX 5050, above 1 means (64 x 4096² does not fit in 8 GB.) On device properties: no constant of this shape fits even this one card. 4 x 2048² needs it above Good catch on the fill mode. |
The n <= 256 cut sent large batches of medium matrices through the serialized
loop: 64x512^2 ran 5x slower than potrfBatched on an RTX 5050, and it picked the
slower path in 14 of 24 measured shapes.
potrfBatched parallelizes across the batch, so it wins until the batch is too
small to keep the device busy at that size. Measured across n in
{256, 512, 1024, 2048, 4096} and batch in {2, 4, 8, 16, 64}, 24 shapes because
64x4096^2 does not fit in 8 GB, the loop only wins for large matrices in small
batches. No constant fits every shape, so 1024 misses 8x4096^2 and 16x4096^2 by
about 10%.
The test shapes move with it: 2x2048 now covers the loop with more than one
matrix, which 2x512 used to do and no longer would.
c4f5c48 to
344f008
Compare
|
Force-pushed: rebased onto main. The previous push was 9 commits behind, which is what made it |
3 participants
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 whennum_matrices * 1024 > n.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.
install_requires, the auditwheel excludes, and theMLX_LOAD_CUDA_LIBS_FROM_PYTHONrpaths. That isnvidia-cusolver-cu12==11.7.*on toolkit12 and
nvidia-cusolver==12.*on toolkit 13, where the wheel is versioned 12.x the sameway
nvidia-cufft==12.*sits besidenvidia-cublas==13.*. The new rpath entry is for thecu12 wheel; the toolkit 13 wheel lands in
nvidia/cu13/lib, already on the list. cusolverdeclares its cusparse/nvJitLink deps itself and finds them 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 and two 2048x2048 through the loop, 16 8x8 through the batched
path, 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. Four rows of the GPU
column were timed on the wrong card and have been re-measured. The CPU column still needs
redoing on an idle machine.
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