Blog Tags GitHub
Back to blog
system-optimization kernel-optimization distributed-training communication gpu-architecture

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

GPU-initiated NVSHMEM redesign of GROMACS domain decomposition halo exchange, fusing data packing and communication to eliminate CPU-GPU sync bottlenecks and improve strong scaling by up to 2x

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

一、论文概述

项目内容
标题Redesigning GROMACS Halo Exchange: Improving Strong Scaling with GPU-initiated NVSHMEM
作者Mahesh Doijade (NVIDIA), Andrey Alekseenko (KTH), Ania Brown (NVIDIA), Alan Gray (NVIDIA), Szilard Páll (KTH)
机构NVIDIA / KTH Royal Institute of Technology
论文https://arxiv.org/abs/2509.21527
代码未公开（已 upstreaming 到 GROMACS）
发布PAW-ATM Workshop at SC 2025, Sept 25, 2025
许可CC BY-SA 4.0
类别cs.DC (Distributed Computing)

二、核心思想

问题定义

分子动力学（MD）模拟的强扩展面临严重的延迟挑战。GROMACS 是高度延迟敏感的应用，峰值迭代率在亚毫秒级（100-200 µs/step）。在异构超级计算机上，MPI 的 CPU 中心化特性在 GPU 驻留应用的关键路径上引入了额外延迟：

  1. CPU-GPU 同步开销：每个 halo exchange 脉冲都需要 GPU kernel 完成 → CPU 等待 → MPI 发送 → CPU 等待 → GPU 接收数据 → GPU kernel 启动
  2. MPI 无法从 GPU kernel 调用：MD 步骤无法完全 GPU 驻留
  3. 序列化脉冲：z → y → x 三维通信阶段必须串行执行，每步都需 CPU-GPU 同步
  4. 内核启动延迟累积：约 20 次启动 API 调用 + 30 次事件管理调用，累积成本可达 CPU 壁时间的 50%

解决方案概述

本文提出了一种基于 NVSHMEM 的 GPU kernel 发起的 GROMACS 域分解 halo exchange 算法重构。核心创新包括：

  • GPU 驻留设计：所有通信控制路径卸载到 GPU，消除 CPU-GPU 同步瓶颈
  • Fused 内核：将数据打包（pack）和通信融合到单个 kernel 中，最大化通信并发
  • 依赖分区策略：将索引映射分为 depOffset 以上（独立数据）和以下（依赖数据），独立数据立即打包发送，依赖数据仅等待必要的脉冲
  • TMA 异步复制引擎：NVLink 路径使用 Tensor Memory Accelerator 实现 warp 级别的异步远程内存操作
  • 自适应互连策略：NVLink 路径使用零拷贝直接远程访问；InfiniBand 路径使用粗粒度的 put-with-signal 和 staging buffer
  • 融合信号与通信：在单个 kernel 内完成数据打包/传输和接收方通知，不同于典型 NVSHMEM 实现

模型在 NVIDIA EOS（576 个 DGX H100 节点）和 GB200 NVL72 系统上验证，NVLink 上 intra-node 最高提升 1.5x，multi-node 最高提升 2x，NVLink+InfiniBand 上提升 1.3x。

三、技术架构

GROMACS 域分解 Halo Exchange

GROMACS GPU-resident MPI schedule

GROMACS 使用中性领土域分解（Neutral-Territory Domain Decomposition）：

  • 模拟盒子划分为空间区域（domains），每个 MPI rank 负责其分配区域内原子的受力计算
  • 非键相互作用由 cutoff 限制范围，边界原子需要相邻域的坐标
  • 采用八壳转发方法（eighth-shell method）：边界数据通过中间 rank 转发，而非直接发送到所有消费者

通信阶段：

  1. np(z) 个脉冲在 z 方向
  2. np(y) 个脉冲在 y 方向
  3. np(x) 个脉冲在 x 方向

转发机制确保 np(x) × np(y) × np(z) - 1 个相邻 rank 的数据交换仅需 np(x) + np(y) + np(z) 步。

NVSHMEM vs MPI 调度对比

GROMACS GPU-resident NVSHMEM schedule

