[CUDA] Cholesky via cuSOLVER

[sashko-zakharchuk]

Proposed changes

First op from the CUDA linalg gap discussed in #1392 (and #1026); inverse would follow.

• Cholesky::eval_gpu in the CUDA backend, backed by cuSOLVER: cusolverDnXpotrf per
  matrix, switching to cusolverDnSpotrfBatched for batches when num_matrices * 1024 > n.
  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 also
  ignores a positive info, so neither path reports a non positive definite input.
• 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: cusolver added to install_requires, the auditwheel excludes, and the
  MLX_LOAD_CUDA_LIBS_FROM_PYTHON rpaths. That is nvidia-cusolver-cu12==11.7.* on toolkit
  12 and nvidia-cusolver==12.* on toolkit 13, where the wheel is versioned 12.x the same
  way nvidia-cufft==12.* sits beside nvidia-cublas==13.*. The new rpath entry is for the
  cu12 wheel; the toolkit 13 wheel lands in nvidia/cu13/lib, already on the list. cusolver
  declares its cusparse/nvJitLink deps itself and finds them through its own rpath, so no
  further pins are needed.
• Windows: the CI toolkit install gains the cusolver subpackages, and the delay-load helper
  learns to resolve cusolver, registering the cusparse/nvjitlink wheel dirs alongside it. I
  have no Windows machine, so that path is only compile tested.
• On CUDA builds the tests add GPU checks against the CPU result on positive definite
  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):

           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.980     7.0x
       2048x2048    13.494     3.060     4.4x
       4096x4096    44.960    10.810     4.2x
        64x16x16     0.240     0.062     3.9x
        64x64x64     0.521     0.125     4.2x
       256x32x32     0.685     0.120     5.7x
      16x256x256     4.608     0.369    12.5x

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:

• For positive semi definite input (an exact zero eigenvalue) the undefined region differs:
  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 (sqrtA there has rank
  2): 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.
• 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, 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

• 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)

[sashko-zakharchuk]

Force-pushed: rebased on main, and fixed the cuda-12.6 failure.

The GPU assertions I had added to test_cholesky ran against that test's existing matrix,
sqrtA.T @ sqrtA / 81, which is singular (sqrtA is rank 2), so the factor past the rank
boundary is undefined. cuSOLVER 12.6 writes NaN there where 12.9 and 13.0 return finite
values, which is why only that one job went red and only on the upper=True assertion.

The existing assertions are back to untouched upstream code on the CPU stream, and the GPU
checks now use positive definite inputs sized to hit both cuSOLVER paths.

[haoxu0]

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 n <= 256 batched threshold may be hardware-dependent. On an L40S (sm_89, datacenter) potrfBatched stays substantially ahead of the loop well past n=256:

shape	batched	serialised loop
256 × 32²	0.089 ms	3.856 ms
64 × 128²	0.223 ms	4.919 ms
16 × 512²	0.721 ms	5.304 ms

So 16 x 512² would be ~7x slower routed through the loop on this GPU. Your crossover was measured on sm_120, which has a very different SM count and bandwidth — might be worth either raising the threshold or gating it on device properties.

The fill mode is asymmetric in cuSOLVER. upper=False (the MLX default) maps to CUBLAS_FILL_MODE_UPPER, which is consistently slower than FILL_MODE_LOWER at every size I measured — 512²: 2.04x vs 1.42x, 2048²: 1.96x vs 1.09x, 4096²: 1.97x vs 0.90x, all relative to torch.linalg.cholesky on the same GPU. Computing the fast mode and transposing is mathematically equivalent (cholesky(A, upper=True).T == cholesky(A, upper=False), verified numerically) but only pays off at large n: +13% at 4096², -8% at 512². Probably not worth doing, but worth knowing the default path is the slower one.

I also have a PyTorch comparison benchmark if that's useful — single matrices land at ~2x torch.linalg.cholesky, batched at 1.0-1.7x. Happy to open that separately or hand it over.

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.

[sashko-zakharchuk]

Corrected the benchmark table in the description. The GPU column was measured on a different card
than the one it names, so four rows were wrong:

      shape      was      now
  1024x1024    0.820    0.980
  2048x2048    1.919    3.060
  4096x4096    4.645   10.810
  256x32x32    0.085    0.120

Speedups on those rows become 7.0x, 4.4x, 4.2x and 5.7x. The other eight rows are unchanged. The
CPU column is unchanged too and still needs redoing on an idle machine.

I also removed the claim that a sweep on a second card landed within noise of these numbers and
that the batched/loop threshold held on both. The threshold does not hold, which is the point
@HaoXuAI raised above.

[sashko-zakharchuk]

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 potrfBatched wins:

   n \ batch      2      4      8     16     64
        256    1.25   2.20   3.44   6.62  11.18
        512    1.21   2.01   3.53   4.65   5.06
       1024    1.21   1.80   2.51   2.91   2.82
       2048    0.94   1.24   1.59   1.65   1.33
       4096    0.65   0.81   0.92   0.92      -

(64 x 4096² does not fit in 8 GB.) n <= 256 picks wrong in 14 of those 24, worst 5x. The batch
relative to the matrix size separates them, so I pushed num_matrices * 1024 > n. Your
16 x 512² is one of the 14 and routes batched under it; your other two already did. It misses
8 x 4096² and 16 x 4096², costing about 10% there.

On device properties: no constant of this shape fits even this one card. 4 x 2048² needs it above
512 and 16 x 4096² needs it at or below 256. Anything in (512, 1024] gets the same 22 of 24, so
the constant is not tuned finer than that. 8 x 4096² and 16 x 4096² on your L40S would show
whether an SM-count term is worth it.

Good catch on the fill mode. upper=True plus a transpose beats the default here: 1.4x at 512²,
2.1x at 1024², 2.2x at 2048², 1.4x at 4096², agreeing to about 5e-07 of the factor's largest
entry. swapaxes is a view, so the transpose itself is
free. That looks like its own change. A separate issue for the PyTorch numbers sounds right.

[sashko-zakharchuk]

Force-pushed: rebased onto main. The previous push was 9 commits behind, which is what made it
unmergeable. The only content change is the setup.py conflict with #3995, which pinned the
toolkit 13 dependencies: cusolver there is now nvidia-cusolver==12.*, matching
nvidia-cufft==12.*, since that wheel is versioned 12.x. Nothing else moved.