Repository navigation
Metal batched Cholesky - #3825
kapellirohith wants to merge 1 commit into
Conversation
Run cholesky on the GPU for batches of float32 matrices where it is faster than the CPU: one matrix per simdgroup for n <= 32, one per threadgroup up to n = 128, and a blocked right-looking factorization (batched panel + simdgroup-matrix SYRK dispatches) above that. Shapes below the measured batch thresholds keep using the CPU implementation via Cholesky::use_fallback, so no shape is slower than before. cholesky now also works on any shape and dtype on a GPU stream instead of raising. Also fixes a division by zero in the CPU implementation for inputs with n = 0.
zcbenz
left a comment
zcbenz
left a comment
There was a problem hiding this comment.
We don't silently fallback to cpu for gpu ops, the fallback is only used when fused ops being slower than unfused ones.
On the MPS cholesky kernel being slower than cpu kernel I'm almost certain that it was because of the overhead of creating the MPS graph. On the batched kernel while it is 2x faster than CPU but that is mostly because the CPU implementation is doing computation sequentially, and it would be a lot faster after we merge #3019 and use its cpu kernel thread pool for cholesky.
|
Makes sense on the fallback, I'll drop it and just run on the gpu stream. Agree the 2x is mostly the serial cpu path and #3019 will close most of it. Is a gpu version worth it in any regime after that (large batches of small matrices), or better to keep these on cpu? I can hold and re-benchmark against the threaded cpu once it lands, otherwise I'll close it. |
|
I think for small matrices the cpu route should be enough, what is needed is a GPU kernel that can efficiently handle large matrices. The MPS cholesky kernel should work with proper cache added like what we do with cuDNN kernels in CUDA backend. |
|
Makes sense, small on cpu is fine. I'll close this. I'll look at whether a cached MPS kernel can beat cpu on large matrices, along the lines of the cuDNN caching. |
Draft to check scope before polishing: is the shape-based dispatch in
linalg::cholesky(GPU for batches where it measurably wins, CPU for everything else) the right mechanism, or would you prefer the GPU kernels behind an explicit opt-in, or as part of a broader linalg-on-GPU plan? The same structure extends to LU next. Happy to adjust.Adds a Metal implementation of
Cholesky::eval_gpufor batched float32 inputs, following the batched-MAGMA design (Dong, Haidar, Tomov & Dongarra): one matrix per simdgroup in registers for n <= 32, one matrix per threadgroup in threadgroup memory up to n = 64, one per threadgroup in device memory up to n = 128, and a blocked right-looking factorization above that, where the host encodes one batched panel dispatch (POTF2 + TRTRI + TRSM as a simdgroup-matrix GEMM) and one batched SYRK dispatch over 32x32 lower-triangle tiles per 32-column panel. Panels are processed in pairs so the big trailing update runs at rank 64, halving the read-modify-write traffic on the trailing matrix.The GPU only wins for large enough batches, so
linalg::choleskyconsultsCholesky::use_fallback(same pattern asScaledDotProductAttention) at graph construction and keeps every other shape on the CPU implementation. Nothing gets slower: single or small-batch factorizations run exactly the CPU path they run today. This also meanscholeskynow works for any shape and dtype on a GPU stream instead of raising. Minimum batch for GPU dispatch, measured on an M3 Pro at both ends of each size range: 1000 for n in [13, 31], 500 for n <= 64, 250 for n <= 128, 128 for n <= 192, 100 for n <= 256, 80 for n <= 512, 48 for n <= 1024, 16 for n <= 2048, and always CPU for n <= 12 or n > 2048.Motivated by #1216, which stalled because a single-matrix GPU Cholesky loses to the CPU. Batching is where the GPU win is, and the fallback keeps the default path from regressing anywhere.
Also fixes a division by zero in the CPU implementation for inputs with n = 0 (
cholesky(zeros((3, 0, 0)))previously aborted the process).Benchmarks
M3 Pro, macOS 26.5, float32, best of 30 (ms). CPU is the existing LAPACK/Accelerate path.
Every GPU-dispatched shape measures >= 1.1x at the dispatch thresholds, checked at both ends of each size range; shapes just below a threshold run the CPU path unchanged.
Testing
Verified against the CPU implementation over n = 1..3000 and batch = 1..1000 (kernel-path boundaries, sizes that are not multiples of 8/32/64, narrow last panels, nested batch dims, upper and lower): worst relative reconstruction error ~2e-7, opposite triangle exactly zero. float64 and small batches fall back to the CPU and match it bitwise; float16/bfloat16 raise as before; zero-size inputs work on both streams. For non positive-definite inputs (documented as undefined behaviour) the GPU produces NaNs where the CPU returns a partial factorization.
test_cholesky_gpucovers each kernel path, the CPU fallbacks, and zero-size inputs. Both the default andMLX_METAL_JIT=ONbuilds pass.