MPI 方案（Fig. 1）：

  • CPU MD 时间步进循环中，紫色框 = kernel 启动，品红色 = GPU pack/unpack 等待，红色 = MPI 通信等待
  • 每个脉冲都需要 CPU-GPU 同步，关键路径暴露大量等待

NVSHMEM 方案（Fig. 2）：

  • CPU 时间步进循环中的 kernel 启动可与计算完全重叠
  • GPU stream 中本地/非本地工作完全重叠，halo exchange 不在关键路径上

核心数据结构（Algorithm 1）

PulseData (per pulse):
  sendRank, recvRank, sendSize, recvSize
  atomOffset, coordShift, indexMap
  remoteCoordDst, remoteForceDst, remoteForceSrc
  sendBuf, recvBuf, blocksForPulse

CommContext:
  totalPulses, signal[p], sigVal
  // Runtime path: NVLink when remoteCoordDst/remoteForceSrc non-null
  // Otherwise InfiniBand with NVSHMEM put-with-signal

Signals: device-visible counters indexed by pulse id
Buffers: coords, forces, per-pulse staging buffers
Barriers: indexMapLoadBarrier, forceBufLoadBarrier
Helper: isNVLinkAccess(ptr) ≔ (ptr ≠ null)

GPU 驻留时间步骨架（Algorithm 2）

// Fused NVSHMEM halo exchange — single launch per coord/force exchange
launch Local Non-bonded F on local stream
launch FusedPackCommX on non-local stream
    // Coordinate halo: pack, NVLink TMA store or InfiniBand put-to-remote, then notify
launch Bonded F on non-local stream
launch Non-Local Non-bonded F on non-local stream
launch FusedCommUnpackF on non-local stream
    // Forces halo: NVLink get-from-remote or InfiniBand put-to-remote; notify; then unpack
launch Integration, constraints on update stream
launch other per-step work: ReduceF, Rolling prune, Clear buffers
proceed to next step without CPU-GPU sync (GPU-resident schedule)

坐标 Halo 交换：FusedPackCommX（Algorithm 3）

// Kernel: FusedPackCommX
// blocks with same blockIdx.y cooperate on current pulse p
for p ← blockIdx.y to ctx.totalPulses-1 step gridDim.y do
    meta ← pulses[p];
    outBuf ← (isNVLinkAccess(meta.remoteCoordDst))
              ? shared-mem scratch : meta.sendBuf;
    needLastBlock ← (meta.blocksForPulse > 1)

    packWithDeps(meta, ctx, p, outBuf)    // Algorithm 4

    if isNVLinkAccess(meta.remoteCoordDst) then
        // As chunks are ready: TMA async store
        tma_st(outBuf → meta.remoteCoordDst + meta.atomOffset)
            // issued by thread 0 of each warp (warp leader)
    end if

    syncAndCommWithDeps(DATA, needLastBlock, ctx, meta, p, pulses,
                        outBuf, null, coords)  // Algorithm 5
end for

依赖感知打包（Algorithm 4）

// Device Function: packWithDeps
pack all entries with meta.indexMap[i] < meta.depOffset into outBuf
    // Apply meta.coordShift as needed

elect a leader at chosen scope
    // warp leader for NVLink path, threadIdx.x=0 otherwise

leader: for k ← p-1 down to firstDependentPulse
    acquire_wait(ctx.signal[k] == ctx.sigVal)

scope_barrier()
    // SYNCWARP() for NVLink; SYNCTHREADS() for InfiniBand

pack remaining entries with meta.indexMap[i] ≥ meta.depOffset into outBuf

关键洞察：通过将索引映射在 depOffset 处分割，独立数据的打包和通信可以立即进行，而依赖数据只需等待前一个脉冲的信号。这使得所有脉冲能够并行推进，同时保持必需的顺序关系。

全局脉冲顺序按 Z→Y→X 维度连接定义：

  • firstDependentPulse(z0) = none
  • firstDependentPulse(y0) = z0
  • firstDependentPulse(x0) = y0

