游乐游手机版
首页/AI热点日报/热点详情

GPU加速分子动力学初始化通信实用指南

类型:热点整理2026-07-12
分子动力学模拟中,GPU间的通信瓶颈限制了强扩展性能。GROMACS采用MPI导致CPU频繁同步。通过NVSHMEM实现GPU原生通信,消除CPU-GPU交接,融合打包、传输与等待于单个内核,并利用NVLink优化直接内存访问,显著提升迭代速率。

分子动力学(MD)模拟是计算科学中负担最重的工作负载之一。借助它,研究人员能以惊人的细节观察原子行为,从蛋白质折叠到药物和材料发现,不一而足。

但代价相当高昂。模拟涉及数十万到上千万个原子,每个时间步前进飞秒级,跨越数十亿步。由于问题规模通常是固定的,并行化必须将任务同时分摊到所有可用资源上,这种模式称为强扩展。

GROMACS 是全球使用最广泛的 MD 软件包之一。现代硬件加上异构 CPU-GPU 并行化,已将其性能推至亚毫秒区间,在多个 GPU 上每个时间步达到 100–200 微秒。要在这一水平上实现强扩展,就需要极低的内核延迟,但 GPU 之间的通信仍然是随着 GPU 数量增加而扩展时的根本瓶颈。

大多数大规模 HPC 应用(包括 GROMACS)都使用消息传递接口(MPI)进行进程间通信。然而,MPI 是为以 CPU 为中心的执行模式设计的。

当 GROMACS 在 GPU 上运行时,通信工作流程迫使 GPU 暂停,由 CPU 协调数据传输,然后再通知 GPU 恢复运行。在 GROMACS 的“光环交换”(halo exchange)——即在相邻 GPU 域之间共享边界原子数据的算法——中,这种交接会在所有三个空间维度上重复进行,在峰值迭代率下消耗超过 50% 的 CPU 墙钟时间,从而限制可扩展性。

在这篇文章中,我们探讨如何通过使用 NVSHMEM 实现 GPU 原生通信来替代 CPU 编排的 MPI,从而消除这一瓶颈。

使用设备发起的远程内存访问加速 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 通信脉冲。上方 CPU 时间线显示内核启动、GPU 同步等待和 MPI 操作。下方 GPU 流显示工作与力通信的重叠

下图展示了带有两个通信脉冲的模拟时间步执行图,显示了内核发起的 RMA、内核融合以及事件驱动的多队列执行如何实现通信与计算的重叠。

