jaxsgp4: GPU-accelerated mega-constellation propagation with batch parallelism

Charlotte Priestley; Will Handley

cs.DCMar 29, 2026

jaxsgp4: GPU-accelerated mega-constellation propagation with batch parallelism

Authors: Charlotte Priestley, Will Handley

Organizations: Kavli Institute for Cosmology, University of Cambridge, Cambridge, UK · Institute of Astronomy, University of Cambridge, Cambridge, UK · PolyChord Ltd., London, UK

Abstract

As the population of anthropogenic space objects transitions from sparse clusters to mega-constellations exceeding 100,000 satellites, traditional orbital propagation techniques face a critical bottleneck. Standard CPU-bound implementations of the Simplified General Perturbations 4 (SGP4) algorithm are less well suited to handle the requisite scale of collision avoidance and Space Situational Awareness (SSA) tasks. This paper introduces jaxsgp4, an open-source high-performance reimplementation of SGP4 utilising the JAX library. JAX has gained traction in the landscape of computational research, offering an easy mechanism for Just-In-Time (JIT) compilation, automatic vectorisation and automatic optimisation of code for CPU, GPU and TPU hardware modalities. By refactoring the algorithm into a pure functional paradigm, we leverage these transformations to execute massively parallel propagations on modern GPUs. We demonstrate that jaxsgp4 can propagate the entire Starlink constellation (9,341 satellites) each to 1,000 future time steps in under 4 ms on a single A100 GPU, representing a speedup of $1500\times$ over traditional C++ baselines. Furthermore, we argue that the use of 32-bit precision for SGP4 propagation tasks offers a principled trade-off, sacrificing negligible precision loss for a substantial gain in throughput on hardware accelerators.

Figures & tables

Figure 1: Scaling performance comparison between JAX /GPU and C++/CPU implementations of the SGP4 algorithm, for a single satellite propagated to multiple times (left) and multiple satellites propagated to a single time (right). jaxsgp4 exploits the potential of modern GPUs to perform massively-parallel computations, exhibiting a flat scaling regime where increased workload does not increase wall-clock time until hardware saturation (top). The ‘break-even’ point for each GPU (Speedup=1, denoted by dashed line in bottom panels), at which the benefit of parallel computation for large batch sizes overcomes the initial dispatch overhead, occurs at batch sizes of

\sim

300-500. In the bottom panels, we also specify the maximum speedup achieved by jaxsgp4 over the C++ baseline. This occurs in the linear hardware saturation regime, at which point the GPUs are operating at maximum computational efficiency.

Figure 2: GPU-accelerated JAX vs C++ SGP4 performance comparison. Blue indicates JAX is faster, Red indicates C++ is faster. The standard C++ implementation is inherently single-threaded and was run on a standard CPU while the JAX benchmark was run on an NVIDIA A100 GPU. Log-spaced axes show the relative performance of each implementation across different scales.

Figure 3: Relative error accumulation over two weeks running jaxsgp4 with 32 bit precision versus the standard SGP4 implementation which uses 64 bit precision. The median and 5 th - 95 th percentile regions are highlighted in blue. The dashed line depicts a conservative lower bound on physical error growth due to limitations of the SGP4 model itself. jaxsgp4 errors are order 1e-9, 1e-12 for position and velocity respectively (floating point) when the same analysis is performed using 64 bit precision.

GPU Acceleration

GPU acceleration uses parallel processors, specialized kernels, and memory-aware systems to speed computation across AI and scientific workloads.

See all 169 papers

WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled $LDL^\top$ Factorization
Jul 13, 2026Henrik Hose, Se Hwan Jeon, Charles Khazoom +2GPU Kernel Optimization GPU Acceleration
Institute for Data Science in Mechanical Engineering (DSME), RWTH Aachen University, Germany · Department of Mechanical Engineering, Massachusetts Institute of Technology, Cambridge MA, USA.
Accelerated Fourier SAT (AFSAT): Fully Realising a GPU-based Symmetric Pseudo-Boolean SAT Solver
Jun 4, 2026Cody J Christopher, Charles GrettonGPU Acceleration High-Performance Computing
School of Computing, Australian National University, Canberra, Australia.
Time-Optimal Collision Avoidance Via a Greedy Polynomial Backward Sweep
May 31, 2026Zeno Pavanello, Frank De Veld, Roberto ArmellinCollision Avoidance Trajectory Optimization
Department of Aerospace Science and Technology (DAER), Politecnico di Milano, via Privata Giuseppe La Masa 34, Milano, 20156, Italy · Netherlands Aerospace Centre (NLR), Anthony Fokkerweg 2, Amsterdam, 1059 CM, The Netherlands · Te Pūnaha ¯Atea Auckland - Space Institute (TPA-SI), The University of Auckland, 20 Symonds Street, Auckland, 1010, New Zealand

Jul 13, 2026J/K move · Enter open · S save

WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled $LDL^\top$ Factorization

Henrik Hose, Se Hwan Jeon, Charles Khazoom, Sangbae Kim+1

Institute for Data Science in Mechanical Engineering (DSME), RWTH Aachen University, Germany · Department of Mechanical Engineering, Massachusetts Institute of Technology, Cambridge MA, USA.

GPU Kernel Optimization GPU Acceleration Nonlinear MPC Alternating Direction Method of Multipliers Quadratic Programming

Open paper AskPDF arXiv

This paper introduces numerical optimizations for maximizing throughput on GPU when solving large batches (10,000 to over 100,000) of sequential quadratic programming (SQP) iterations, where all problems have the same structure. The optimizations are implemented in a toolbox WarpMPC for model-predictive control (MPC) in JAX and Warp. Based on the insight that all MPC problem instances in a batch share the same sparsity in time, cost, and constraints, we propose unrolling sparse linear factorizations and solves, which dominate alternating direction method of multipliers (ADMM) solver runtime. We avoid memory access bottlenecks and wasting computations via optimized memory layout, padding-reducing segmentation of the unrolled factorization, and dependency level scheduled backsolves, additionally accelerating sensitivity computation. We achieve throughputs of 8,000 to 250,000 SQP iterations per second on nonlinear cartpole, quadrotor, and humanoid robot benchmarks, outperforming baselines by 3

\times

to 25

\times

. We illustrate practical usefulness by synthesizing a dataset and training a neural network approximation of an MPC in under 4 minutes that stabilizes a nano quadrotor in hardware experiments.

Abstract

GPU Acceleration

Explore similar work

WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled LDL⊤LDL^\topLDL⊤ Factorization

Accelerated Fourier SAT (AFSAT): Fully Realising a GPU-based Symmetric Pseudo-Boolean SAT Solver

Time-Optimal Collision Avoidance Via a Greedy Polynomial Backward Sweep

WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled $LDL^\top$ Factorization