NVLink 上，kernel 使用 warp 级 TMA 远程存储将打包与通信流水线化：

  1. 每个 threadblock 内的每个 warp 将坐标的专用 chunk 打包到共享内存
  2. 尊重 TMA 合并的 128-byte 对齐要求和 float3 布局
  3. 作为 chunk 打包完成，选出的 warp leader 发出异步 TMA 存储操作
  4. 其他 warp 可以继续打包 — 解耦了 warp 间的进度，避免 block 级同步

使用 cuda::ptx::cp_async_bulk 驱动 TMA 引擎。TMA 自主处理传输，释放 SM 资源用于计算。

InfiniBand 路径：staging buffer + put-with-signal

对于 InfiniBand，kernel 仍使用 packWithDeps 基于依赖分区工作：

  1. 将坐标打包到全局内存中的 staging buffer
  2. 通过 nvshmem_float_put_signal_nbi 发出远程 put
  3. 独立数据立即打包，依赖数据通过信号同步等待
  4. 重叠独立数据打包与依赖解析

Force Halo 交换：FusedCommUnpackF（Algorithm 6）

// Kernel: FusedCommUnpackF
for p ← ctx.totalPulses - 1 - blockIdx.y down to 0 step -gridDim.y do
    meta ← pulses[p]  // reverse send/recv roles vs coordinates

    if p = ctx.totalPulses-1 and first thread then
        syncAndCommWithDeps(DATA, false, ctx, meta, p, pulses,
                            meta.recvBuf, &forces[meta.atomOffset], null)
    end if

    dataSize ← (meta.remoteForceSrc ? meta.recvSize : meta.sendSize)
    bufLength ← 2048; chunkOffset ← blockIdx.x × bufLength
    chunkSize ← min(bufLength, dataSize - chunkOffset)

    Initialize shared memory buffers and barriers

    if threadIdx.x == 0 and chunkOffset < dataSize then
        tma_ld(smemMap ← meta.indexMap + chunkOffset, chunkSize,
               indexMapLoadBarrier)
        acquire_wait(ctx.signal[p] == ctx.sigVal)
    end if

    if isNVLinkAccess(meta.remoteForceSrc) then
        tma_ld(smemF ← meta.remoteForceSrc + meta.atomOffset + chunkOffset,
               bufLength, forceBufLoadBarrier)
        forceBufLoadBarrier.arrive_and_wait()
    end if

    indexMapLoadBarrier.arrive_and_wait()

    // Parallel accumulation across threads
    for each received entry: map to target atom index using index map
        write/accumulate into forces (use atomicAdd when accumulating)

    if p > 0 then
        syncAndCommWithDeps(DEP_MGMT, true, ctx, meta, p, pulses,
                            null, &forces[0], null)
    end if
end for

关键设计：从最后一个脉冲的已计算受力开始，沿依赖链向后工作。使用 atomicAdd 在所有脉冲上并行解包，最小化依赖 prior pulse 的数据的等待时间。

融合信号与内存排序

不同于典型 NVSHMEM 实现的重要创新：将数据打包/传输与接收方通知融合在单个 kernel 中（通过 syncAndCommWithDeps）：

// syncAndCommWithDeps (Algorithm 5)
block_barrier(); only threadIdx.x = 0 proceeds to notify

if needLastBlock then
    old ← atomicIncReleaseGpu(blockCompletionCounter[p])
    if old ≠ meta.blocksForPulse-1 return
end if

if isNVLink then
    hasDataWrites ? system_release_store(receiverSignal, ctx.sigVal)
                  : system_relaxed_store(receiverSignal, ctx.sigVal)
else  // InfiniBand
    nvshmem_float_put_signal_nbi(dest, src, size, receiverSignal,
                                  ctx.sigVal, sendRank)
  • NVLink：system_release_store 确保前置数据写入可见；system_relaxed_store 用于无需刷新的场景
  • InfiniBand：nvshmem_float_put_signal_nbi 非阻塞 put-with-signal
  • 仅从最后一个 block 发射 system-scope 操作，最小化开销

层级同步优化

Receiver Notification = Data Transfer Kernel 内部融合
↓
Per-pulse ordering enforced through:
  - GPU-scope release atomic increment (block completion counter)
  - System-scope acquire/release stores (NVLink)
  - Put-with-signal (InfiniBand)

