NVSHMEM 编程模型：PE、对称内存、Put/Get 操作，与 MPI 的对比

好，咱们今天聊聊 NVSHMEM 的编程模型。说实话，很多从 MPI 转过来的朋友，第一次看到 NVSHMEM 的 API 都会觉得眼熟，但又说不上哪里不一样。我当年刚接触的时候也是这种感觉——明明都是通信库，怎么用起来感觉完全两码事？

别急，咱们一步步拆开看。

1. PE：不只是进程，更是 GPU 的“代言人”

NVSHMEM 里有个核心概念叫 PE（Processing Element）。说白了，每个 PE 就是一个参与通信的“计算单元”。

在 MPI 里，我们习惯叫 rank。每个 rank 对应一个进程。但在 NVSHMEM 里，每个 PE 通常对应 一个 GPU + 一个 CPU 线程。你想想看，这其实更贴近真实硬件——GPU 才是真正干活的主力，CPU 线程只是负责“发号施令”。

关键区别：MPI 的 rank 是进程级别的抽象，而 NVSHMEM 的 PE 是 GPU 级别的抽象。一个节点上如果有 4 块 GPU，那通常就有 4 个 PE。

我记得第一次在 DGX-1 上跑 NVSHMEM 程序时，我习惯性地开了 8 个 MPI 进程，结果发现每个进程都试图初始化 8 块 GPU...嗯，那场面相当混乱。后来才明白，PE 的数量应该等于 GPU 的数量，而不是 CPU 核心数。

获取当前 PE 编号和总 PE 数量的方式很简单：

int my_pe = nvshmem_my_pe();
int n_pes = nvshmem_n_pes();

是不是跟 MPI 的 MPI_Comm_rank 和 MPI_Comm_size 很像？但底层逻辑完全不同。

2. 对称内存：所有 GPU 都能看到的“共享空间”

对称内存是 NVSHMEM 最让我惊艳的设计之一。什么叫对称？就是 每个 PE 都分配一块同样大小的内存，而且所有 PE 都能直接访问对方的这块内存。

你想想看，这相当于在多个 GPU 之间建立了一个“虚拟的共享地址空间”。PE 0 上的 GPU 可以直接读写 PE 1 上的对称内存，不需要 CPU 参与中转。

分配对称内存的代码长这样：

// 每个 PE 分配 1024 个 float
float *sym_mem = (float*)nvshmem_malloc(1024 * sizeof(float));

// 初始化
for (int i = 0; i < 1024; i++) {
    sym_mem[i] = (float)my_pe;
}

// 所有 PE 同步，确保数据就绪
nvshmem_barrier_all();

避坑指南：我曾经在分配对称内存时忘了调用 nvshmem_barrier_all()，结果 PE 1 去读 PE 0 的数据时，PE 0 还没写完。嗯，数据全是乱的。记住：对称内存的“对称”是空间上的，不是时间上的——你需要自己保证同步。

对称内存的地址在所有 PE 上都是“有效”的。PE 0 上的 sym_mem 指针，在 PE 1 上也能用——当然，PE 1 用这个指针访问的是 PE 0 的那块内存。这就是 NVSHMEM 的魔力所在。

3. Put/Get 操作：GPU 直读直写，CPU 靠边站

Put 和 Get 是 NVSHMEM 最核心的通信操作。说白了：

  • Put：把数据从“我”的对称内存，写到“别人”的对称内存
  • Get：把数据从“别人”的对称内存，读到“我”的对称内存

而且这些操作都是 GPU 直接发起的，不需要 CPU 参与。这意味着你的 GPU kernel 里可以直接调用 Put/Get，实现真正的 GPU 间直通。

看个例子：

// 在 GPU kernel 内部
__global__ void my_kernel(float *dest, float *src, int nelems) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid < nelems) {
        // PE 0 把数据 Put 到 PE 1
        if (nvshmem_my_pe() == 0) {
            nvshmem_float_put(&dest[tid], &src[tid], 1, 1);
        }
    }
}

