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

2025-09-25

AI Summary

To address the limited strong scaling of GROMACS on heterogeneous supercomputers—caused by MPI communication latency (sub-millisecond) and low GPU utilization—this work proposes a novel GPU-kernel-initiated communication paradigm. We redesign the domain-decomposition halo exchange algorithm to tightly integrate data packing and NVSHMEM communication within GPU kernels, enabling aggressive overlap of computation and communication. Leveraging NVLink’s asynchronous copy engine and kernel fusion techniques, we bypass the CPU-centric MPI communication path, thereby significantly reducing latency and improving bandwidth efficiency. Experimental results demonstrate 1.5× speedup on single-node NVLink systems, 2× on multi-node NVLink clusters, and 1.3× on hybrid NVLink–InfiniBand networks. This approach effectively overcomes the heterogeneous scalability bottleneck of traditional MPI-based molecular dynamics simulations.

Technology Category

• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems
• Machine Learning: Scalability of ML Systems
• Multiagent Systems: Agent Communication

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Reduces MPI latency in GROMACS molecular dynamics simulations
• Improves strong scaling via GPU-initiated NVSHMEM communication
• Enhances communication-computation overlap for latency-sensitive applications

Innovation

• GPU-initiated NVSHMEM redesigns halo exchange
• Kernel fusion enables data packing and communication overlap
• Asynchronous NVLink copy engine optimizes latency bandwidth