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

2026-07-13

AI Summary

This work addresses the efficient GPU-based solution of large-scale homogeneous model predictive control (MPC) problems by proposing an ADMM-based framework that integrates unrolled sparse LDL^⊤ factorization, a fill-reducing partitioning strategy, and a dependency-aware back-substitution algorithm to jointly optimize memory layout and computation scheduling. Implemented in JAX and Warp, the system significantly accelerates SQP iterations and sensitivity computations, achieving 8,000 to 250,000 SQP iterations per second on nonlinear benchmarks including inverted pendulums, quadrotors, and humanoid robots—yielding a 3–25× speedup over baselines. The approach enables end-to-end MPC dataset generation and neural network training within four minutes and has been successfully deployed on a real quadrotor platform.

Abstract

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.

Problem

• Model Predictive Control
• Large-Batch Optimization
• GPU Acceleration
• Sequential Quadratic Programming
• Sparsity Exploitation

Innovation

• unrolled LDL^T factorization
• large-batch MPC
• GPU acceleration
• ADMM solver
• sparse linear algebra