这里 nvshmem_float_put 的四个参数分别是：目标地址、源地址、元素个数、目标 PE 编号。是不是很直观？

对应的 Get 操作：

// PE 0 从 PE 1 的对称内存中 Get 数据
nvshmem_float_get(&my_data, &remote_data, 1, 1);

注意：Put/Get 操作是 单边通信。也就是说，发起方直接读写目标方的内存，目标方完全不知道这件事。这在 MPI 里是做不到的——MPI 的 Send/Recv 必须双方配对。单边通信的好处是延迟更低，但坏处是你要自己保证数据一致性。

4. 与 MPI 的对比：一张表说清楚

我经常被问到：“NVSHMEM 和 MPI 到底选哪个？”我的回答是：看场景。下面这张表是我根据自己的项目经验总结的：

特性 MPI NVSHMEM
通信模式 双边（Send/Recv 配对） 单边（Put/Get 直读直写）
数据路径 CPU 发起，经过系统内存 GPU 直接发起，GPU 间直通
内存模型 分布式，每个进程私有 对称内存，全局可寻址
同步方式 隐式同步（Send 完成后自动同步） 显式同步（需要 barrier 或 fence）
延迟（小消息） 微秒级（受 CPU 调度影响） 纳秒级（GPU 直接执行）
带宽（大消息） 受 PCIe/NVLink 带宽限制 充分利用 NVLink 直连带宽
编程复杂度 较低，API 成熟稳定 中等，需要理解 GPU 内存模型
适用场景 传统 HPC，CPU 主导的异构计算 GPU 密集通信，如深度学习训练、GNN

我个人习惯是：如果通信模式是“规整的、可预测的”，比如每个 GPU 定期交换边界数据，那 MPI 就够用了。但如果通信是“稀疏的、动态的”，比如图神经网络里每个节点随机访问邻居的特征，那 NVSHMEM 的单边通信优势就体现出来了。

5. 核心逻辑框架图

下面这张图展示了 NVSHMEM 编程模型的核心逻辑。我画的时候特意把 PE、对称内存和 Put/Get 的关系理清楚了：

NVSHMEM 编程模型核心逻辑 PE 0 CPU 线程 0 GPU 0 对称内存块 0 地址：0x7f...000 PE 1 CPU 线程 1 GPU 1 对称内存块 1 地址：0x7f...100 Put GPU 0 → GPU 1 Get GPU 1 → GPU 0 对称内存：全局可寻址 PE 0 的指针在 PE 1 上同样有效，指向 PE 0 的内存 与 MPI 的关键区别 MPI：双边通信，Send/Recv 必须配对，CPU 参与数据搬运 NVSHMEM：单边通信，Put/Get 直读直写，GPU 直接执行

6. 实际项目中的选择建议

我在做图神经网络训练框架时，遇到过这样一个场景：每个 GPU 负责处理一部分图节点，训练时需要频繁访问邻居节点的特征。如果用 MPI，每个 GPU 都要先发请求、等回复，通信开销占了训练时间的 40%。

后来改成 NVSHMEM，直接用 Put/Get 在 GPU kernel 里读取邻居特征，通信开销降到了 10% 以下。为什么？因为 MPI 的 Send/Recv 需要 CPU 参与调度，而 NVSHMEM 的 Put/Get 是 GPU 直接执行的，省掉了 CPU 这个“中间商”。

我的建议：如果你的程序里通信模式是“一对多”或“多对一”，而且通信频率很高，那 NVSHMEM 的单边通信会带来显著的性能提升。但如果你的通信模式是规整的“邻居交换”，MPI 的集体操作（Allreduce、Bcast 等）可能更高效——毕竟 MPI 对这些操作做了深度优化。

嗯，NVSHMEM 的编程模型其实就这三个核心概念：PE 是通信的主体，对称内存是通信的媒介，Put/Get 是通信的手段。理解这三者的关系，你就掌握了 NVSHMEM 的精髓。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →