Beyond Single-GPU: Scaling PDLP to Distributed Multi-GPU Systems

TL;DR

This work presents a distributed implementation of the Primal-Dual Hybrid Gradient algorithm, built upon cuPDLPx, which achieves strong scalability and high performance while preserving full FP64 numerical accuracy.

Abstract

In this work, we present a distributed implementation of the Primal-Dual Hybrid Gradient (PDHG) algorithm for solving massive-scale linear programming (LP) problems. Although PDHG-based solvers have shown strong performance on single-node GPU architectures, their applicability to industrial-scale instances is often limited by GPU memory capacity and computational throughput. To overcome these challenges, we extend the PDHG framework to a distributed-memory setting via a practical two-dimensional grid partitioning of the constraint matrix, enabling scalable execution across multiple GPUs. Our implementation leverages the NCCL communication backend to efficiently synchronize primal-dual updates across devices. To improve load balance and computational efficiency, we introduce a block-wise random shuffling strategy combined with nonzero-aware data distribution, and further accelerate computation through fused CUDA kernels. By distributing both memory and computation, the proposed framework not only overcomes the single-GPU memory bottleneck but also achieves substantial speedups by exploiting multi-GPU parallelism with relatively low communication overhead. Extensive experiments on standard LP benchmarks, including MIPLIB and Hans' instances, as well as large-scale real-world datasets, show that our distributed implementation, built upon cuPDLPx, achieves strong scalability and high performance while preserving full FP64 numerical accuracy.

Figures (2)

• Figure 1: Distributed memory layout under 2D grid partitioning. Left: A logical $2 \times 2$ device mesh illustrating the global topology. Right: Local data stored on the device at grid coordinate $(i,j)$. The constraint matrix is stored as block $A_{[i,j]}$, primal variables $x_{[j]}$, and dual variables $y_{[i]}$. This layout enables local computation of $A_{[i,j]} x_{[j]}$ and $A_{[i,j]}^\top y_{[i]}$ prior to reduction.
• Figure 2: Comparison of matrix permutation strategies for distributed sparse optimization. Left: Natural ordering preserves local dense structures and favors memory coalescing, but causes severe load imbalance across devices. Middle: Full random permutation achieves uniform load distribution but destroys local sparsity, leading to irregular memory access and degraded SpMV performance. Right: Block-wise random permutation balances load while preserving dense micro-structures, enabling scalable and efficient GPU execution under 2D partitioning.