metal-linalg

eigh and the SVD for batches of mid-size matrices

Status: analysed, not built (2026-10-08). The CPU path is 2-10x ahead of every GPU backend here, and what a GPU kernel could reach is about level with it, not ahead; see Why not now.

What

Batches of matrices of about 96 to 512 go to the CPU for eigh and the SVD, which spreads a batch over every core. QR’s batches of these sizes went to the GPU once it had a blocked kernel of its own (qr-mid-size-kernel.md); the question is whether eigh and the SVD could have the same.

The measurements

M5 Pro, 2.16.0, sweep_eigh and sweep_svd with MLX’s buffer cache on, median, ms:

batch x shape	CPU	best GPU backend
eigh, 1024 x 96^2	22.5	75 (block Jacobi)
eigh, 256 x 128^2	10.1	41 (block Jacobi)
eigh, 1024 x 128^2	38.7	174 (block Jacobi)
eigh, 64 x 256^2	12.2	74 (block Jacobi)
eigh, 256 x 256^2	44.5	328 (block Jacobi)
eigh, 16 x 512^2	16.4	89 (tridiag)
eigvalsh, 1024 x 96^2	12.0	64 (block Jacobi)
eigvalsh, 1024 x 128^2	20.7	144 (block Jacobi)
eigvalsh, 256 x 256^2	21.0	250 (block Jacobi)
eigvalsh, 16 x 512^2	7.5	58 (band)
SVD, 1024 x 96^2	42.5	84 (block Jacobi)
SVD, 1024 x 128^2	77.2	187 (block Jacobi)
SVD, 256 x 256^2	76.8	359 (block Jacobi)
svdvals, 1024 x 128^2	52.9	1626 (bidiag)
svdvals, 256 x 256^2	39.1	867 (bidiag)

At 256 x 256^2 the CPU does eigh’s roughly 9 N^3 flops a matrix at about 0.9 TFLOP/s: Accelerate runs on the CPU’s matrix units, a batch spread over 18 cores.

Why not now

QR’s mid-size kernel works because QR is one-sided: its updates are matrix products (8 x 8 simdgroup MMA, 2.7 TFLOP/s) and its panels a short chain. eigh and the SVD are two-sided, and their GPU kernels for small matrices (ql, golub_kahan) rely on the whole matrix, and then the eigenvectors, in threadgroup memory: that ends at N = 87 and k = 83 (32 KB). Beyond it:

• The QL or implicit-QR phase applies about N^2 dependent rotations to the vectors, each touching two columns. In threadgroup memory a sweep is one barrier; with the vectors in device memory it is some 50 MB of traffic for one 128 x 128 matrix (about 25 thousand rotations of two 128-float columns), 50 GB for a batch of 1024: far past the CPU’s 39 ms.
• The vectors in registers, a row a thread (128 floats a thread at 128), as QR’s register kernel keeps its matrix: a rotation’s columns are indexed at run time, which registers do not allow; unrolling every sweep over all N column pairs (as QR’s kernels unroll their columns) costs N checks a pair, and the tridiagonalization would need the rotate-to-index-0 scheme of QR’s register kernel as well. A rewrite of both kernels, at the occupancy of 130 registers a thread.
• The ideal: at 1024 x 96^2 eigh is about 7.7 MFLOP a matrix, 7.9 GFLOP in all; at the 0.5 TFLOP/s the ql kernel reaches on its latency-bound phases, 16 ms against the CPU’s 22.5 (1.4x), and only where everything fits.
• A two-stage reduction in one threadgroup (to a band by products, then the band chased to tridiagonal in threadgroup memory, then QL) would replace the BLAS2 traffic of the tridiagonalization, but leaves the QL phase’s rotations, the larger part with vectors.

The values alone (eigvalsh, svdvals) need no vectors, so the packed lower triangle (N (N + 1) / 2 floats) in threadgroup memory would take ql’s values-only path to N = 123; but the CPU’s values are faster still (12 ms at 1024 x 96^2), and ql_vals reaches about 1.6x the CPU only at 4096 x 32^2.

If someone picks it up

The narrow window is 88-128 with vectors, at batches of 1024 or more: the ql and golub_kahan kernels with the vectors in registers (a row a thread) and their rotations’ column pairs unrolled. Expected at best 1.2-1.5x the CPU there; measure ql at 64 x 64 against the CPU first, since the extension can only be slower per flop. Effort: 3-5 days.

This site is open source. Improve this page.