Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

JAXMg: A multi-GPU linear solver in JAX

Roeland Wiersema

TL;DR

The paper tackles the challenge of scaling dense linear algebra across multiple GPUs within JAX, focusing on solving large SPD systems and eigenvalue problems that exceed single-GPU memory. It introduces JAXMg, a JAX interface to NVIDIA cuSOLVERMg via an XLA FFI, exposing JIT-compatible primitives for potrs, potri, and syevd on distributed GPUs using a 1D block-cyclic data layout. Key design elements include 1D cyclic data distribution and memory-management strategies for SPDM/MPDM modes, enabling in-JAX execution while delegating compute to cuSOLVERMg. Benchmark results on an 8-GPU node show improved scalability and higher throughput than native single-GPU cuSOLVER routines, enabling larger matrices and preserving JAX's composability in end-to-end workflows.

Abstract

Solving large dense linear systems and eigenvalue problems is a core requirement in many areas of scientific computing, but scaling these operations beyond a single GPU remains challenging within modern programming frameworks. While highly optimized multi-GPU solver libraries exist, they are typically difficult to integrate into composable, just-in-time (JIT) compiled Python workflows. JAXMg provides multi-GPU dense linear algebra for JAX, enabling Cholesky-based linear solves and symmetric eigendecompositions for matrices that exceed single-GPU memory limits. By interfacing JAX with NVIDIA's cuSOLVERMg through an XLA Foreign Function Interface, JAXMg exposes distributed GPU solvers as JIT-compatible JAX primitives. This design allows scalable linear algebra to be embedded directly within JAX programs, preserving composability with JAX transformations and enabling multi-GPU execution in end-to-end scientific workflows.

JAXMg: A multi-GPU linear solver in JAX

TL;DR

The paper tackles the challenge of scaling dense linear algebra across multiple GPUs within JAX, focusing on solving large SPD systems and eigenvalue problems that exceed single-GPU memory. It introduces JAXMg, a JAX interface to NVIDIA cuSOLVERMg via an XLA FFI, exposing JIT-compatible primitives for potrs, potri, and syevd on distributed GPUs using a 1D block-cyclic data layout. Key design elements include 1D cyclic data distribution and memory-management strategies for SPDM/MPDM modes, enabling in-JAX execution while delegating compute to cuSOLVERMg. Benchmark results on an 8-GPU node show improved scalability and higher throughput than native single-GPU cuSOLVER routines, enabling larger matrices and preserving JAX's composability in end-to-end workflows.

Abstract

Solving large dense linear systems and eigenvalue problems is a core requirement in many areas of scientific computing, but scaling these operations beyond a single GPU remains challenging within modern programming frameworks. While highly optimized multi-GPU solver libraries exist, they are typically difficult to integrate into composable, just-in-time (JIT) compiled Python workflows. JAXMg provides multi-GPU dense linear algebra for JAX, enabling Cholesky-based linear solves and symmetric eigendecompositions for matrices that exceed single-GPU memory limits. By interfacing JAX with NVIDIA's cuSOLVERMg through an XLA Foreign Function Interface, JAXMg exposes distributed GPU solvers as JIT-compatible JAX primitives. This design allows scalable linear algebra to be embedded directly within JAX programs, preserving composability with JAX transformations and enabling multi-GPU execution in end-to-end scientific workflows.
Paper Structure (6 sections, 3 figures)

This paper contains 6 sections, 3 figures.

Table of Contents

  1. Statement of need
  2. Software design
  3. 1D Cyclic Data Distribution
  4. Memory management
  5. Benchmark
  6. Acknowledgements

Figures (3)

  • Figure 1: Arrays are row-wise sharded and put into the round-robin 1D cyclic form illustrated here.
  • Figure 2: (Left) In SPMD mode, each GPU is controlled by a separate thread. Since these threads share the same virtual memory space, we can share device pointers between them. (Right) In MPMD mode, we use the cudaIpc API to transport the device pointers to a single array in process 0.
  • Figure 3: Benchmark comparing the native single-GPU JAX routines (which call cuSOLVERDn) to JAXMg. Reported timings include array allocation, which is negligible compared to the runtime. Larger tile sizes improve performance only once the problem size is sufficiently large, consistent with a GPU-utilization effect. Tile size has negligible impact for syevd, while potri shows a strong dependence on $T_A$. (a) Comparison of jaxmg.potrs with jax.scipy.linalg.cho_factor + jax.scipy.linalg.cho_solve for a float32 matrix. The largest solvable problem is N=524288, which utilizes >1 TB of memory. (b) Comparison of jaxmg.potri with jax.numpy.linalg.inv for a complex128 matrix. (c) Comparison of jaxmg.syevd with jax.numpy.linalg.eigh for a float64 matrix.