Metal Cholesky

[kapellirohith]

Adds a Metal implementation of Cholesky::eval_gpu (float32), following up on the request in #3825 for a GPU kernel that efficiently handles large matrices. A GPU stream runs on the GPU for every input shape; there is no CPU rerouting anywhere in the op.

The factorization is blocked right-looking with one algorithm for all sizes: strips of 512 columns are factored 32 columns at a time (POTF2 on one simdgroup with simd_shuffle, then TRSM against the inverted diagonal block on simdgroup_matrix fragments, then a strip-bounded SYRK on 32x32 simdgroup-matrix tiles), and the trailing update beyond the strip, which is where nearly all the FLOPs are for a large matrix, runs on the existing steel GEMM (steel_matmul_regular_axpby with alpha = -1, beta = 1 over strided submatrix views, no copies). The GEMM accumulates in place (C and D alias): each threadgroup reads its C tile in the epilogue after the K loop and then stores D to the same tile, so no threadgroup touches another's region; flagging this since it relies on that property of the fused kernel. The three phases are separate dispatches ordered by buffer barriers, so no memory location is read and written inside the same dispatch. Four kernels total: cholesky_potf2, cholesky_trsm, cholesky_syrk, cholesky_fixup.

On the MPS route suggested in #3825: I measured MPSMatrixDecompositionCholesky first, with the kernel object created once and reused (the cuDNN-style caching). It is execution bound, not setup bound, and loses to the CPU at every large size, so a custom kernel was the only path to a win. Standalone measurement on an M3 Pro (status = success, factor verified against LAPACK):

n	MPS cached (ms)	CPU (ms)	MPS/CPU
1024	41.8	0.7	0.02x
2048	163.3	5.4	0.03x
4096	856.5	101.7	0.12x
8192	4131.9	271.1	0.07x

Benchmarks

M3 Pro, macOS 26.5, float32, best of N, CPU is the existing LAPACK/Accelerate path. GFLOP/s counts n^3/3.

Single large matrices, the case #3825 asked for:

shape	CPU ms	GPU ms	speedup	GPU GFLOP/s
4096	61.1	32.0	1.9x	716
6144	129.0	73.0	1.8x	1059
8192	279.3	140.5	2.0x	1304

The CPU time varies between runs (61 to 73 ms at n = 4096 across sessions); the GPU time is stable, so the large-matrix speedup lands between 1.9x and 2.3x depending on the run.

Batched and mid-size shapes also currently win:

shape	CPU ms	GPU ms	speedup
1024 x16	14.6	10.2	1.4x
2048 x4	19.2	17.1	1.1x
32 x1000	1.8	0.8	2.2x
64 x1000	5.9	2.0	2.9x
256 x120	4.0	2.7	1.5x

The threaded CPU pool from #3019 is expected to erode these batched and mid-size margins once the CPU loops over the batch in parallel, while the single large-matrix win above is unaffected by it.

Honest limits: single matrices in the 1024 to 3584 range lose to the CPU (0.2x to 0.8x); the panel factorization is a serial dependency chain and Accelerate's AMX is very strong there. Tiny shapes on a GPU stream are correct but slow (a single 256 matrix is about 0.06x); as discussed in #3825, small factorizations are expected to use the CPU stream. For inputs that are not positive definite (documented as undefined behaviour) the GPU produces NaNs where the CPU returns a partial factorization. Numbers are from an M3 Pro only; I do not have other Apple silicon to measure.

Testing

Verified against the CPU implementation over n = 1 to 8192 (powers of two, non-multiples of the 32 and 512 block sizes, panel and strip boundaries), single and batched, upper and lower, non-contiguous inputs, and zero-size shapes: worst relative reconstruction error about 2e-7 (float32), opposite triangle exactly zero. float64 on a GPU stream raises instead of silently running on the CPU. test_cholesky_gpu covers the size boundaries, batching, non-contiguous input, float64, and zero-size shapes. Both the default and MLX_METAL_JIT=ON builds pass. Also fixes a division by zero in the CPU implementation for n = 0 inputs, which previously aborted the process.