CUDA on Fain的Blog

LLM学习日志 #8 NVSHMEM初探

Fri, 29 May 2026 17:03:31 +0800

最近在研究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大人的整理和解释，请读者注意甄别。

一、基础设施类

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

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

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

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

集合通信Collectives：barrier/sync（barrier_all/sync_all）、broadcast、fcollect（即allgather）、alltoall，以及归约：sum/prod/min/max/and/or/xor。语义上和NCCL的allreduce/allgather/alltoall是对位的。
点对点同步与内存序：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的函数指针。如果不这样，会在运行时报错。

LLM 学习日志 #6 巨型内核：MegaKernel

Thu, 21 May 2026 23:12:39 +0800

这是最近笔者很感兴趣的一个新技术，因此对其进行了一些论文的阅读了解。在此整理为一篇简短的日志，以备忘和供读者参考和了解相关的技术。

什么是MegaKernel

MegaKernel（巨型内核）是一种 GPU 推理优化技术，其核心思想是把整个LLM的推理过程，包括所有计算、跨 GPU 通信融合(fuse)进一个单独的 GPU kernel 中执行，而不是像传统方式那样把每个算子(operator)启动成一个独立的 kernel。

它在某种程度上和 CUDA Graph 在思想上有相似性（减少launch开销），但比 CUDA Graph 的录制-重放方案激进的多也彻底的多，当然也因此被施加了更明显的限制和工程复杂性。

传统 LLM 推理通常依赖一长串 GPU kernel launch 和外部通信调用，这导致硬件利用率不足。巨型内核通过端到端的 GPU 融合方式,把跨多层、多迭代、多 GPU 的操作全部融合进一个 kernel，从而消除启动开销，实现细粒度的软件流水线,并让计算与通信重叠。

最近Megakernel也开始逐渐被用于训练的生产环境，这是一个值得注意的进展。

值得关注的社区进展

虽然说是新技术，其实也不是很新了。它是在2025年由几个团队几乎同期推动起来的方向。具体而言又分为两条方向的线索：

1. 斯坦福 Hazy Research 团队(Tri Dao 等人所在组)

2025 年 3 月发布了 ThunderMLA，这是一个针对 DeepSeek MLA 解码的融合 megakernel，在多种工作负载下比 DeepSeek 的 FlashMLA 快 20–35%。随后在 5 月发布了著名的 "No Bubbles" 博客《Designing a Low-Latency Megakernel for Llama-1B》，指出 vLLM、SGLang 等主流推理引擎在 H100 上跑 Llama-3.2-1B 单序列时，GPU 带宽利用率最多只有 50%，而 megakernel 方法可以大幅突破这个瓶颈。这部分工作主要注重开发手工设计的高性能megakernel。

2. CMU 贾志豪(Zhihao Jia)团队

主要开发的是Mirage Persistent Kernel (MPK)，注重系统性的自动化编译器路线。该项目由 CMU、UW、Berkeley、NVIDIA、清华的联合团队开发，2025 年 6 月开源，核心成员包括 Xinhao Cheng、Mengdi Wu、Xupeng Miao、陈天奇(Tianqi Chen)、贾志豪等等知名研究者。MPK 是首个能将多 GPU 模型推理自动转换为单个高性能 megakernel 的编译器和运行时系统，引入了 SM 级别(streaming multiprocessor)的图表示来捕获数据依赖，从而实现跨算子软件流水线、细粒度 kernel 重叠等以前不可行的 GPU 优化。

值得关注的论文：

《Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs》(arXiv:2512.22219)

《Mirage: A Multi-Level Superoptimizer for Tensor Programs》

《Event Tensor: A Unified Abstraction for Compiling Dynamic Megakernel》 Megakernel编译抽象

《UniEP: Unified Expert-Parallel MoE MegaKernel for LLM Training》(arXiv:2604.19241) 关于MoE+训练场景+MegaKernel

MegaKernel的核心抽象和核心优势

SM 级任务图(SM-level task graph)

MPK 最有标志性的抽象。众所周知传统编译器(PyTorch、TVM、Triton)的最小调度单位是整张 GPU 上的一个 kernel。MPK 提出在单个流式多处理器(SM)粒度上表示计算和 GPU 间通信,引入了一个叫 tGraph 的 SM 级图表示，节点是跑在单个 SM 上的任务，边是任务间的细粒度依赖。

这意味着调度器/编译器处理的粒度不再是类似"matmul - attention - matmul"的结构，而是"matmul 的第 5 个 tile 在 SM 17 完成 - attention 中依赖这个 tile 的任务现在可以在 SM 17 开始"这样的结构。它本身可以消除一部分jitter和流水线空泡，但是和下文的动态调度结合起来更有威力。

片上解释器(on-GPU interpreter)/持久内核

持久内核作为 GPU 性能优化概念已经被广泛认识和采用，但 MegaKernel 明显将其推进到进一步的高度。整个 megakernel 启动一次后,每个 SM 都初始化一个解释器，开始执行一串指令序列(序列在 launch 之前调度到每个 SM 的队列中)。解释器通过类似编译器的指令重排的方式进行流水线填充，从而消除操作之间的内存气泡。

静态和动态调度

该说不说，GPU 终于发展到了能让 CPU 失业的临界点。如果动态调度成为 GPU 的标配方案，CPU 可能基本上就会变成 GPU 的启动器。让我们观察它的未来发展。

静态调度 在编译期就把指令切分到每个 SM 的私有队列。好处是可以预取并把多条指令直接载入共享内存,取指开销极低；坏处是要求编译器提前静态划分指令，而指令的延迟是可变的，SM 之间还存在抖动(jitter)，会造成负载不均。

动态调度 用一个全局指令队列让所有 SM 抢占式领取。取指开销更高(需要 global memory 往返加原子锁)，但可以隐藏在上一条指令执行期间，不需要编译器提前分配指令。当一个 SM 准备好时，从全局队列中取出一条指令执行。这个方案的好处是自然对 jitter 自适应，快的 SM 得到更多工作，因此具有负反馈的自均衡特性。当静态调度的空泡累积的较多的时候，优势就会很明显。

MegaKernel的发展前景

这套机制的代价主要是工程复杂度和动态性支持。现代 LLM 服务负载本身是动态和不可预测的，因此需要 continuous batching 甚至变长输入之类的方案。要在单个 megakernel 里支持这些动态形状，常常需要为每种可能的形状重新生成或重新编译 kernel，启动延迟会变得不可接受，甚至当形状空间过大时根本做不到。现有的 megakernel 框架目前主要面向 single-batch 推理，在动态形状或数据相关负载下还不能与通用推理系统公平对比。这被认为是未来一两年这个方向的研究重点。