学习日志

LLM学习日志 #8 NVSHMEM初探

最近在研究Kernel内的通信发起方式，调研到的结果是NVSHMEM。尝试了一下，感觉它的api还挺奇怪的，和CUDA本身的还不完全一样。在这里简单记录一下几个需要额外记忆的点，以避免遗忘，也供大家参考。

NVSHMEM是什么

NVSHMEM是同时结合了RDMA的单边操作（one-sided communication）和CUDA kernel内部的GPU发起通信模式的高性能网络并行编程接口，基于C语言。用更科班的方法说，它是NVIDIA基于OpenSHMEM标准实现的一套并行编程接口，核心抽象模型是基于对称堆的分块全局地址空间（PGAS，Partitioned Global Address Space）。通过在每个GPU上分配完全对称的内存空间，每块GPU就能够直接读写其他GPU的显存，无论对端是通过高速的NVLink进行节点内的互联，还是通过InfiniBand等等高性能网络连接。

它的主要开发目的是解决通信当中的CPU关键路径开销问题。如果写过MPI、pytorch或者各种普通CUDA Kernel的模型的话大概就知道，早期最最传统的流程是GPU算完一批数据，先把结果搬回CPU，CPU通过MPI发给另一台机器，对方的CPU把数据搬到它的GPU上。这需要基于PCIe通信的多次往返，因此会造成一个很大的开销。进一步的说，后来现代一些的CUDA-Aware MPI，以及著名的NCCL，其工作方式主要是让CPU发起GPU的GPUDirect RDMA，在那之后GPU自己负责从一张卡将数据直接搬运到NIC或者通过NVLink搬运到另一张卡，CPU的内存不负责搬运数据。NVSHMEM提供一种更进一步的手段，即让通信的发起也是在GPU侧的（因为是写在CUDA Kernel里的代码触发的），而CPU对于整个流程几乎完全没有介入的地方。

顺带一提，另一个很重要的类似地位的东西叫NCCL GIN（NCCL GPU-Initiated Networking，GPU 发起网络通信），一个比NVSHMEM晚一些出现的API接口。GIN（GPU-Initiated Networking）是NCCL 2.28引入的Device API里的一个模式。这套Device API有三种模式：LSA（Load/Store Accessible）走节点内的NVLink/PCIe，Multimem走NVLink SHARP做硬件多播，而GIN负责跨节点的InfiniBand/RoCE网络RDMA。它要解决的问题和NVSHMEM其实是一样的，同样是让GPU从CUDA kernel内部直接发起网络操作，省掉CPU协调的开销。但同时，GIN存在于NCCL自己的统一运行时里，能把device端发起的细粒度操作和NCCL原有的集合算法、生产级基础设施整合在一起，从而允许用户不用离开NCCL，也不用整个迁移到SHMEM的编程模型。有空也得学一下。关键概念是GDAKI（GPUDirect Async Kernel-Initiated），指"由GPU kernel自己发起网络传输"这个能力本身；IBGDA（InfiniBand GPUDirect Async） ，是NVSHMEM的一个InfiniBand transport，也就是GDAKI的具体落地；IBRC（InfiniBand Reliable Connection） 是NVSHMEM另一个、也更早的InfiniBand transport，建在RC（可靠连接）队列对上，在IBGDA出现之前，IBRC靠一个CPU上的proxy thread来管理通信：kernel里发起的操作先把一个work descriptor写进host内存的proxy buffer，CPU上的代理线程发现后，再去触发对应的网络操作。NCCL GIN和NVSHMEM的底层技术是完全相同的，都是所谓的GDAKI。NVSHMEM是OpenSHMEM/PGAS模式，对称堆、基于指针的对称寻址；GIN采用的是MPI RMA window模型，用ncclCommWindowRegister把各rank自己的buffer集合注册进一个个内存窗口、拿到指向所有peer的远端key，寻址是基于window的而不是NVSHMEM那种指针加atomics；而且GIN的window允许各rank注册不同大小，不强求严格对称，换言之更接近RDMA verbs那一套。

无论如何，通过将CPU发起通信控制指令的负担从CPU上转移到GPU本身上，这就可以省下相当一部分通信的Overhead开销。当然，它的代价是GPU编程的复杂性和GPU的SM会需要更多操作来发起通信。不过绝大多数时候是值得的。

Helloworld流程

和MPI、torch.distributed一样，NVSHMEM需要init、需要找到自己的world_rank，需要知道world_size，需要在最后finalize。其他领域的常见的那个rank概念在NVSHMEM中被称作PE，全称叫Processing Element，处理单元。具体来说的对应关系是，nvshmem_my_pe()返回的那个mype就是world_rank，nvshmem_n_pes()返回的npes就是world_size。

不过，后面的差异就非常大了。对于MPI、torch.distributed等，init之后就可以直接开始调用各种操作。这是因为其他通信库里默认是可以做rank到GPU的多重映射的，一个GPU可以同时是多个rank，因此系统可以自动进行默认的映射和分配。但是，PE到GPU需要进行一个显式绑定才能达到相同的可操作状态。这似乎是因为一个PE必须绑一块GPU、整个任务期间不变，同时一块GPU也不能被多个PE共享。进一步的，这又和对称堆的分配和PGAS模型有关系：对称堆的内存地址是物理的，无法虚拟化和映射，因此同一块GPU就只能同时支持一个PE，进一步的就需要手动指定每个cudaSetDevice才能够正常工作了。

另一个值得注意的地方是它需要部分的侵入式介入CUDA正常的Memory allocation流程。为了使用NVSHMEM，必须使用它自己的内存分配函数，这是一个集合通信性质的函数，需要所有PE同时调用（想想也是，不然就不叫“对称堆”了）。先看它们的malloc的函数签名：

1
2
3
4
5

// CUDA：返回值是错误码，指针通过出参回填，也就是"非赋值"形式
cudaError_t err = cudaMalloc(&ptr, size);

// NVSHMEM：直接把指针返回回来，跟libc的malloc一个样，也就是"赋值"形式
ptr = nvshmem_malloc(size);

可以看到前者是专门的CUDA风格，后者反而和普通的libc的malloc风格一样，因此在代码里看惯了第一种的读者/工程师可能会很不适应。它不能直接使用malloc，因为malloc并没有“对称堆”的“对称”这个语义保证，毕竟万一你写了什么if xx malloc else do nothing呢？

此外值得一提的是为什么有这么奇怪的风格差异。这其实是个历史遗留问题，因为它压根不是NVIDIA从零设计的API，而是OpenSHMEM这套标准的一个实现，而OpenSHMEM比CUDA老得多，换句话说血统继承不一样。NVSHMEM前缀虽然挂着nv，但my_pe、n_pes、put、get、barrier、malloc这些名字和语义基本是照搬OpenSHMEM的，再往上能一直追到上世纪90年代Cray的SHMEM。这套API在CUDA出生之前就定型了。所以不是它不想和CUDA统一风格，而是它"不能"；API签名已经是它follow的标准的一部分，单方面改掉的话就不配叫OpenSHMEM实现了。（别让老标准害了你啊.jpg）

说它的api设计的很奇怪也是因为这个。如果想对一段代码进行适配NVSHMEM的修改，就不得不改变很多基本性的东西，几乎不可能向后兼容或者无功能性破坏。与之相对的是NCCL的模型，其照常用cudaMalloc分配tensor，要通信时把需要buffer传入ncclAllReduce，对于内存模型没有任何侵入性的改动。NVSHMEM则不同，要先整体接受它的内存模型（重构整个分配的代码），然后才能使用它的GPU端通信能力。它的内存模型和通信模型是焊死在一起的。NVSHMEM本质上还是更接近于把SHMEM那一整套编程模型嫁接到GPU上，从而允许一些老代码方便的对GPU做扩展，反而不是给CUDA runtime本身去做扩展。

NVSHMEM的主要API操作

这方面我无法靠自己完全写出来（我熟悉的就寥寥几个），因此依靠的是ai大人的整理和解释，请读者注意甄别。

