3. NVSHMEM基础：对称内存模型、PE概念、Put/Get操作
好，咱们进入NVSHMEM的核心地带。说实话，我第一次接触NVSHMEM的时候，最困惑的就是它那个“对称内存”到底是个啥玩意儿。跟普通的MPI或者NCCL完全不是一个路子。今天我就带你把它掰开揉碎了讲清楚。
3.1 对称内存模型：所有PE共享的“虚拟连续空间”
先说说对称内存模型。你想想看，在分布式训练里，每个GPU都有自己的显存，数据是分散的。NVSHMEM的做法很巧妙——它把所有参与计算的GPU显存，逻辑上“拼接”成一个巨大的、连续的虚拟地址空间。
什么意思呢？
- 每个PE（Processing Element，处理单元，通常对应一个GPU）都有一块“对称堆”（Symmetric Heap）。
- 这块堆在每个PE上的虚拟地址是完全一样的。
- 你在PE 0上分配了一个变量，它在PE 1、PE 2...上都有对应的“镜像”。
说白了，就是“同一个地址，不同物理位置”。
核心要点：对称内存的地址在所有PE上是相同的。你不需要知道数据在哪个GPU上，只需要知道它的“全局地址”就行。NVSHMEM底层帮你搞定路由和传输。
我在项目中遇到过一个问题：刚开始用NVSHMEM时，我习惯性地在每个PE上单独malloc，结果发现数据根本对不上。后来才意识到，必须用nvshmem_malloc来分配对称内存。嗯，这个坑我替你踩过了。
3.2 PE（Processing Element）概念：分布式计算的基本单元
PE是NVSHMEM里的基本执行单元。你可以把它理解为一个“工人”，每个工人有自己的计算资源（GPU）和存储空间（显存中的对称堆）。
PE的编号从0到npes-1，其中npes是总的PE数量。每个PE可以通过nvshmem_my_pe()获取自己的编号，通过nvshmem_n_pes()获取总PE数。
| 函数 | 作用 | 我常用的场景 |
|---|---|---|
nvshmem_my_pe() |
返回当前PE的ID | 判断自己是不是主节点（PE 0） |
nvshmem_n_pes() |
返回总PE数量 | 计算数据分片大小 |
nvshmem_barrier_all() |
所有PE同步 | 确保数据就绪后再操作 |
你可能会问：“PE和GPU是一一对应的吗？” 嗯，通常是的，但也不绝对。NVSHMEM允许一个GPU上跑多个PE（通过CUDA streams实现），但我个人建议一个PE绑定一个GPU，这样逻辑清晰，性能也最好。
3.3 Put/Get操作：远程内存访问的“读写指令”
Put和Get是NVSHMEM最核心的通信原语。它们允许一个PE直接读写另一个PE的对称内存，而不需要对方参与。这跟NCCL的集体通信（AllReduce、Broadcast）完全不同——NCCL需要所有节点协同，而NVSHMEM的Put/Get是“单边操作”。
3.3.1 Put操作：把数据“推”到远程PE
Put操作，说白了就是“我把我的数据写到你那边去”。
// 示例：PE 0 向 PE 1 的对称内存写入数据
float *src = (float*)nvshmem_malloc(1024 * sizeof(float));
float *dst = (float*)nvshmem_malloc(1024 * sizeof(float));
// 初始化src（只在PE 0上）
if (nvshmem_my_pe() == 0) {
for (int i = 0; i < 1024; i++) src[i] = i * 1.0f;
}
// PE 0 把src的数据put到PE 1的dst
nvshmem_float_put(dst, src, 1024, 1); // 目标PE是1
nvshmem_barrier_all(); // 确保所有PE都完成
// 现在PE 1上的dst数组就有了数据
if (nvshmem_my_pe() == 1) {
printf("dst[0] = %f\n", dst[0]); // 输出 0.0
}这里要注意：nvshmem_float_put的第一个参数是远程PE上的目标地址，第二个参数是本地源地址。我曾经搞反过，结果数据写到了自己头上，排查了半天才发现。
我的习惯：写Put/Get代码时，我会在注释里明确标注“本地”和“远程”，避免混淆。比如：// dst: 远程PE 1上的地址
3.3.2 Get操作：从远程PE“拉”数据
Get操作正好相反——我从别人那里把数据拿过来。
// 示例：PE 1 从 PE 0 获取数据
float *local_buf = (float*)nvshmem_malloc(1024 * sizeof(float));
float *remote_buf = (float*)nvshmem_malloc(1024 * sizeof(float));
// 假设PE 0已经初始化了remote_buf
nvshmem_barrier_all();
// PE 1 从 PE 0 的remote_buf读取数据到自己的local_buf
if (nvshmem_my_pe() == 1) {
nvshmem_float_get(local_buf, remote_buf, 1024, 0); // 源PE是0
printf("Got data: %f\n", local_buf[0]);
}Get操作有个好处：你可以在需要的时候才去拉数据，不用提前push。这在动态数据流场景下特别有用。
3.3.3 Put vs Get：什么时候用哪个？
| 操作 | 适用场景 | 我个人的选择倾向 |
|---|---|---|
| Put | 主动推送数据（如梯度更新） | 数据生产者明确知道消费者是谁时 |
| Get | 按需拉取数据（如参数查询） | 消费者不确定数据何时更新时 |
举个例子：在分布式训练中，梯度同步我倾向于用Put——每个GPU算完梯度后主动push到其他GPU。而参数读取我倾向于用Get——每个GPU在需要时才去拉最新的参数。
避坑指南：我曾经在循环里频繁调用Get操作，结果发现性能极差。后来才意识到，Get是阻塞的，会等待数据从远程传输完成。如果数据量小、频率高，建议用Put配合非阻塞接口（如nvshmem_float_put_nbi）来提升吞吐。
3.4 知识体系总览
下面这张图是我自己总结的NVSHMEM基础架构，你看一眼就能明白这三者的关系：
你看，每个PE的对称堆地址都是0x1000，但物理上分布在不同的GPU显存里。Put操作就像你把自己的东西快递给别人，Get操作就像你去别人家取东西。NVSHMEM负责搞定快递路线（NVLink、PCIe、网络），你只需要告诉它“寄什么、寄给谁、寄多少”。
好了，对称内存模型、PE概念、Put/Get操作，这三块就是NVSHMEM的基石。我个人觉得，理解对称内存是第一步，搞懂PE是第二步，熟练运用Put/Get是第三步。这三步走稳了，后面再讲原子操作、集体通信、动态共享内存，你就不会觉得吃力了。