GPU-Initiated Halo Exchange Cuts CPU Bottlenecks in GROMACS

NVIDIA published a guide on July 9, 2026, describing how GROMACS molecular dynamics simulations can replace CPU-orchestrated MPI communication with GPU-initiated remote memory access using the NVSHMEM library. The work targets the halo exchange, the algorithm that shares boundary atom data between neighboring GPU domains, which the authors identify as a scalability bottleneck at sub-millisecond timestep rates.

The CPU bottleneck in halo exchange

GROMACS simulations advance in steps of 100 to 200 microseconds across multiple GPUs, a regime that demands extremely low kernel latency. The traditional MPI-based halo exchange forces the GPU to pause while the CPU orchestrates data transfers and then signals the GPU to resume. In a 3D decomposition with one pulse per dimension, this handoff repeats across all three spatial dimensions, consuming more than 50% of total CPU wall time at peak iteration rates.

The existing implementation uses a staged forwarding mechanism where data relays through intermediate ranks within each spatial dimension across one or more communication steps, or pulses. Each pulse begins with a pack kernel that gathers boundary atoms into a contiguous send buffer, creating a chain of dependencies between communication phases in the order Z to Y to X. The CPU-side control flow wraps each pulse with blocking synchronizations: one before the MPI invocation to guarantee the pack kernel completed, and one before the next pulse can consume forwarded data. With one pulse per dimension, that is six blocking waits per timestep for coordinates and six more for forces, 12 in total.

Fusing pack, put, and wait into one kernel

The first design folds packing, the remote put, and the completion wait into a single kernel per communication pulse. The host still issues launches in pulse order but drops the CPU-GPU barriers between pulses, with correct execution order guaranteed by the CUDA stream alone. All GPU thread blocks arrive at a device-scoped barrier after packing, and block 0 waits on that barrier before issuing the NVSHMEM put operation that transfers packed data to the peer and sets a signal on the peer to confirm completion. The receiver waits on that signal before consuming the data.

This version still processes pulses sequentially because stream ordering enforces it, and each pulse incurs a separate kernel launch. A further optimization addresses this with fine-grained per-pulse dependencies expressed directly in a fused kernel, replacing coarse phase-level serialization with per-pulse signals. This allows independent halo-exchange data to transmit immediately while dependent data waits only on those fine-grained signals.

Interconnect-aware transport

When the peer GPU is NVLink-connected, the design probes for a non-null device pointer using the NVSHMEM pointer query and branches accordingly. With NVLink, the kernel packs directly into the peer's coordinate array, eliminating the separate global memory round-trip of staging into a local buffer and issuing a put. Without NVLink, the kernel falls back to packing into a local staging buffer and using NVSHMEM put operations over InfiniBand or Slingshot.

The guide reports benchmarking on the NVIDIA Eos supercomputer showing up to 1.5x better intra-node performance than GPU-aware MPI on four to eight GPUs. Gains extend to 2x on GB200 NVL72 multi-node NVLink clusters and 1.3x on standard InfiniBand clusters. The authors suggest trying the GPU-initiated halo exchange in GROMACS 2026 to measure strong-scaling gains and exploring NVSHMEM to identify CPU-orchestrated communication on the critical path. Full design and benchmark details are in an SC'25 workshop paper.