一、基础设施类

  1. 初始化与查询：nvshmem_init/nvshmemx_init_attr/nvshmem_finalize，以及nvshmem_my_pe、nvshmem_n_pes、版本/名称查询和线程支持（nvshmem_init_thread/query_thread）。
  2. 对称内存管理：nvshmem_malloc/free/align/calloc等内存分配函数；外加nvshmem_ptr，在peer可直接P2P访问时拿到对方对象的裸指针，从而用普通load/store代替put/get。
  3. 团队管理：nvshmem_team_split_*、team_my_pe/team_n_pes/team_translate_pe，以及一组预定义team（如TEAM_WORLD、TEAMX_NODE等）。team就是OpenSHMEM里communicator/subgroup的对应物，取代了老的active set。

二、通信类（数据怎么读写）

  1. 单边远程读写RMA：put/get两大族——块传输putmem/getmem、单元素nvshmem_TYPE_p/_g、带步长的iput/iget，每种再分阻塞与非阻塞（_nbi后缀）。这是NVSHMEM的主干。
  2. 原子操作AMO：atomic_add/inc/fetch_add/fetch_inc/compare_swap/swap/set/and/or/xor，分取回值和不取回两类，用来在远端做无锁更新。
  3. 信号操作Signaling：put-with-signal（nvshmem_TYPE_put_signal，含非阻塞版）在把数据写到远端后顺带更新一个远端flag，再配合nvshmemx_signal_op和signal-fetch。这是搭配wait/test做点对点同步的利器，DeepEP一类的MoE通信重度依赖它。

三、同步与排序类（怎么保证"看到"和"完成"）

  1. 集合通信Collectives：barrier/sync（barrier_all/sync_all）、broadcast、fcollect（即allgather）、alltoall，以及归约：sum/prod/min/max/and/or/xor。语义上和NCCL的allreduce/allgather/alltoall是对位的。
  2. 点对点同步与内存序：wait_until（带all/any/some/vector变体）、test系列、signal_wait_until用来等远端写入；而nvshmem_fence负责给一连串put排序、nvshmem_quiet负责确保它们真正完成。在CPU上发的fence/quiet只管CPU发起的操作，在GPU上发的只管GPU发起的，两边不互相兜底。

然后是调用方式轴，也是NVSHMEM区别于纯OpenSHMEM的地方。上面绝大多数功能都同时存在于三种形态：

  • 主机端（nvshmem_*）：CPU调用。
  • 流序（nvshmemx_*_on_stream）：CPU把操作按CUDA流的顺序入队，给每个函数末尾多加一个cudaStream_t参数，在GPU上按流序执行。
  • 设备端（kernel内直接调）：GPU发起，而且带协作粒度——同一个集合/批量操作有_warp和_block两种后缀变体，分别由一个warp或一个线程块里的所有线程协同完成，不带后缀则是单线程粒度。

它的api复杂程度甚至远超RDMA的verbs，想全记住的人有福了。

代码示范

一段典型的NVSHMEM的Helloworld代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

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

__global__ void simple_shift(int *dst) {
    int mype = nvshmem_my_pe();
    int npes = nvshmem_n_pes();
    int peer = (mype + 1) % npes;
    nvshmem_int_p(dst, mype, peer);   // 把自己的编号直接写进peer的显存
}

int main(void) {
    nvshmem_init();
    int mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE);
    cudaSetDevice(mype_node);          // 关键：在任何NVSHMEM device调用前绑好GPU

    cudaStream_t stream;
    cudaStreamCreate(&stream);
    int *dst = (int *)nvshmem_malloc(sizeof(int));   // 对称内存

    simple_shift<<<1, 1, 0, stream>>>(dst);
    nvshmemx_barrier_all_on_stream(stream);          // 等所有PE都写完

    int msg;
    cudaMemcpyAsync(&msg, dst, sizeof(int), cudaMemcpyDeviceToHost, stream);
    cudaStreamSynchronize(stream);
    printf("PE %d/%d received %d\n", nvshmem_my_pe(), nvshmem_n_pes(), msg);

    nvshmem_free(dst);
    nvshmem_finalize();
    return 0;
}

值得注意的是，如果Kernel内部需要进行NVSHMEM专属的barrier等操作，就不能直接使用kernelname<<<...,...>>>这种launch方式了，需要使用NVSHMEM专属的collective launch函数，传入原本的Kernel的函数指针。如果不这样，会在运行时报错。