四、核心公式

通信脉冲数

对于 D 维域分解（np(x), np(y), np(z)）： Total Pulses=np(x)+np(y)+np(z)\text{Total Pulses} = n_p(x) + n_p(y) + n_p(z)

依赖关系

firstDependentPulse(z₀) = none
firstDependentPulse(y₀) = z₀
firstDependentPulse(x₀) = y₀

TMA 合并要求

  • 对齐粒度：128 bytes
  • 数据类型：float3
  • Warp leader 模式：每个 warp 的 thread 0 发出 TMA 操作

内存排序语义

// NVLink: system scope acquire-release
system_release_store(receiverSignal, ctx.sigVal)   // data writes visible
system_relaxed_store(receiverSignal, ctx.sigVal)   // no prior writes to flush

// InfiniBand: put-with-signal
nvshmem_float_put_signal_nbi(dest, src, size, signal, sigVal, rank)

// GPU scope atomic
atomicIncReleaseGpu(blockCompletionCounter[p])

关键性能指标

  • Local work：local non-bonded kernel 的开始到结束
  • Non-local work：第一个 pack 开始到最后一个 unpack 结束
  • Non-overlap：local non-bonded kernel 结束到最后一个 unpack 结束（clamp 到 0）
  • Time per step：总运行时间减去 Domain Decomposition 和 Neighbor Search 壁时间（每 200 步执行一次）

五、核心创新

创新点说明效果
GPU 发起的 halo exchange首个将 GROMACS DD halo exchange 完全迁移到 GPU kernel 的方案消除 CPU-GPU 同步瓶颈
Fused 多脉冲内核单个 kernel 并行处理所有维度脉冲，而非串行减少 kernel 启动开销
依赖分区策略索引映射在 depOffset 分割，独立数据立即发送最大化通信并发
TMA 流水线打包+传输Warp 级 TMA 远程存储，解耦 warp 间进度NVLink 上亚微秒级延迟
融合信号与通信数据打包/传输和接收方通知在单个 kernel 内不同于典型 NVSHMEM 实现
自适应互连策略运行时根据 NVLink/InfiniBand 自动切换通信方式跨硬件配置最优性能

六、实验结果

测试平台

NVIDIA EOS 超级计算机：

  • 576 个 DGX H100 节点
  • 每节点：2× Intel Xeon 8480C CPU + 8× H100 GPU
  • 总计：4,608 个 GPU
  • 节点内：NVLink 4.0 + NVSwitch
  • 节点间：ConnectX-7 NDR 400G InfiniBand，rail-optimized 非阻塞全树拓扑
  • CUDA 12.6.3, NVSHMEM 3.1.7 (IBRC transport)

GB200 NVL72 多节点 NVLink：

  • 36×2 配置（72 个 GPU）
  • 多节点 NVLink 直连

基准系统

Grappa 基准集（水-乙醇混合物），系统大小 45,000 到 4600 万原子：

  • 45k, 90k, 180k, 360k, 720k, 1440k, 2880k, 5760k, 23040k 原子

使用反应场静电模型以聚焦短程交互和 halo exchange 分析。

Intra-node 性能对比（Fig. 3）

系统大小4 GPUs8 GPUs
45k (11.25k/GPU)NVSHMEM +46% (1649 vs 1126 ns/day)-
180k (45k/GPU)NVSHMEM +4% (1103 vs 1058 ns/day)NVSHMEM +28% (1249 vs 973 ns/day)
360k (90k/GPU)基本持平 (671 vs 670 ns/day)NVSHMEM +17% (910 vs 779 ns/day)

趋势：从小系统（通信主导）到大系统（计算主导），NVSHMEM 优势从 46% 降至 17%，反映 TMA 批量传输在计算不足以隐藏通信开销时的最大收益。

GB200 NVL72 强扩展（Fig. 4）

以单节点性能为基线（720k: 492 ns/day, 1440k: 272 ns/day）：

720k 系统：

节点数GPU数ns/day并行效率
28~41484%
416~27055%
832~15732%

1440k 系统：

节点数GPU数ns/day并行效率
28~23988%
416~19371%
832~13148%

