WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled $LDL^\top$ Factorization
WarpMPC uses ADMM with unrolled LDL⊤ factorization for large-batch GPU MPC, achieving 8,000–250,000 SQP iterations/sec.
Key Findings
Methodology
This paper introduces a GPU-optimized large-batch MPC solver based on ADMM and unrolled sparse LDL⊤ factorization. By exploiting shared sparsity patterns across batch instances, it unrolls factorization and backsolve kernels, avoiding memory bottlenecks. Using CasADi for symbolic derivation and automatic translation to JAX, the framework ensures efficient sparse computations. The core innovations include fixed sparsity pattern kernels, segmented kernel design, and dependency level scheduling, which collectively enable over 100,000 parallel SQP iterations per second. The approach combines problem-specific symbolic analysis with low-level kernel optimization, significantly reducing redundant computations and memory access overheads.
Key Results
- Benchmarking on nonlinear Cartpole, quadrotor, and humanoid models shows throughput of 8,000 to 250,000 SQP iterations/sec, outperforming existing GPU solvers by 3× to 25×. The optimized memory layout and segmentation strategies effectively minimize padding and idle computation, leading to high GPU utilization.
- In hardware experiments, a neural network trained via this framework in under 4 minutes successfully stabilized a nano quadrotor, demonstrating real-world applicability. The large batch capability enables rapid data synthesis and model training, accelerating reinforcement learning workflows.
- Compared to traditional iterative solvers like cuDSS or dense reformulations, WarpMPC maintains high throughput for large batch sizes, especially in high-dimensional, long-horizon nonlinear MPC problems, marking a significant step forward in real-time control scalability.
Significance
This work addresses a critical bottleneck in deploying MPC for high-dimensional robotic systems and learning-based control. By leveraging GPU parallelism with problem-specific sparsity exploitation, it enables real-time, large-scale batch processing that was previously infeasible. The approach bridges the gap between classical optimization and modern machine learning, facilitating faster training, simulation, and deployment of complex control policies. It opens new avenues for integrating MPC into deep learning pipelines, making high-dimensional, nonlinear control more accessible and scalable for industry and research. The framework’s ability to handle hundreds of thousands of parallel solves on a single GPU marks a paradigm shift in control optimization, promising transformative impacts in autonomous robotics, automotive, and industrial automation.
Technical Contribution
The paper introduces a fixed sparsity pattern approach for LDL⊤ factorization, combined with problem-specific kernel generation and layered scheduling, to maximize GPU throughput. It automates symbolic analysis from CasADi, translating it into optimized JAX code, enabling high-performance sparse linear algebra. The segmentation of kernel operations reduces padding overhead, while dependency level scheduling minimizes sequential bottlenecks in backsolve steps. These innovations collectively enable GPU-based batch processing of over 100,000 MPC problems per second, a significant leap from existing iterative or dense methods. The framework also supports sensitivity analysis and differentiability, facilitating integration with learning algorithms.
Novelty
This is the first work to unify unrolled sparse LDL⊤ factorization with ADMM for large-batch MPC on GPU, leveraging fixed sparsity structures for maximum efficiency. Unlike prior methods that rely on iterative solvers or dense reformulations, this approach exploits problem structure at the kernel level, achieving unprecedented throughput. The automatic symbolic-to-numeric pipeline and layered scheduling constitute a novel engineering paradigm for sparse linear algebra in high-performance computing, setting a new standard for real-time, large-scale MPC.
Limitations
- The approach assumes fixed sparsity patterns, limiting applicability to problems with dynamic or adaptive structures. Recomputing symbolic analysis for changing models incurs overhead.
- Kernel precompilation and storage demand significant memory, especially for very high-dimensional systems, potentially impacting scalability.
- Current implementation targets single GPU environments; multi-GPU or distributed extensions require further development to handle communication overhead and synchronization.
Future Work
Future directions include developing adaptive kernel generation for dynamic models, extending the framework to multi-GPU systems, and integrating learning-based models for end-to-end control. Additionally, exploring automatic structure detection and real-time kernel recompilation could broaden applicability to more complex, changing environments. Further research aims to reduce memory footprint and improve scalability, enabling deployment in larger robotic systems and industrial settings.
AI Executive Summary
This paper presents WarpMPC, a GPU-optimized large-batch MPC solver based on ADMM and unrolled sparse LDL⊤ factorization. By exploiting shared sparsity patterns across instances, the framework unrolls kernels for factorization and backsolve, avoiding memory bottlenecks and redundant computations. Using symbolic analysis from CasADi and automatic translation to JAX, the authors generate highly efficient, problem-specific kernels that leverage fixed sparsity structures. Core innovations include segmented kernel design and dependency level scheduling, which significantly reduce padding and sequential dependencies, enabling over 100,000 parallel SQP iterations per second on a single GPU. Experimental results on nonlinear Cartpole, quadrotor, and humanoid models demonstrate throughput improvements of 3× to 25× over existing GPU solvers, with real-world hardware validation showing rapid neural network training and stable quadrotor control in under 4 minutes. This breakthrough addresses a longstanding bottleneck in deploying MPC for high-dimensional, real-time robotic applications, bridging classical optimization with modern machine learning. The approach’s scalability and efficiency open new horizons for autonomous systems, industrial automation, and reinforcement learning, where large-scale parallel optimization is essential. Despite current limitations related to fixed sparsity assumptions and single-GPU deployment, future work aims to extend the framework’s adaptability, multi-GPU support, and integration with learning-based control, promising a transformative impact on the field of real-time optimal control.
Deep Analysis
Background
近年来，模型预测控制（MPC）在机器人、自动驾驶等领域得到广泛应用，但高维非线性系统的实时求解仍面临巨大挑战。传统CPU求解器如OSQP、Gurobi在单机性能上已取得一定突破，但难以满足大规模批处理和实时控制需求。GPU的高并行能力提供了潜在解决方案，但现有GPU稀疏线性求解器多依赖迭代方法，受限于分支和存储不规则性，吞吐率有限。近年来，结合ADMM的稀疏LDL⊤分解成为研究热点，旨在利用稀疏结构提升求解效率，但在GPU上的实现仍受制于内存访问瓶颈和核函数设计。本文在此背景下，提出未展开的稀疏LDL⊤分解策略，结合自动符号分析和层级调度，突破GPU稀疏线性求解瓶颈，推动大规模批处理在实时控制中的应用。
Core Problem
核心问题在于如何在GPU上高效处理大规模稀疏线性系统，尤其在批处理场景中。传统方法受限于内存访问不规则和分支，导致吞吐率低。现有GPU稀疏线性求解器难以满足每秒数万次迭代的需求，限制了大规模非线性MPC的实时应用。如何设计一种结构固定、计算高效、可扩展的求解框架，是推动GPU在控制中的关键难题。
Innovation
本文的创新点包括：1）提出未展开的稀疏LDL⊤核函数，避免重复因子化，提升效率；2）利用固定稀疏结构，自动生成核函数，减少内存访问和无效计算；3）引入分段核函数设计，减少padding带来的浪费；4）采用层级调度，显著降低依赖层数，加快后向求解速度。这些创新结合在一起，使GPU可以在大批量场景中实现每秒数万到数十万次的SQP迭代，远超现有方案。
Methodology
- �� 利用CasADi自动生成稀疏表达式，转换为JAX代码，确保结构一致性。• 预处理固定稀疏模式的符号信息，包括核函数、分段策略和依赖层级。• 在GPU上实现未展开的LDL⊤核函数，结合存储优化，减少内存访问。• 采用分段核设计，根据列非零数动态调整核宽度，避免无效计算。• 通过层级调度，减少依赖层数，支持并行后向求解。• 在每次SQP迭代中，利用预编译核函数快速完成线性系统分解和回代。• 结合线搜索策略，确保收敛性和鲁棒性。
Experiments
在非线性Cartpole、四旋翼和 humanoid机器人任务中，采用不同规模的测试集验证性能。对比传统GPU稀疏线性求解器，展示吞吐率提升3到25倍。通过硬件实验验证训练的神经网络模型在微型四旋翼上的稳定性，验证方法的实用性。参数设置包括批量大小、核段数、调度层级等，确保在不同场景下的性能表现。采用标准指标如每秒SQP迭代数、求解时间和控制精度进行评估。
Results
实验结果显示，WarpMPC在非线性控制任务中达到了每秒8,000至25万次SQP迭代，明显优于传统GPU稀疏求解器。存储布局和核分段优化显著减少了无效计算，提升GPU利用率。硬件验证中，神经网络近似模型在不到4分钟内完成训练，成功实现微型四旋翼的稳定控制。这些结果证明了方法在高维、长时域控制中的优越性能和实用价值。
Applications
该技术适用于机器人自主控制、自动驾驶、工业自动化等场景，特别是在需要大规模批量并行优化的实时系统中。通过GPU实现高吞吐率，降低硬件成本，提升系统响应速度，为深度学习与控制的融合提供技术基础。未来可结合学习模型，进一步优化端到端的控制策略，实现智能自主系统。
Limitations & Outlook
当前方法依赖于问题的固定稀疏结构，动态变化的模型可能需要重新分析和核函数重构。核函数的预编译和存储在高维复杂系统中可能带来较大开销。多GPU或分布式环境的扩展尚未实现，未来需解决通信和同步问题。此外，极端高维系统的存储和计算成本仍是挑战。
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.