feat: multi-GPU belief propagation backend

[Panadestein]

Summary

Adds a MultiGPUBeliefPropagationCache that distributes a single BP instance across multiple CUDA devices. The implementation is not an embarrassingly parallel fan-out of independent jobs: all devices jointly maintain one shared BP fixed point and one evolving quantum state. Gate application is unchanged and still runs sequentially; only the BP sweep is parallelised.

Algorithm

Partitioning

The tensor network graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ is partitioned with KaHyPar into $N_\text{dev}$ device-local subsets $\mathcal{V}^{(p)}$. Two objectives are exposed:

• memory-balanced (default): vertex weight = degree, spreads high-coordination hot spots evenly across devices.
• min-cut (uniform weights): minimises the number of boundary edges and therefore the inter-device communication volume.

Each device holds in VRAM the tensors ${T_v : v\in\mathcal{V}^{(p)}}$ and every directed message incident to $\mathcal{V}^{(p)}$. Cut edges are recorded as the messages that device must produce and ship.

Distributed sweep

The undirected edges are partitioned into proper color classes so that no two edges in a class share a vertex. One BP iteration then:

1. Parallel phase (Jacobi within a color): all devices update their owned messages for the current color simultaneously via Threads.@spawn; no write conflicts because the local environments are vertex-disjoint.
2. Halo exchange (Gauss–Seidel across colors): a peer-to-peer (P2P) NVLink/PCIe transfer ships each freshly computed cut message directly to its consumer device before the next color begins.

Gate application

Gates are applied sequentially on a chosen home device. For a two-site gate spanning a partition boundary, the partner tensor and surrounding environment messages are staged onto the home device, the gate is absorbed, a truncated QR/SVD is performed, the resulting singular-value spectrum is written as the new seed message on the gate edge and broadcast to all devices tracking it. At large bond dimension, where a single QR/SVD can transiently exceed device memory, the factorization falls back transparently to host RAM and the result is moved back to the GPU.

Validation and performance

Validated against the single-GPU reference on a 56-qubit heavy-hex circuit with 10675 gates. In double precision the two implementations reach the same BP fixed point to within $10^{-8}$ in message fidelity; in the single-precision production regime observables agree to $10^{-5}$ – $10^{-4}$ across $\chi=50$ – $500$, consistent with the BP convergence floor rather than any algorithmic bias.

Strong-scaling speedup (4× A100-64 GB vs. single A100, same circuit):

$\chi$	speedup
50	≈ 1.6×
200	≈ 3.1×
400	≈ 6.9×
500	≈ 6.3×

The single-GPU run exhausts memory at $\chi\approx550$; the four-GPU run continues, reaching $\chi=700$ in roughly one hour of wall time. At eight H200 GPUs (140 GB each) the same circuit has been run at $\chi=980$.

New public interface

# Construct from an existing single-device BP cache
cache = MultiGPUBeliefPropagationCache(BeliefPropagationCache(ψ); n_devices = 4)

# BP sweep (dispatches to the "multigpu_bp" update algorithm)
cache = update(cache; maxiter, tolerance, α)

# Gate application (same layer-level interface as the single-device path)
cache, errors = apply_gates(layer, cache; apply_kwargs)

# Collect the distributed state back to a single host for measurement
host = collect_to_cpu!(cache)

The partitioning objective is selected via the partition_alg keyword ("memory_balanced" or "min_cut").

__

CC @fpietra @davidemateria feel free to jump into the discussion, review, and answer questions.

[JoeyT1994]

Thanks! This is very impressive and useful.

As this relies on CUDA.jl which I don't want as a fixed dependency, this will probably need to be written as a package extension (which then only loads if CUDA.jl is loaded), see: https://discourse.julialang.org/t/quick-tutorial-on-package-extensions/130923.

In the short term we could have this as a branch off of main which I can then accept as a PR and people can load into easily / use for the time being and I can even have the README point too

[Panadestein]

Thanks! This is very impressive and useful.

As this relies on CUDA.jl which I don't want as a fixed dependency, this will probably need to be written as a package extension (which then only loads if CUDA.jl is loaded), see: https://discourse.julialang.org/t/quick-tutorial-on-package-extensions/130923.

In the short term we could have this as a branch off of main which I can then accept as a PR and people can load into easily / use for the time being and I can even have the README point too

Thanks!

Makes sense to make it an ext, I even tried when starting the PR but found it cumbersome to use in supercomputers due to the environment build. The rewrite shouldn't be that complication from the current point I think. I'll push that to this PR. Short term feel free to make it a branch.

[Panadestein]

Hi @JoeyT1994, the refactoring to extension is done, so this may not need the interim branch DistributedBP after all.

Everything CUDA/KaHyPar-specific now lives in a single ext/TensorNetworkQuantumSimulatorCUDAExt.jl, with both packages in [weakdeps]. src/ only keeps the MultiGPUBeliefPropagationCache type and its interface, so nothing extra is loaded or installed unless the user does using CUDA, KaHyPar.

I validated it with the new tests and the 56-qubit heavy-hex benchmark from the description over BD 50–400 vs single GPU, same errors as before. So the refactor is behaviour-neutral and adds no dependencies for non-GPU users. Would you be happy to retarget this at main?