Multi-node 性能对比（EOS, Fig. 5）

系统节点数NVSHMEMMPI提升
720k81103944+17%
1440k16--+12%, 效率 33% vs 31%
5760k128--1.3x, 效率 15% vs 11%
23040k288716633+13%

设备侧时序分析（Fig. 6-8）

Intra-node（Fig. 6, 4 GPUs, 1D DD）：

系统原子/GPULocal workNon-local (MPI)Non-local (NVSHMEM)Non-overlap
45k11.25k~22 µs116 µs64 µs116 µs → 64 µs
180k45k~38 µs101 µs94 µs-
360k90k~152 µs~152 µs~152 µs接近完美重叠

Multi-node 11.25k/GPU（Fig. 7, 1D→2D→3D）：

  • 1D→2D：non-local work 变化 <11%（尽管脉冲数翻倍）
  • 2D→3D：non-local time 增加 ~45%（与脉冲数 1.5x 增长一致）

Multi-node 90k/GPU（Fig. 8）：

  • 1D：NVSHMEM 完全重叠 local 和 non-local work
  • 2D：NVSHMEM 比 MPI 快 ~24 µs（尽管 local kernel 慢 ~16 µs，但不在关键路径）
  • 3D：NVSHMEM 比 MPI 快 50-60 µs（尽管 local kernel 慢 ~10 µs）

扩展极限分析

  • 最优扩展点：~20k 原子/GPU（NVLink 和 NVLink+InfiniBand 均观察到）
  • GPU 利用率限制：<10-25k 原子/GPU 时，不足以饱和现代 H100（>100 SMs）
  • NVSHMEM 代价：SM 资源共享导致 local work 减慢 10-16 µs，但在大多数场景下被 non-local work 减少所补偿

七、相关工作

工作关系
GROMACS thread-MPI之前实现单节点异步通信，但不支持多节点
cuFFTMp (GROMACS 2023)使用 NVSHMEM 的 PME 分解，为本工作奠定基础
Stream-aware MPI粗粒度方法，难以利用 NVLink 细粒度能力
NAMD + Charm++早期使用 PGAS 降低 MD 延迟
LAMMPS SHMEM先前探索过 SHMEM 在 MD 中的应用，但有对称分配缺点
QUDA / PETSc / Kokkos其他使用 NVSHMEM 的应用，但未针对 MD halo exchange 优化

八、总结

核心贡献

  1. GPU 发起的 halo exchange：基于 NVSHMEM 的 GROMACS 中性领土域分解算法重构，融合 pack/unpack 和数据传输
  2. Fused 多脉冲设计：跨所有分解维度的通信阶段融合，通过独立/依赖数据的分离处理最大化并发
  3. 层级同步优化：接收方通知融合到数据传输 kernel 中
  4. 自适应互连策略：基于 NVLink/InfiniBand 的运行时自适应通信，TMA 异步复制引擎
  5. 全面性能评估：在 EOS（576 H100 节点）和 GB200 NVL72 上的强扩展研究

技术影响

  • GPU 发起通信的普适价值：证明 GPU-initiated communication 对广泛延迟敏感应用的强扩展有深远益处
  • 可复用性：类似 halo exchange 算法常见于 CFD、天体物理等领域，本工作的优化可被其他 HPC 代码复用
  • GROMACS 未来方向：计划将此方法扩展到 PME 任务的坐标/力通信

局限性

  • Rank specialization 不兼容：当前实现无法与多 PME rank 的 cuFFTMp 结合使用（NVSHMEM 全局对称分配模型限制）
  • SM 资源竞争：GPU SM 用于同步和通信会与重叠计算产生资源竞争，在大输入小节点数场景下 MPI 略优
  • Fully GPU-resident 未完成：目前某些不平衡 sync 仍需 CPU 同步，计划设计完全 GPU 驻留方案
  • 仅单精度：基准使用反应场模型聚焦短程交互，长程 PME 交互未在本工作中覆盖
  • 验证平台有限：主要在 EOS (H100) 和 GB200 NVL72 上验证，HPE Slingshot 11 仅在较小规模观察

九、参考资源