图 2. 二维域分解中 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
__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<<>>(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 调用之前以确保打包内核完成,另一次在下一个脉冲可以消费转发数据之前强制执行。如果采用三维分解且每个维度一个脉冲,那么每个时间步坐标需要六次阻塞等待,力也需要六次——总共 12 次。每次都会向关键路径增加延迟,影响 GROMACS 的迭代速率。

用 NVSHMEM 替代主机端通信

基于这一思路,第一个设计将打包、远程 put 和完成等待合并到每个通信脉冲的单个内核中。主机仍然按脉冲顺序发出启动,但去掉了脉冲之间的 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<<>>(
        coords, dataPacked, indexMap, sendSize,
        sendRank, atomOffsetInSendRank,
        signalReceiverRank + pulseOffset, signalCounter, bar, recvSize);
    pulseOffset++;
```

整个打包-put-等待操作融合为单个设备端启动。改编自 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* 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 可以直接加载/存储到对方的内存中。与其先打包到本地缓冲区再发出 put,不如直接打包到对端的坐标数组中,从而消除一次全局内存往返。

nvshmem_ptr(remotePtr, peerRank) 在对端可通过 NVLink 访问时返回非空设备指针,否则返回空。我们在内核中查询一次,并据此分支,如下所示:

```cuda
__global__ void packSendBufAndPutNvshmemKernel(
    float3* coords, float3* dataPacked, int* map,
    int sendSize, int sendRank, int atomOffsetInSendRank,
    uint64_t* signalReceiverRank, uint64_t signalCounter,
    cuda::barrier* 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 已完成打包,它就会使用系统作用域 release 存储(st.release.sys)通知对端,这确保了在对端读取信号之前,所有前面的直接写入都可见。无需额外的数据 put 操作,因为数据已经就位。

nvshmem_ptr 返回空时,行为与第一个设计相同。数据被打包到本地暂存缓冲区,然后交给 nvshmemx_float_put_signal_nbi_block,由 NVSHMEM 传输层通过 InfiniBand、Slingshot 或其他 RDMA 结构交付。

这个优化突显了两条路径之间的一个重要不对称性。NVLink 路径需要大量 CUDA 核心(SM)同时发出存储以饱和 NVLink 带宽,而 InfiniBand 路径只需要单个线程(CTA)使用 NVSHMEM 传输发起批量传输。剩下的瓶颈是每个脉冲的内核启动开销以及脉冲之间严格的顺序排序,即使某些数据没有跨脉冲依赖。

通过内核融合实现脉冲间并行

关键洞察是:并非脉冲中的所有数据都依赖于先前的脉冲。大多数脉冲中的原子可以立即打包和发送;只有那些接收来自前一个脉冲数据的“转发”原子才需要等待。这通过将每个脉冲的索引映射在 dependencyAtomOffset 处划分来实现:

  1. 独立原子(atomIndex < dependencyAtomOffset): 一旦脉冲运行,就可以安全地打包和发送
  2. 依赖原子(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
__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 等待 z0x0 等待两者。这就是 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(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(dest, data, map, sendSize, gridStride, 
                                    threadIndex, dependencyAtomOffset, unused);
        }
    }
    else
    {
        // Pulse 0: no dependencies, pack everything unconditionally
        int unused = 0;
        packSubset(dest, data, map, sendSize, gridStride, 
                            threadIndex, 0, unused);
    }

    __syncthreads();
}
```

一个实用说明:loadRelaxedSys 使用 PTX ld.relaxed.sys.global.u64——系统作用域是因为信号由远程 GPU 写入,relaxed 顺序是因为这是一个自旋循环。等待之后的 __syncthreads() 使得受保护的数据对所有线程可见。在生产代码中,可以使用 libcu++ 的 cuda::atomic_ref 配合 cuda::memory_order_relaxed 提供等效语义。

融合内核将每个脉冲分配给一个 blockIdx.y 行,该行的所有块通过网格步长循环(由 blockIdx.x 索引)合作处理一个脉冲。预计算的远程指针在运行时选择传输方式:非空表示 NVLink 直接存储,空则回退到 NVSHMEM put 路径(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 作用域的 release 结合起来,刷新所有前面的写入。只有最后一个到达的块才执行系统作用域存储或 NVSHMEM put;其他块立即退出。

NVLink 路径:基于 TMA 的直接存储

在 NVLink 路径上,进一步的优化是将数据移动从 SM 卸载到张量内存翻跟斗(TMA)异步复制引擎(NVIDIA Hopper 及后续架构),从而释放 SM 资源用于计算。TMA 变体在逐 warp 粒度上工作——每个 warp 独立发出自己的 TMA 存储:

  1. 每个 warp 将其独立原子打包到共享内存中
  2. 如果 warp 的块包含依赖原子,则 warp leader 在打包剩余原子之前等待前一个脉冲的信号——完全填充共享内存缓冲区
  3. warp leader 向对端发出单个 TMA 存储(cuda::ptx::cp_async_bulk),完成整个缓冲区的传输;其他 warp 独立进行,无需块范围的同步
```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 处理远程存储,释放计算资源。每个 warp 在完成打包后立即发出自己的存储;只有独立原子的 warp 完全跳过依赖等待。共享内存用作暂存区,满足 TMA 的对齐要求,同时确保在 TMA 批量复制之前缓冲区已完全填充。

力的光环交换:反转依赖链

在力的光环交换期间,边界原子上的计算力被发送回最初拥有这些原子的秩。这里,依赖链被反转。最后一个脉冲的力最先可用,早期的脉冲在转发累加的力之前等待后面的脉冲。融合内核在单个内核中处理力通信和解包,随着数据向后流过脉冲链进行处理。

SM 占用率与内核调度

重叠收益还依赖于 GPU 调度的一个微妙之处。计算密集的本地非成键力内核会用数百个线程块填满每个 SM,而调度器无法抢占正在运行的块,因此高优先级的通信内核必须等到一些本地块完成并释放它们的 SM 后才能启动。使用 MPI 时,每个脉冲启动的打包和解包内核都会独立承担这种调度延迟。而融合内核只承担一次:单个内核处理所有脉冲,依赖感知分区使其在调度后立即打包并发送独立数据,同时依赖数据仍在飞行中。更少的内核边界意味着整个非本地流上的重新调度事件更少,因此内核融合不仅减少了启动开销,还提高了与本地计算的重叠。

策略总结

设计基于四个理念:

  • 内核融合与感知依赖的划分:一个融合内核取代了每个脉冲的打包/发送/等待循环,在一次启动中处理所有脉冲。每个脉冲的索引映射被拆分为独立原子(立即发送)和依赖尾部(仅等待细粒度的逐脉冲信号),将每个时间步的六次启动减少为一次。
  • 减少线程块同步:通知通过层级化的 release 作用域融合到内核中。每个块执行一个较轻的 GPU 作用域 release 原子操作,而每个脉冲仅由最后一个块发出一次系统作用域的跨 GPU release。
  • 事件驱动的无主机执行:跨内核依赖表示为 GPU 可见的信号,因此本地和非本地流水线可以背靠背启动,无需 cudaStreamSynchronize
  • 互连感知的传输:NVLink 对端接收直接存储加上 release 信号;否则内核回退到通过 RDMA 互连(InfiniBand、Slingshot)的 nvshmemx_float_put_signal_nbi_block

可扩展性能的飞跃

基准测试在 NVIDIA Eos 超级计算机上进行,该集群包含 576 个 NVIDIA DGX H100 节点,节点内通过 NVLink 4.0 连接,节点间通过 400 Gb/s InfiniBand NDR 连接。测试使用了从 45,000 到 2300 万个原子的水-乙醇混合物,代表典型的生物分子工作负载。

节点内性能比 GPU 感知的 MPI 高出最高 1.5 倍,最大的增益出现在较小的、受延迟限制的系统上。一个 45,000 个原子的系统在四个 GPU 上获得了 46% 的提升(1,649 vs. 1,126 ns/day)。在 NVIDIA GB200 NVL72 多节点 NVLink 集群上,收益扩展到 2 倍;在标准 InfiniBand 集群上为 1.3 倍。

唯一的一个折衷是设备发起的通信会消耗少量 GPU 资源,因此在非常大的、计算密集型配置且节点数较少时,传统 MPI 有 1–3% 的微弱优势。然而,在与生产级 MD 研究最相关的扩展区间内,新方法持续胜出。

以下图表展示了使用 MPI 和 NVSHMEM 的模拟性能数据——图 3 为节点内 4 和 8 GPU 执行,图 4 为每节点 4 GPU 的多节点执行,采用达到的模拟长度(ns/天)和迭代速率(ms/步)指标。

图 3. GROMACS 节点内 MPI 和 NVSHMEM 在 4/8 GPU 上的比较,显示模拟性能(ns/天)和迭代速率(ms/步)。迭代速率为每个时间步的平均墙钟时间,跨所有时间步计算
图 4. 多节点 MPI 和 NVSHMEM 强扩展比较,显示模拟性能(ns/天)和迭代速率(ms/步)随节点数的变化。迭代速率为每个时间步的平均墙钟时间,跨所有时间步计算

未来:开放与可移植的解决方案

这展示了延迟敏感的 HPC 应用在处理 GPU 间通信时的一种更广泛的转变。那些在微秒级内核时间内进行强扩展的应用需要一种与以 CPU 为中心的 MPI 不同的模型。GPU 发起的通信是一个引人注目的答案。

这些技术并非分子动力学所独有。光环交换模式出现在整个计算科学领域(CFD、天体物理学、格点 QCD),并且该方法可以直接迁移到任何依赖于分解域之间边界数据交换的应用中。

一个自然的下一步是将这些想法提升到更广泛的 HPC 生态系统中。诸如 Kokkos Remote Spaces 之类的分布式数据抽象和 std::execution 之类的执行模型已经开始将这些模式封装在可移植接口背后。

更广泛的采用也需要标准化:这里使用的 GPU 原生扩展目前不属于核心 OpenSHMEM 规范,需要正式采纳设备发起的通信、细粒度信令和直接对等内存访问,以实现统一的可用性。

目标是未来来自 GROMACS 这类专门代码的创新成为任何在异构硬件上扩展的应用的标准构建块。如果你运行 GROMACS,不妨在 GROMACS 2026 中尝试 GPU 发起的光环交换,并测量你自己模拟中的强扩展收益。

要将相同的模式应用到其他地方:

  • 探索 NVSHMEM,寻找关键路径上的 CPU 编排通信。
  • 阅读 SC‘25 研讨会论文,了解完整设计和基准测试。
来源:https://www.bestblogs.dev/article/85a83ab8?utm_source=rss&utm_medium=feed&utm_campaign=resources&entry=rss_article_item

相关热点

继续查看同栏目近期热点。

延伸阅读

补充最近整理过的热点入口。