分子动力学 (MD) 模拟是计算科学领域要求最严苛的工作负载之一。使用它们，研究人员可以非常详细地观察原子行为，从蛋白质折叠到药物和材料发现。
但成本高昂。仿真模拟了数十万到数千万个原子，在数十亿个步骤中每步提升飞秒。由于问题规模通常是固定不变的，因此并行化必须同时将工作分散到所有可用资源中，这种机制被称为强扩展。
GROMACS 是全球应用最广泛的 MD 软件包之一。现代硬件和 异构 CPU-GPU 并行化 已将性能提升至亚毫秒级，使多个 GPU 的每个时间步长达到 100 – 200 微秒。实现此级别的强大扩展需要极低的内核延迟，但随着模拟在更多 GPU 上的扩展，GPU 到 GPU 的通信仍然是一个基本瓶颈。
大多数大规模 HPC 应用程序 (包括 GROMACS) 都使用消息传递接口 (MPI) 进行进程间通信。但是，MPI 专为以 CPU 为中心的执行而设计。
当 GROMACS 在 GPU 上运行时，通信工作流会强制 GPU 暂停，而 CPU 会在发出 GPU 恢复信号之前编排数据传输。在 GROMACS 的光环交换 (在相邻 GPU 域之间共享边界原子数据的算法) 中，这种交接会在所有三个空间维度中重复进行，在迭代速率峰值和可扩展性上限下，消耗的 CPU 总时间超过 50%。
在本文中，我们研究了通过使用 NVSHMEM 将 CPU 编排的 MPI 替换为 GPU 原生通信来消除此瓶颈的策略。
使用设备发起的远程内存访问来加速 GROMACS
加速 GROMACS 的关键在于通过 GPU 原生通信消除 CPU-GPU 传递。NVIDIA NVSHMEM 是一个基于 OpenSHMEM 分区全局地址空间模型实现远程内存访问的库。NVSHMEM 使 GPU 内核能够直接启动数据传输，从而将 CPU 从关键路径中移除，并实现更好的计算和通信重叠。
GROMACS 光环交换使用分阶段转发机制：数据通过每个空间维度内的中间秩在一个或多个通信步骤或脉冲中进行转发。每个脉冲都从一个内核包开始，该内核将边界原子收集到连续的发送缓冲区中。这在通信阶段 (Z* Y* X) 之间创建了依赖关系链，而传统实现会通过粗略相位级序列化来处理这些依赖关系。重新编写后，将这些粗屏障替换为直接在融合核函数中表示的细粒度每脉冲信号。
一个自然的起点是流触发 API (nvshmemx_put_signal_nbi_on_stream) ，它可以消除 CPU-GPU 同步障碍，并尽可能减少代码重组。它有两个限制：通信只能从内核边界开始，防止脉冲内的数据包/ 通信重叠，而流控制的测序意味着从属转发链 (Z* Y* X) 无法融合。
我们在“替换主机端通信”和“跨脉冲启用并行性”部分中通过细粒度每脉冲依赖关系解决了这两个问题。
图 1. GROMACS GPU 驻留 MPI 调度，用于具有两个 X/ F 通信脉冲的 2D DD。CPU 时间轴 (顶部) 显示了内核启动、GPU 同步等待和 MPI。GPU 流 (底部) 显示重叠作业和强制通信下图显示了具有两个通信脉冲的仿真时间步的执行图，展示了内核启动的 RMA、内核融合和事件驱动的多队列执行如何重叠通信和计算。
图 2. 2D DD 的 GROMACS NVSHMEM 调度。重叠的 CPU 内核启动和 GPU 计算流可消除关键路径中的光环交换从 MPI 基准到优化的 GPU 原生光晕交换
现有的基于 MPI 的光环交换通过相邻秩传递数据，以更大限度地减少通信量，但代价是三个阶段 (Z = Y* X) 之间的顺序依赖关系，每个阶段都包含一个或多个脉冲。以下代码显示了坐标光环交换包内核及其在主机侧从 GROMACS 源代码调用的内容。
```cuda
// GPU kernel: pack boundary atom coordinates into a contiguous send buffer
__global__ void packSendBufKernel(float3* dataPacked, // packed coordinate buffer
const float3* data, // full coordinate array
const int* map, // index map: which atoms to pack
const int mapSize)
{
int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
if (threadIndex < mapSize)
{
dataPacked[threadIndex] = data[map[threadIndex]];
}
}
// GPU kernel: unpack received halo forces, scatter back into the force array via index map
template<bool accumulate>
__global__ void unpackRecvBufKernel(float3* forces, // full force array
const float3* dataPacked, // received packed forces
const int* map,
const int mapSize)
{
int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
if (threadIndex < mapSize)
{
if (accumulate)
forces[map[threadIndex]] += dataPacked[threadIndex];
else
forces[map[threadIndex]] = dataPacked[threadIndex];
}
}
```
```c
// CPU host code: coordinate halo exchange (MPI path)
for each dimension d in [Z, Y, X]:
for each pulse p in dimension d:
// 1. Launch GPU pack kernel on non-local stream
packSendBufKernel<<<grid, block, 0, nonLocalStream>>>(
sendBuf, coords, indexMap, mapSize);
// 2. CPU blocks until GPU pack completes
cudaStreamSynchronize(nonLocalStream); // *** BLOCKING ***
// 3. CPU orchestrates MPI transfer
MPI_Isend(sendBuf, sendSize, sendRank, ...);
MPI_Irecv(coords + atomOffset, recvSize, recvRank, ...);
MPI_Waitall(...); // *** BLOCKING ***
// Received data lands directly
// in the coordinate buffer at the correct offset and is consumed
// by the non-local non-bonded force kernel as-is.
```
虽然这些内核很简单，但瓶颈在于封装它们的 CPU 端控制流。每个脉冲都需要 CPU 在 GPU 上阻塞，启动 MPI 传输，并指示 GPU 为所有三个维度中的每个脉冲恢复序列化。
力光环交换遵循相同的模式，但方向相反 (X+ Y+ Z) 。发送坐标的秩现在会接收在这些光晕原子上计算的力，并且强制需要通过 unpackRecvBufKernel 执行散射 – 解包步骤，因为它们必须累积到全力数组中的正确位置。
每个脉冲都需要两次阻塞 CPU – GPU 同步：一次在 MPI 调用之前，以保证包内核已完成，另一次在下一次脉冲使用转发数据之前执行。借助 3D 分解和每个维度一个脉冲，对于坐标，每个时间步需要等待六个块，对于力则需要另外六个块，总共 12 个。每种方法都会增加关键路径的延迟，从而影响 GROMACS 迭代率。
将主机侧通信替换为 NVSHMEM
在此理念的基础上，第一个设计将打包、远程放置和完成等待折叠到每个通信脉冲的单个内核中。主机仍然无法按脉冲顺序启动，但会消除脉冲之间的 CPU – GPU 屏障。现在，仅 NVIDIA CUDA 流就能保证正确的执行顺序。主机调用融合的 pack-put-wait 内核及其实现如下：
```c
// CPU host code: one kernel launch per pulse, no CPU-GPU sync in between
int pulseOffset = 0;
for each dimension d in [Z, Y, X]:
for each pulse p in dimension d:
// Pack + send + wait are fused into a single GPU kernel
packSendBufAndPutNvshmemKernel<<<grid, block, 0, nonLocalStream>>>(
coords, dataPacked, indexMap, sendSize,
sendRank, atomOffsetInSendRank,
signalReceiverRank + pulseOffset, signalCounter, bar, recvSize);
pulseOffset++;
```
整个“包 – 放 – 等待”操作融合到单个设备端启动中。内核改编自 GROMACS 实现，如下所示：
```cuda
__global__ void packSendBufAndPutNvshmemKernel(
float3* coords, float3* dataPacked, int* map,
int sendSize, int sendRank, int atomOffsetInSendRank,
uint64_t* signalReceiverRank, uint64_t signalCounter,
cuda::barrier<cuda::thread_scope_device>* bar, int recvSize)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (sendSize > 0)
{
// Pack: grid-stride gather into contiguous send buffer
for (int i = tid; i < sendSize; i += blockDim.x * gridDim.x)
dataPacked[i] = coords[map[i]];
// Device-scoped barrier: all CTAs must finish packing before put
auto token = bar->arrive();
if (blockIdx.x == 0)
{
bar->wait(std::move(token));
// Put packed data into peer's coordinate array
nvshmemx_float_put_signal_nbi_block(
&coords[atomOffsetInSendRank], // peer destination
dataPacked, // local source
sendSize * 3, // float count
signalReceiverRank, signalCounter,
NVSHMEM_SIGNAL_SET, sendRank);
}
}
// Wait for incoming data from recvRank before exiting
if (tid == 0 && recvSize > 0)
nvshmem_signal_wait_until(signalReceiverRank, NVSHMEM_CMP_EQ, signalCounter);
}
```
此实现消除了关键路径中的所有 CPU-GPU 同步。CPU 仅执行核函数启动，这可能与 GPU 执行重叠。打包后，所有 GPU 线程 (CTA) 都会到达设备范围内的屏障，块 0 会在发出 nvshmemx_float_put_signal_nbi_block 调用之前等待该屏障，以确保所有线程均已完成。该调用将打包的数据传输到对端，并在对端设置 signalReceiverRank 以确认完成；接收器在使用数据之前等待该信号。
但是，请注意，此版本仍按顺序处理脉冲 (按流顺序执行) ，并且每个脉冲仍会产生单独的核函数启动。
优化跨 NVIDIA NVLink 连接的 GPU 的高效传输
在所述设计中，数据始终通过 NVSHMEM 传输 (nvshmemx_float_put_signal_nbi_block) 进行路由，这适用于任何互连，但当对端通过 NVLink 连接时，GPU 可以直接加载/ 存储到彼此的内存中。与其打包到本地缓冲区并发出看涨期权，不如直接打包到节点的坐标数组中，从而消除单独的全局内存往返行程。
当可通过 NVLink 访问点对点时，nvshmem_ptr(remotePtr, peerRank) 会返回非空设备指针，否则返回 null。我们会在内核和分支中相应地查询一次，如下所示：
```cuda
__global__ void packSendBufAndPutNvshmemKernel(
float3* coords, float3* dataPacked, int* map,
int sendSize, int sendRank, int atomOffsetInSendRank,
uint64_t* signalReceiverRank, uint64_t signalCounter,
cuda::barrier<cuda::thread_scope_device>* bar, int recvSize)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (sendSize > 0)
{
// Probe for NVLink: non-null means direct store is possible
float3* remotePtr = (float3*)nvshmem_ptr(coords, sendRank);
bool isNVLink = (remotePtr != nullptr);
// NVLink: pack directly into peer's coordinate buffer
// IB: pack into local staging buffer for later put
float3* dest = isNVLink ? (remotePtr + atomOffsetInSendRank)
: dataPacked;
for (int i = tid; i < sendSize; i += blockDim.x * gridDim.x)
dest[i] = coords[map[i]];
auto token = bar->arrive();
if (blockIdx.x == 0)
{
bar->wait(std::move(token));
if (isNVLink)
{
// Data already in peer memory — just signal completion
uint64_t* peerSignal = (uint64_t*)nvshmem_ptr(
signalReceiverRank, sendRank);
storeReleaseSysAsm(peerSignal, signalCounter);
}
else
{
// IB: combined data transfer + signal notification
nvshmemx_float_put_signal_nbi_block(
(float*)&coords[atomOffsetInSendRank],
(float*)dataPacked,
sendSize * 3,
signalReceiverRank, signalCounter,
NVSHMEM_SIGNAL_SET, sendRank);
}
}
}
if (tid == 0 && recvSize > 0)
nvshmem_signal_wait_until(signalReceiverRank, NVSHMEM_CMP_EQ, signalCounter);
}
```
在 NVLink 路径上，打包循环直接写入点对点坐标数组中的确切位置。在块 0 确认所有 CTA 均已使用设备范围内的屏障完成打包后，它会通过系统范围的发布存储 (st.release.sys) 向对等方发送信号，以确保在对等方读取信号之前，所有之前的直接写入都对其可见。由于数据已经就位，因此不需要额外的数据放置操作。
当 nvshmem_ptr 返回 null 时，其行为与第一个设计相匹配。数据被打包到本地暂存缓冲区中，并传递给 nvshmemx_float_put_signal_nbi_block，由 NVSHMEM 传输通过 InfiniBand、Slingshot 或其他 RDMA 网络传输。
这种优化凸显了这两种路径之间的重要不对称。NVLink 路径需要多个 CUDA 核心 (SM) 并发发送存储，以使 NVLink 带宽饱和，而 InfiniBand 路径只需要一个线程 (CTA) ，即可使用 NVSHMEM 传输启动批量传输。剩下的瓶颈是每个脉冲内核启动用度和脉冲之间的刚性顺序排序，即使一些数据没有交叉脉冲依赖项也是如此。
通过内核融合实现跨脉冲的并行性
关键的认识是，脉冲中的所有数据并非都依赖于先前的脉冲。脉冲中的大多数原子都可以被打包并立即发送；只有接收前一个脉冲数据的“转发”原子必须等待。这是通过在 dependencyAtomOffset 对每个脉冲的索引图进行分区来实现的：
- 独立原子 (atomIndex =
dependencyAtomOffset)： 一旦脉冲运行，即可安全地打包和发送 - 相关原子 (atomIndex ≥
dependencyAtomOffset)： 必须等待相关的前脉冲信号
借助此分区，单个内核可同时处理所有脉冲。相关工作仅在其所需的特定脉冲信号上等待，而不是在整个相位上等待。
依赖关系感知打包
packSubset 设备函数有选择地按值封装原子，首先针对独立原子，然后在等待信号后针对相关原子进行第二次封装。
```cuda
// Templated pack: selectively gather atoms based on threshold
// packAll=true: pack everything (no dependency check needed)
// packAll=false, packLessThan=true: pack only atomIndex < threshold (independent)
// packAll=false, packLessThan=false: pack only atomIndex >= threshold (dependent)
template<bool packAll, bool packLessThan>
__device__ void packSubset(float3* dest,
float3* data,
const int* map,
int sendSize,
int gridStride,
int threadIndex,
int threshold,
int& hasDependencyAtoms)
{
for (int idx = threadIndex; idx < sendSize; idx += gridStride)
{
int atomIndex = map[idx];
float3 srcVal = data[atomIndex];
if constexpr (packAll)
{
dest[idx] = srcVal;
}
else
{
bool packNow = packLessThan ? (atomIndex < threshold)
: (atomIndex >= threshold);
if (packNow)
dest[idx] = srcVal;
else if (packLessThan)
hasDependencyAtoms++; // track deferred atoms in first pass
}
}
}
```
脉冲按 Z→Y→X 顺序全局编号。每个原子都会等待所有之前的脉冲信号，在填充其相关原子之前。每个维度只有一个脉冲：[z0, y0, x0]：z0 没有依赖项；y0 等待 z0；x0 等待两者。packHaloCoords 中的循环 for (int i = currPulse; i > 0; i--) 可实现以下目标：
```cuda
// Dependency-aware packing: pack independent data immediately,
// wait for signals from prior pulses, then pack dependent data
__device__ void packHaloCoords(float3* dest,
int sendSize,
float3* data,
const int* map,
int currPulse,
uint64_t signalCounter,
uint64_t* signalReceiverRankCurr,
int dependencyAtomOffset)
{
int gridStride = blockDim.x * gridDim.x;
int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
if (currPulse > 0)
{
// Pass 1: pack atoms below dependency threshold (no waiting)
int hasDependencyAtoms = 0;
packSubset<false, true>(dest, data, map, sendSize, gridStride,
threadIndex, dependencyAtomOffset,
hasDependencyAtoms);
// Thread 0 waits for all prior pulses to deliver their data
if (threadIdx.x == 0)
{
for (int i = currPulse; i > 0; i--)
while (loadRelaxedSys(signalReceiverRankCurr - i) != signalCounter)
;
}
constexpr unsigned fullWarpMask = 0xFFFFFFFF;
int need_to_wait = __any_sync(fullWarpMask, hasDependencyAtoms > 0);
__syncthreads();
// Pass 2: pack the deferred atoms (>= threshold)
if (need_to_wait)
{
int unused = 0;
packSubset<false, false>(dest, data, map, sendSize, gridStride,
threadIndex, dependencyAtomOffset,
unused);
}
}
else
{
// Pulse 0: no dependencies, pack everything unconditionally
int unused = 0;
packSubset<true, true>(dest, data, map, sendSize, gridStride,
threadIndex, 0, unused);
}
__syncthreads();
}
```
需要注意的是：loadRelaxedSys 使用 PTX ld.relaxed.sys.global.u64 – 系统范围，因为信号是由远程 GPU 编写的，由于它是旋转循环，因此可以放松排序。等待后的 __syncthreads() 会对所有线程显示受保护的数据。在产品级代码中，使用 libcu++ 中的 cuda::memory_order_relaxed 的 cuda::atomic_ref<uint64_t, cuda::thread_scope_system> 可提供等效的语义。
融合后的内核将每个脉冲分配给 blockIdx.y 行，该行中的所有块通过网格跨度循环 (由 blockIdx.x 索引) 在一个脉冲上进行协作。预先计算的远程指针会在运行时选择传输：非 null 表示 NVLink 直接存储，null 会回退至 NVSHMEM 传递路径 (InfiniBand) ：
```cuda
// One kernel launch replaces all per-pulse pack + send operations
// Grid: (blocksPerPulse, totalNumPulses, 1) — one blockIdx.y per pulse
__global__ void fusedPulsesPackAndSendKernel(
float3* data, // full coordinate array
HaloExchangeData* pulseData, // per-pulse metadata array
uint64_t* signalReceiverRank,
uint64_t signalCounter,
uint32_t* gridSync, // per-pulse block-completion counters
int totalNumPulses,
int dependencyAtomOffset)
{
int currPulse = blockIdx.y;
if (currPulse < totalNumPulses)
{
HaloExchangeData halo = pulseData[currPulse];
// Precomputed by nvshmem_ptr; null => IB/RDMA path
float3* remotePtr = halo.remoteCoordsPutPtr;
bool isNVLink = (remotePtr != nullptr);
// Choose destination: peer memory (NVLink) or local staging buffer (IB)
float3* dest = isNVLink ? (remotePtr + halo.atomOffsetInPeer)
: halo.d_sendBuf;
if (halo.sendSize > 0)
{
// Two-pass dependency-aware packing
packHaloCoords(dest, halo.sendSize, data, halo.indexMap,
currPulse, signalCounter,
signalReceiverRank + currPulse,
dependencyAtomOffset);
// Last-block sync and peer notification
if (threadIdx.x == 0)
signalPeerOnLastBlock(
isNVLink,
gridSync + currPulse, signalReceiverRank + currPulse,
signalCounter, halo.sendRank, gridDim.x,
dest, halo.recvPtr, halo.sendSize);
}
}
}
```
一种微妙但对性能至关重要的优化是将接收端通知融合到同一内核中。典型的 NVSHMEM 模式使用单独的内核或主机侧调用向接收端发送信号。我们使用 signalPeerOnLastBlock 来避免这种情况，它使用分层块完成计数器来选择最后一个块，如下所示。
```cuda
// Custom PTX wrappers for precise memory ordering control
__device__ uint32_t atomicIncReleaseGpu(uint32_t* addr, int32_t mod) {
uint32_t old;
asm("atom.inc.release.gpu.global.u32 %0,[%1],%2;"
: "=r"(old) : "l"(addr), "r"(mod) : "memory");
return old;
}
__device__ void storeReleaseSys(uint64_t* ptr, uint64_t val) {
asm("st.release.sys.global.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");
}
__device__ void signalPeerOnLastBlock(
bool isNVLink, // selects NVLink direct-store path vs. InfiniBand put
uint32_t* gridSync, // per-pulse block-completion counter
uint64_t* signalCurr, // peer's signal slot for this pulse
uint64_t signalCounter, // value to store into the signal slot
int sendRank, // peer rank to notify
int numBlocks,
const float* dataPacked, // staging buffer (InfiniBand path only)
void* recvPtr, // peer's receive buffer (InfiniBand path only)
int sendSize) // payload size in float3 elements (InfiniBand only)
{
// GPU-scope release atomic: flushes this block's writes to GPU scope
// and returns the old counter value
uint32_t old = atomicIncReleaseGpu(gridSync, numBlocks - 1);
// Only the last arriving block proceeds
if (old != numBlocks - 1)
return;
if (isNVLink)
{
// NVLink: data already written directly — just signal the peer
// st.release.sys ensures all prior writes visible before signal
uint64_t* peerSignal = (uint64_t*)nvshmem_ptr(signalCurr, sendRank);
storeReleaseSys(peerSignal, signalCounter);
}
else
{
// InfiniBand: combined data transfer + notification
nvshmem_float_put_signal_nbi(
recvPtr, dataPacked, sendSize * 3,
signalCurr, signalCounter, NVSHMEM_SIGNAL_SET, sendRank);
}
}
```
atomicIncReleaseGpu 包装 PTX atom.inc.release.gpu.global.u32，将块完成跟踪与 GPU 范围版本相结合，刷新之前的所有写入。只有最后一个到达的块支付系统范围存储或 NVSHMEM 的费用；所有其他块都会立即退出。
NVLink 路径：基于 TMA 的直接存储
在 NVLink 路径上，进一步优化将数据移动从 SM 卸载到 Tensor Memory Accelerator (TMA) 异步复制引擎 ( NVIDIA Hopper 及更高版本) ，从而释放 SM 资源用于计算。TMA 变体以每个线程束粒度工作，每个线程束独立地发布自己的 TMA 存储：
- 每个线程束都将其独立的原子打包到共享内存中
- 如果线程束的数据块包含相关原子，线程束引导器会在所有通道都填充剩余原子之前等待先前脉冲信号，从而完全填充共享内存缓冲区
- 线程束引导器向节点发送完整缓冲区的单个 TMA 存储 (
cuda::ptx::cp_async_bulk) ，而其他线程束则独立进行，无需全块同步
```cuda
// NVLink TMA path: per-warp pipelining of pack + remote store
int warpId = threadIdx.x / 32;
int laneId = threadIdx.x % 32;
int warpChunk = SMEM_BUFFER_LENGTH / numWarpsPerBlock;
float3* mySmem = &sharedBuf[warpId * warpChunk];
int chunkOffset = blockIdx.x * SMEM_BUFFER_LENGTH + warpId * warpChunk;
bool isWarpLeader = warpElect(); // elect one thread per warp
// Pass 1: pack independent atoms into shared memory
for (int i = laneId; i < warpChunk; i += 32) {
int atomIdx = indexMap[chunkOffset + i];
if (atomIdx < dependencyAtomOffset)
mySmem[i] = data[atomIdx];
}
// Pass 2 only when this warp's chunk contains dependent atoms
if (hasDependencyAtoms) {
// Warp leader: wait on prior pulse signals before packing dependent atoms
if (isWarpLeader) {
// ... wait on prior pulse signals at warp granularity ...
}
__syncwarp();
for (int i = laneId; i < warpChunk; i += 32) {
int atomIdx = indexMap[chunkOffset + i];
if (atomIdx >= dependencyAtomOffset)
mySmem[i] = data[atomIdx];
}
}
// Warp leader issues TMA async store of the full contiguous buffer to peer
if (isWarpLeader)
cuda::ptx::cp_async_bulk(remoteDst + chunkOffset, mySmem,
warpChunk * sizeof(float3));
// Other warps continue independently — no syncthreads() between warps
```
TMA 引擎独立于 SM 处理远程存储，从而释放计算资源。每个线程束在完成打包后立即发布其存储；只有独立原子的线程束会跳过依赖项等待。共享内存可用作满足 TMA 对齐要求的暂存区，同时确保在 TMA 批量复制之前有一个完全填充的缓冲区。
强制光环交换：反向依赖链
在力晕交换过程中，计算出的边界原子受的力会返回到最初拥有这些原子的秩。此时，依赖项链反向。上一个脉冲的作用力首先可用，而较早的脉冲先等待较晚的脉冲，然后再转发累积的作用力。融合内核在单个内核中处理强制通信和解包，在数据通过脉冲链反向流动时进行处理。
SM 占用和内核调度
重叠的优势还取决于 GPU 调度的微妙之处。计算密集型本地非键合力内核使用数百个线程块填充每个 SM，并且调度程序无法抢占运行中的线程块，因此在一些本地线程块完成并释放其 SM 之前，高优先级通信内核无法启动。借助 MPI，在每次启动时，为每个脉冲启动的打包和解压缩核函数都可以独立支付这种调度延迟。融合后的内核只需付出一次代价：单个内核可处理所有脉冲，依赖感知分区允许其在调度时立即打包和发送独立数据，而依赖数据仍在运行中。更少的内核边界意味着整个非本地流中的重新调度事件会更少，因此内核融合可以改善与本地计算的重叠，而不仅仅是减少启动开销。
策略概述
该设计基于四个理念：
- 具有依赖感知分区的内核融合：一个融合的内核可取代每个脉冲包/发送/等待周期，在一次启动中处理所有脉冲。每个脉冲的索引图可分为独立的原子(立即发送)和相关尾部(仅等待每个脉冲的细粒度信号)，因此每个时间步长可将6次启动时间缩短为1次。
- 减少线程块同步: 使用分层发布范围将通知融合到内核中。每个块执行较轻的 GPU 范围释放原子，而最终块的系统范围跨 GPU 释放每个脉冲仅发出一次。
- 事件驱动、无主机执行：跨核依赖项表示为 GPU 可见信号，因此本地和非本地工作流在不使用
cudaStreamSynchronize的情况下连续启动。 - 互连感知传输：NVLink 对等端接收直接存储和释放信号，否则，通过 RDMA 互连 ( InfiniBand、Slingshot) ，内核回退至
nvshmemx_float_put_signal_nbi_block。
扩展性能的飞跃
在 NVIDIA Eos 超级计算机上执行了基准测试，这是一个由 576 个 NVIDIA DGX H100 节点组成的集群，通过 NVLink 4.0 内部节点和 400 Gb/s InfiniBand NDR 节点间连接，使用从 45000 到 23M 个原子的水 – 酒精混合物来表示典型的生物分子工作负载。
与 GPU 感知型 MPI 相比，节点内的性能提升高达 1.5 倍，在受延迟限制的小型系统上，性能提升最大。一个基于四个 GPU 的 45000 原子系统的性能提升了 46% ( 1649 纳秒/ 天比 1126 纳秒/ 天) 。在 NVIDIA GB200 NVL72 多节点 NVLink 集群上，优势扩展到 2 倍，在标准 InfiniBand 集群上扩展到 1.3 倍。
需要权衡的一个问题是，设备初始化通信消耗的 GPU 资源并不多，在节点数较少的情况下，在非常大、计算受限的配置上，传统 MPI 可获得 1 – 3% 的窄边缘优势。然而，在与生产 MD 研究最相关的扩展机制中，新方法始终如一地获胜。
下图显示了使用 MPI 和 NVSHMEM 在 4 个和 8 个 GPU 上进行节点内执行的模拟性能 (图 3) ，以及使用实现的模拟长度 (纳秒/ 天) 和迭代速率 (毫秒/ 步长) 指标 (每节点 4 个 GPU) 执行的多节点性能 (图 4) 。
图 3. 4/ 8 GPU 上的 GROMACS 节点内 MPI 和 NVSHMEM 比较，显示模拟性能 (纳秒/ 天) 和迭代速率 (毫秒/ 步长) 。迭代率是指在所有时间步中计算的每个时间步的平均壁时s
图 4. 多节点 MPI 和 NVSHMEM 强扩展显示模拟性能 (纳秒/ 天) 和迭代速率 (毫秒/ 步长) 作为节点数量的函数。迭代率是指在所有时间步长中计算出的每个时间步长的平均壁时未来的开放式和便携式解决方案
这表明，易受延迟影响的 HPC 应用程序处理 GPU 间通信的方式发生了更广泛的转变。强扩展到微秒级内核时间的应用程序需要与以 CPU 为中心的 MPI 不同的模型。GPU 发起的通信是一个令人信服的答案。这些技术并不局限于分子动力学。光环交换模式出现在整个计算科学 ( CFD、天体物理学、晶格 QCD) 中，并且该方法直接传输到任何依赖分解域之间边界数据交换的应用程序。
下一步自然是将这些想法引入更广泛的 HPC 生态系统。Kokkos 远程空间等分布式数据抽象和 std::execution 等执行模型开始将这些模式封装在可移植接口后面。
更广泛的采用还需要标准化：这里使用的 GPU 原生扩展目前不属于核心 OpenSHMEM 规范，并且需要正式采用设备初始化通信、细粒度信令和直接对等内存访问才能实现统一可用性。
我们的目标是在未来，GROMACS 等专用代码的创新将成为跨异构硬件扩展任何应用程序的标准构建块。如果您运行 GROMACS，请尝试 GROMACS 2026 中由 GPU 发起的光环交换，并测量您自己的模拟的强扩展收益。
若要在其他位置应用相同模式：
- 探索 NVSHMEM 并在您的关键路径上寻找 CPU 编排的通信。
- 阅读包含完整设计和基准的 SC’25 研讨会论文。
