[Panadestein]

🤖 AI text below 🤖

What

Two performance changes for large sweeps on GPUs:

1. Large thin SVDs on the device (cherry-picked from bench/kicked-ising, c592cb0). thin_svd runs SVDs with a side of at least 512 on the GPU (cuSOLVER gesvd) and keeps smaller ones on the host. At n = 2048, complex128, an A100 takes 0.63 s against 2.2 s for NumPy's OpenBLAS gesdd, and gesvd converges on ill-conditioned factors where gesdd fails.

2. Distributed CholeskyQR of the sketch on several GPUs. At every site of the right-to-left pass, the sketch used to be gathered on rank 0, factored there by a Householder QR, and its rows scattered, while the other GPUs waited. That serial step is why 4 GPUs only reached 1.5x of one at chi = 2048. From chi_out = 1024 on, the sketch is now factored in place:

   • an all-to-all turns each GPU's block of sketch columns into a block of rows;
   • shifted CholeskyQR3 (Fukaya et al., SIAM J. Sci. Comput. 42, A477 (2020)) orthonormalizes the rows, exchanging only chi_out x chi_out Gram matrices; the second plain pass is skipped once the first sees an orthonormal input;
   • a second all-to-all gives each GPU its rows of the output core, which it also copies to the host itself (the D2H copy is no longer serialized on rank 0).

   Sketches too ill-conditioned for CholeskyQR (numerically rank-deficient: early in an evolution at cutoff 0, where the bonds are padded to chi_out) fall back to the Householder path. Rank deficiency spans runs of neighbouring sites, so after two fallbacks in a row the next 2, 4, ... (at most 16) sites skip the attempt. One GPU, and chi_out < 1024, are unchanged.

New pieces: utils.cholesky_qr, utils.solve_triangular (in _backend, since the array API has no triangular solve), DeviceGroup.alltoall. The docs (features/large-problems, contributing/testing) describe the new path.

Benchmarks

Leonardo, one node with 4x A100-SXM-64GB (NVLink), complex128, IBM 127-qubit kicked Ising Heisenberg MPO (src(U^dag, O, U) per sub-layer, theta_h = 8pi/32). Each pair below compresses the same input with the same seed.

Saturated operator (random MPO with bonds min(chi, 4^d), one compression per sub-layer):

chi_out	Householder on rank 0	CholeskyQR	ratio
512	2.05 s	2.53 s	1.25 (warm compressions; hence the 1024 threshold)
1024	8.2 s	5.9 s	0.72
2048	34.8 s (one Trotter step: 209 s)	23.8 s (143 s)	0.68
4096	179 s	134 s	0.75 (two runs: 0.76, 0.74)
2048, cutoff 1e-12	122 s	108 s	0.89 (the device SVD of R dominates)

Real evolution, 20 Trotter steps of <Z62> from Z62, chi = 1024, cutoff 0, each QR path in its own job: 969 s Householder vs 851 s CholeskyQR (0.88x). Steps 0-6 are rank-deficient everywhere and cost 1.02-1.04x, thanks to the backoff (1.14x without it). From step 8 on, the operator saturates, CholeskyQR takes over, and steps run 0.73-0.86x.

Accuracy

• Every output core is orthonormal to 1e-14 (chi 512-2048) and 1e-13 (4096), the same as Householder.
• Per compression, the two paths agree to rounding: the norms agree to 1e-14, <0|O|0> to 1e-11-1e-13 relative, and the Frobenius distance sits at the 3e-8 floor of the inner-product measurement.
• Over the 20-step evolution, <Z62> agrees to 13 digits while the compressions are exact (steps 0-6). From step 8 it drifts to about 3e-3 (0.1874 vs 0.1861 at step 20). This drift is not CholeskyQR's. The same Householder code on 2 GPUs instead of 4, where only the rounding differs, drifts from the 4-GPU run the same way: it diverges from step 6, reaches 1e-4 by step 8, and stays at 3e-4 to 5e-3 afterwards. Once truncation sets in, the 20-step chain amplifies any rounding-level change to the truncation-error scale. CholeskyQR's differences from the 4-GPU Householder run stay inside that envelope (3e-4 to 6e-3). Against a plain 1-GPU run, the difference at step 20 is 5.5e-4 for 4-GPU Householder, 8.8e-4 for 2-GPU Householder and 1.8e-3 for 4-GPU CholeskyQR; across the steps, all three lie in the same 1e-5 to 8e-3 band. For scale, chi = 1024 and 2048 differ by 4e-3 at step 20, and the 2D reference by 8e-2.
• Tests cover cholesky_qr against truncated_qr for condition numbers 1 to 1e12, with a cutoff, and on rank-deficient input (which must be declined on every rank). The multi-device equivalence tests run on both QR paths, plus two-GPU tests on hardware. All pass on Leonardo (tests/test_gpu_backend.py 74/74 with the linalg and group tests).

Notes for review

• Every rank factors bitwise-identical Gram sums (gathered and summed in rank order), so the ranks take the same decisions and get the same factors without a broadcast.
• alltoall and the host-side writes reuse the existing event-ordered peer copies; there is no new dependency (no NCCL).

🤖 Generated with Claude Code

[github-actions]

Test Results

    5 files  ±  0      5 suites  ±0   4m 38s ⏱️ +9s
  440 tests + 43    430 ✅ + 41  10 💤 +2  0 ❌ ±0 
2 102 runs  +203  2 081 ✅ +199  21 💤 +4  0 ❌ ±0 

Results for commit e2f010c. ± Comparison against base commit 3f7114a.

♻️ This comment has been updated with latest results.

[robertodr]

does it make sense to use the triangular solve in scipy? https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.solve_triangular.html#scipy.linalg.solve_triangular

or are only the tiny linear systems solved on the CPU?

[Panadestein]

Yes, I would say that for sizes where the triangular system arises, i.e. it gets triggered by the Cholesky QR, which is at BD>1024, running on CPU will be potentially very slow. And of course for small systems I don't think the difference will matter much between scipy and numpy, but I could measure. We don't have scipy as a strong dependency now.