1. NVSHMEM 概述：从共享内存到GPU直通

各位同学好，我是老张。在HPC领域摸爬滚打了十几年，从MPI时代一路走到GPU集群。今天咱们聊聊NVSHMEM——这个让GPU之间可以直接“对话”的通信库。

说实话，我第一次接触NVSHMEM时，心里是有点抵触的。心想：“又来个新框架？”但真正用上之后，嗯，真香。它解决了一个困扰我多年的痛点：GPU之间的数据搬运太慢了。

什么是NVSHMEM？

NVSHMEM，全称NVIDIA Shared Memory。它是一套基于PGAS（Partitioned Global Address Space）模型的通信库。说白了，它让多个GPU看起来像在共享同一块内存。

你想想看，传统做法是：GPU A要数据，得先拷回CPU，再发给GPU B。NVSHMEM直接让GPU A往GPU B的内存里写数据——一步到位。

核心思想：每个GPU都有一块“对称内存”，所有GPU都能直接读写这块内存。就像大家共用一块黑板，谁都能在上面写字。

NVSHMEM与CUDA的关系

这个问题我经常被问到。其实很简单：NVSHMEM是建立在CUDA之上的。

  • CUDA 是基础，负责GPU上的计算和基本内存管理
  • NVSHMEM 是扩展，专门处理多GPU之间的通信

我在项目中遇到过一种情况：用CUDA的cudaMemcpy做多GPU通信，代码写得跟蜘蛛网似的。后来换成NVSHMEM，代码量直接砍掉一半。为什么？因为NVSHMEM把通信细节封装好了。

我的建议：计算部分用CUDA，通信部分用NVSHMEM。各司其职，别混着用。

NVSHMEM与MPI的区别

这个问题值得好好说说。MPI是HPC的老大哥，NVSHMEM是后起之秀。它们有什么区别？我画了张图，一看就明白。

MPI 通信模型 CPU 0 内存 CPU 1 内存 显式消息传递 GPU 0 GPU 1 需要CPU参与中转 通信延迟高 NVSHMEM 通信模型 GPU 0 对称内存 GPU 1 对称内存 直接GPU↔GPU 全局共享地址空间 GPU 0 和 GPU 1 看到相同的内存布局 无需CPU参与 延迟低，带宽高

从图上能看出来：MPI是“你发我收”的模式，NVSHMEM是“我直接写你的内存”。

对比项 MPI NVSHMEM
通信模式 双边（Send/Recv） 单边（Put/Get）
CPU参与 必须 可选
延迟 较高（微秒级） 低（纳秒级）
编程复杂度 中等 较低
适用场景 CPU集群、异构系统 多GPU系统

注意：NVSHMEM不是要取代MPI。我见过有人非要二选一，结果把自己坑了。实际项目中，MPI负责节点间通信，NVSHMEM负责节点内GPU通信，两者配合才是王道。

核心概念：PE、对称内存、Put/Get

好，现在咱们来拆解NVSHMEM的三个核心概念。这些是基本功，必须吃透。

PE（Processing Element）

PE就是“处理单元”。在NVSHMEM的世界里，每个GPU就是一个PE。比如你有4块GPU，那就是4个PE，编号从0到3。

我习惯把PE想象成“工位”。每个工位上有一个人（GPU），有自己的工具（内存），还能直接伸手去拿隔壁工位的东西。

关键点：PE的编号是全局的。PE 0可以访问PE 2的内存，不需要经过PE 1。这就是“单边通信”的精髓。

对称内存

对称内存，这是NVSHMEM最巧妙的设计。每个PE都有一块内存，大小相同，地址对齐。所有PE上的对称内存组成一个“全局共享地址空间”。

举个例子：

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

// PE 0 可以直接写 PE 1 的内存
if (my_pe == 0) {
    nvshmem_float_put(&sym_mem[0], &data, 1, 1);  // 把data写到PE 1的sym_mem[0]
}

我曾经踩过一个坑：对称内存必须所有PE同时分配。如果你在PE 0上分配了，PE 1没分配，那程序直接崩。嗯，血的教训。

避坑指南：对称内存的分配和释放，必须在所有PE上同步进行。我一般用nvshmem_barrier_all()来保证同步。

Put/Get操作

Put和Get是NVSHMEM的两种基本通信操作：

  • Put：把数据写到远程PE的内存里
  • Get：从远程PE的内存里读取数据

这两种操作都是“单边”的——发起方说了算，远程PE不需要配合。这跟MPI的Send/Recv完全不同。

看个完整的例子：

#include <nvshmem.h>
#include <stdio.h>

int main(void) {
    nvshmem_init();
    int my_pe = nvshmem_my_pe();
    int npes = nvshmem_n_pes();
    
    // 分配对称内存
    int *buf = (int*)nvshmem_malloc(sizeof(int));
    *buf = my_pe;  // 每个PE存自己的编号
    
    nvshmem_barrier_all();  // 同步，确保所有PE都准备好了
    
    // PE 0 读取 PE 1 的数据
    if (my_pe == 0) {
        int remote_val;
        nvshmem_int_get(&remote_val, buf, 1, 1);  // 从PE 1读取
        printf("PE 0 got value %d from PE 1\n", remote_val);
    }
    
    // PE 1 把数据写到 PE 0
    if (my_pe == 1) {
        int val = 42;
        nvshmem_int_put(buf, &val, 1, 0);  // 写到PE 0
    }
    
    nvshmem_barrier_all();
    nvshmem_free(buf);
    nvshmem_finalize();
    return 0;
}

核心要点：

  • Put操作：源在本地，目标在远程
  • Get操作：源在远程，目标在本地
  • 两者都是异步的，需要显式同步

说到异步，我得提醒一句：Put/Get操作默认是异步的。也就是说，函数返回时数据不一定已经到达。如果你需要确保数据到位，得调用nvshmem_quiet()或者nvshmem_barrier_all()。

我刚开始用NVSHMEM时，就因为这个异步特性吃过亏。数据还没写完，我就开始读了，结果读到一堆垃圾。后来养成了习惯：每次Put/Get之后，如果后续操作依赖这个数据，一定加个同步。

重要：NVSHMEM的Put/Get操作不保证顺序。如果你连续发了两个Put到同一个PE，第二个可能先到。需要顺序保证的话，用nvshmem_fence()。

好了，第一章的内容就到这里。NVSHMEM的核心思想其实不复杂：让GPU之间能直接读写内存，省去中间环节。后面的章节我们会深入每个操作的细节，包括性能优化、同步机制、以及实际项目中的应用技巧。

记住一句话：NVSHMEM不是魔法，它只是把GPU之间的通信路径缩短了。但就是这一步缩短，能带来数量级的性能提升。

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

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