Redesigning GROMACS Halo Exchange: Improving Strong Scaling with GPU-initiated NVSHMEM

2025/09/25 by Doijade, Mahesh, Alekseenko, Andrey, Brown, Ania +2
#Computational Physics (physics.comp-ph) #Distributed #FOS: Computer and information sciences #FOS: Physical sciences #Parallel #Performance (cs.PF) #and Cluster Computing (cs.DC)

paper · doi:10.48550/arxiv.2509.21527

Abstract

Improving time-to-solution in molecular dynamics simulations often requires strong scaling due to fixed-sized problems. GROMACS is highly latency-sensitive, with peak iteration rates in the sub-millisecond, making scalability on heterogeneous supercomputers challenging. MPI's CPU-centric nature introduces additional latencies on GPU-resident applications' critical path, hindering GPU utilization and scalability. To address these limitations, we present an NVSHMEM-based GPU kernel-initiated redesign of the GROMACS domain decomposition halo-exchange algorithm. Highly tuned GPU kernels fuse data packing and communication, leveraging hardware latency-hiding for fine-grained overlap. We employ kernel fusion across overlapped data forwarding communication phases and utilize the asynchronous copy engine over NVLink to optimize latency and bandwidth. Our GPU-resident formulation greatly increases communication-computation overlap, improving GROMACS strong scaling performance across NVLink by up to 1.5x (intra-node) and 2x (multi-node), and up to 1.3x multi-node over NVLink+InfiniBand. This demonstrates the profound benefits of GPU-initiated communication for strong-scaling a broad range of latency-sensitive applications.

Citations

• Scaling of the GROMACS Molecular Dynamics Code to 65k CPU Cores on an HPC Cluster
• Intel(R) SHMEM: GPU-initiated OpenSHMEM using SYCL
• Understanding GPU Triggering APIs for MPI+X Communication
• GROMACS on AMD GPU-Based HPC Platforms: Using SYCL for Performance and Portability
• Exploring Fully Offloaded GPU Stream-Aware Message Passing
• Two-link Staggered Quark Smearing in QUDA
• MGG: Accelerating Graph Neural Networks with Fine-grained intra-kernel Communication-Computation Pipelining on Multi-GPU Platforms
• Fast and Scalable Sparse Triangular Solver for Multi-GPU Based HPC Architectures
• ACEMD: Accelerating bio-molecular dynamics in the microsecond time-scale
• GROMACS 3.0: a package for molecular simulation and trajectory analysis

Related

• GROMACS on AMD GPU-Based HPC Platforms: Using SYCL for Performance and Portability
• A Comparison of the Performance of the Molecular Dynamics Simulation Package GROMACS Implemented in the SYCL and CUDA Programming Models
• Scaling of the GROMACS Molecular Dynamics Code to 65k CPU Cores on an HPC Cluster
• GROMACS Unplugged: How Power Capping and Frequency Shapes Performance on GPUs
• Cyclic Data Streaming on GPUs for Short Range Stencils Applied to Molecular Dynamics