一、NVSHMEM 是什么：从起源到定位
大家好，我是你们这趟 NVSHMEM 之旅的向导。今天咱们聊聊 NVSHMEM 的来龙去脉。
先问个问题：你写 GPU 程序时，最头疼的是什么？
我个人觉得，是数据搬运。你把数据从 CPU 拷到 GPU，算完再拷回来。单卡还好，多卡协作时，那叫一个痛苦。每个 GPU 有自己的显存，数据不能直接访问邻居的。你得手动调用 cudaMemcpy，还得操心同步。
嗯，NVSHMEM 就是来解决这个痛点的。
NVSHMEM 的起源与背景
NVSHMEM 的全称是 NVIDIA Shared Memory，但它不是我们常说的 GPU 上的共享内存（Shared Memory）。它借鉴了另一个经典库——SHMEM。
SHMEM 是什么？它诞生于 90 年代，专为 Cray 超级计算机设计。核心思想很简单：每个处理器有自己的局部内存，但你可以通过一套 API 直接读写别人的内存。不需要像 MPI 那样发消息、收消息，直接 put/get 就行。
NVIDIA 在 2019 年把 SHMEM 搬到了 GPU 上，推出了 NVSHMEM。说白了，就是让 GPU 之间能像访问自己的显存一样，直接读写对方的显存。
核心要点：NVSHMEM 是一套用于 GPU 集群的通信库，它实现了 PGAS（分区全局地址空间） 编程模型。你写代码时，感觉所有 GPU 的显存是一个整体，但实际上每个 GPU 只管理自己那一块。
我记得第一次接触 NVSHMEM 时，心里犯嘀咕：这不就是 MPI 吗？后来发现完全不是一回事。MPI 是消息传递，你发我收。NVSHMEM 是内存访问，你直接写我的地址。性能差距很大。
PGAS 编程模型简介
PGAS，全称 Partitioned Global Address Space。名字挺长，拆开看就懂了：
- Global Address Space：全局地址空间。所有 GPU 的显存被映射到一个统一的虚拟地址空间里。
- Partitioned：分区的。每个 GPU 只拥有其中一部分，这部分叫“本地分区”。
你想想看，写代码时你只需要知道数据的“全局地址”，就能直接读写它。至于数据在哪个 GPU 上，库帮你搞定。
举个例子：
// 每个 PE（Processing Element，即 GPU）分配 1024 个 float
float *ptr = (float*)nvshmem_malloc(1024 * sizeof(float));
// PE 0 把数据写到 PE 1 的地址上
if (my_pe == 0) {
nvshmem_float_put(&ptr[0], &src_data, 1, 1); // 把 src_data 写到 PE 1 的 ptr[0]
}
// PE 1 直接读取自己的 ptr[0]，数据已经在了
if (my_pe == 1) {
float val = ptr[0]; // 无需同步，数据已就位
}看到了吗？PE 0 直接往 PE 1 的内存里写数据。PE 1 不需要主动接收。这就是 PGAS 的精髓——通信变成了内存访问。
避坑指南：我曾经犯过一个低级错误——以为 nvshmem_malloc 分配的内存是全局可见的，就直接在 CPU 端访问。不行！NVSHMEM 的分配只在 GPU 端有效。CPU 端要访问，得用 cudaMemcpy 或者 nvshmemx_ 开头的对称分配接口。切记。
PGAS 模型有几种变体：
| 模型 | 特点 | 代表 |
|---|---|---|
| UPC（统一并行 C） | 语言扩展，编译时确定数据分布 | UPC、CAF |
| SHMEM | 库形式，运行时动态通信 | NVSHMEM、OpenSHMEM |
| GASNet | 底层通信层，支持多种模型 | GASNet-EX |
NVSHMEM 属于 SHMEM 这一支。它不改变你的编程语言，只提供一套 API。你写 CUDA C/C++ 代码，直接调用 nvshmem_put、nvshmem_get 就行。
NVSHMEM 在 GPU 集群中的定位
现在 GPU 集群越来越常见。一个节点插 4 张、8 张 GPU，甚至跨节点互联。怎么让这些 GPU 高效协作？
传统方案有：
- MPI + CUDA：CPU 端发消息，GPU 端计算。数据要经过 CPU 中转，延迟高。
- NCCL：专为 GPU 设计的集合通信库，适合 AllReduce、Broadcast 等操作。但它是“集体”的，不适合点对点随机访问。
- CUDA 感知的 MPI：GPU 可以直接传递指针给 MPI，但底层还是消息传递。
NVSHMEM 的定位是：填补点对点、细粒度、低延迟通信的空白。
我举个例子你就明白了。假设你在做图计算，每个 GPU 负责一部分节点。计算时，一个节点可能需要频繁访问邻居节点的数据。邻居可能在别的 GPU 上。
用 MPI 怎么做？你得先打包数据，发送，对方接收，解包。每次通信开销很大。
用 NVSHMEM 呢？直接 put/get。就像访问本地数组一样。延迟能降到微秒级。
NVSHMEM 的典型应用场景：
- 稀疏线性代数（SpMV、SpGEMM）
- 图计算（BFS、PageRank）
- 粒子模拟（PIC、N-body）
- 深度学习中的模型并行（参数服务器模式）
当然，NVSHMEM 不是万能的。它不适合大数据块传输（那是 NCCL 的强项）。它也不适合 CPU 和 GPU 之间的通信（那是 cudaMemcpy 的事）。
说白了，NVSHMEM 是给那些需要 GPU 之间频繁、小数据量、低延迟通信的场景准备的。
注意：NVSHMEM 要求所有 GPU 必须通过 NVLink 或 InfiniBand 互联。如果只是通过 PCIe 连接，性能会大打折扣。我在一个项目中试过 PCIe 下的 NVSHMEM，延迟比 NVLink 高了 5 倍。所以，硬件条件很重要。
最后，我用一张图来总结 NVSHMEM 在整个 GPU 编程生态中的位置：
从这张图你能看到，NVSHMEM 位于编程模型层和通信库层之间。它向上提供 PGAS 抽象，向下直接操作 GPU 显存和互联网络。说白了，它让 GPU 集群编程变得像单机多线程一样简单。
好了，这一章就到这里。NVSHMEM 是什么，你应该心里有数了。下一章我们聊聊怎么搭建环境，让你能亲手跑第一个 NVSHMEM 程序。