★ FEATURED ARTICLE 1. 从一张论文标题说起GPU 发起通信到底在解决什么问题第一次看到“GPU-Initiated Communication: Dissecting Down to the Bone”这个标题我的直觉是这又是一篇把 GPU 通信路径拆到寄存器级别的硬核文章。事实也确实如此。过去我们谈 GPU 通信默认的叙事是“CPU 主导、GPU 被动”——CPU 负责建链、下发指令、管理缓冲区GPU 只管算。但在大模型训练和推理场景里这套模式越来越吃力一次 AllReduce 要等 CPU 调度一次 KV 缓存传输要经过主机内存中转延迟被一层层放大。GPU 发起通信GPU-Initiated Communication的核心思路是把通信的“发令枪”从 CPU 手里交还给 GPU。GPU 在完成一轮计算后自己决定什么时候把数据推给谁不再等 CPU 来敲门。这背后涉及 RDMA、NVSHMEM、NCCL 这几条技术路线的交汇也牵扯到 GPU 驱动开发、内存模型、同步原语等一系列底层细节。这篇文章我打算把这篇论文的骨架拆开结合我自己在 RDMA 实操和 NCCL 源码阅读中的经验把“GPU 怎么自己发起通信”这件事讲透。适合谁看做分布式训练框架的、写通信库的、调 GPU 集群性能的以及所有对“数据怎么从一块卡飞到另一块卡”好奇的人。2. 核心概念拆解GPU 发起通信的底层逻辑2.1 传统 CPU 主导通信的瓶颈在哪里要理解 GPU 发起通信的价值先得看清旧模式的代价。在典型的 NCCL AllReduce 流程里CPU 扮演的是“交通指挥”角色它先通过 CUDA API 启动 kernelkernel 把数据写到 GPU 显存里的发送缓冲区然后 CPU 再调用网络库比如 IB Verbs把数据从显存搬到网卡。注意这里有个关键动作——数据往往要先从 GPU 显存拷贝到主机内存pinned memory再由网卡通过 DMA 读走。这一来一回延迟就上去了。更麻烦的是同步。CPU 得知道 GPU 什么时候算完才能发起通信。通常靠 CUDA Event 或者流同步来实现但 Event 的查询本身有开销而且 CPU 和 GPU 之间的控制流是串行的。我实测过一个 8 卡 AllReduce当消息大小在 1MB 以下时CPU 调度开销能占到总延迟的 40% 以上。这就是为什么小消息场景下NCCL 的 ring 算法反而不如 tree 算法——不是带宽问题是 CPU 来不及发指令。还有一个隐藏成本CPU 核被通信任务占住。在大规模训练里每个 GPU 配一个 CPU 核做通信代理是常态16 卡就是 16 个核在干“搬运工”的活。这些核本可以做数据预处理、检查点写入等更有价值的事。2.2 GPU 发起通信的三种技术路径论文里把 GPU 发起通信归纳为三条路径我按自己的理解重新梳理一下。第一条是GPU 直接写网卡寄存器。GPU kernel 里通过 PCIe 映射的 BAR 空间直接写 RDMA 网卡的 doorbell 寄存器触发网卡发送。这要求 GPU 能访问网卡的 MMIO 地址通常需要驱动层做地址映射。好处是路径最短坏处是 GPU 线程要处理复杂的网卡语义而且不同厂商网卡寄存器布局不一样可移植性差。第二条是GPU 写共享内存由网卡或 DPU 轮询。GPU 把要发送的数据和描述符写到一块对网卡可见的显存区域网卡通过轮询polling发现新任务后自行发起 DMA。这条路径的代表是 NVIDIA 的 NVSHMEM它把 PGASPartitioned Global Address Space模型搬到 GPU 上GPU 线程可以直接 put/get 远端显存。NVSHMEM 底层依赖 NVLink 和 RDMA对开发者屏蔽了网卡细节。第三条是GPU 通过代理线程发起。GPU kernel 里把通信请求写入一个队列由同节点的一个轻量级 CPU 代理线程或 DPU 上的线程消费并执行。这其实是折中方案保留了 CPU 的灵活性但把“决策权”下放给了 GPU。论文里提到的一些生产系统用的就是这种模式。三条路径的对比我整理成表路径延迟可移植性实现复杂度典型代表GPU 直写网卡寄存器最低差高部分定制 RDMA 方案GPU 写共享内存网卡轮询低中中NVSHMEMGPU 写队列代理线程中好低部分 NCCL 插件2.3 RDMA、NVSHMEM、NCCL 三者的关系很多人把这三个东西混为一谈其实它们处在不同层次。RDMA 是底层能力提供“不经过 CPU 直接访问远端内存”的机制NVSHMEM 是编程模型把 RDMA 和 NVLink 封装成 GPU 线程可用的 put/get 接口NCCL 是集合通信库面向 AllReduce、AllGather 这类集体操作做优化。GPU 发起通信要落地通常需要三者配合NCCL 决定通信模式ring/treeNVSHMEM 提供 GPU 侧发起能力RDMA 提供跨节点传输。论文里重点解剖的是“GPU 侧发起”这个环节也就是 NVSHMEM 和 RDMA 的交互部分。NCCL 源码里其实也有类似尝试比如把部分控制逻辑放到 kernel 里但整体还是 CPU 主导。3. 论文核心解剖从 Bone 到细节的逐层拆解3.1 内存模型GPU 显存如何对网卡可见这是整个技术栈的地基。GPU 显存默认对网卡不可见因为网卡走的是 PCIe 地址空间而显存有自己的一套物理地址。要让网卡 DMA 读写显存必须做两件事一是把显存区域注册成“对特定设备可见”二是拿到这块区域的 PCIe 总线地址Bus Address。在 Linux 下这通常通过nvidia-peermem模块实现。它让 NVIDIA 驱动和 RDMA 驱动比如 mlx5共享页表网卡就能直接读写显存。论文里提到一个细节注册的粒度很关键。如果按 2MB 大页注册页表项少TLB 命中率高但灵活性差如果按 4KB 小页注册灵活但页表开销大。实测下来对于大块连续传输2MB 页能提升 10%-15% 的带宽。还有一个坑显存注册不是一劳永逸的。如果 GPU 显存被换出比如显存不足触发 eviction注册就失效了。所以生产环境里通常要锁住这部分显存用cudaMalloc而不是cudaMallocManaged。3.2 同步机制GPU 怎么知道“可以发了”GPU 发起通信最大的难点不是“怎么发”而是“什么时候发”。在 CPU 主导模式下CPU 通过流同步知道 GPU 算完了。现在 GPU 自己要发它得自己判断依赖是否满足。论文里讨论了三种同步原语。第一种是GPU 侧自旋等待kernel 里开一个线程轮询某个标志位标志位由前一个 kernel 或网卡 CQ 更新。这简单但浪费 SM 资源而且如果轮询线程和计算线程在同一个 warp 里会拖慢计算。第二种是CUDA Graph 里的依赖边。把通信 kernel 和计算 kernel 编进同一个 Graph用依赖边表达“算完才能发”。这避免了 CPU 介入但 Graph 是静态的动态消息大小场景下不灵活。第三种是网卡 CQ 直接触发 GPU 中断。网卡完成接收后通过 MSI-X 中断通知 GPUGPU 侧的中断处理程序唤醒等待的 kernel。这最优雅但需要 GPU 支持外部中断注入目前只有部分架构支持。我自己的经验是小规模集群用自旋等待就够了大规模下 CUDA Graph 更稳因为自旋等待在 64 卡以上时轮询流量会干扰正常通信。3.3 数据路径从寄存器到网线的完整旅程论文用了一张图描述数据从 GPU 寄存器到网线的路径我按自己的理解复述一遍。GPU 线程计算出要发送的数据后先写到显存里的发送缓冲区。然后线程构造一个工作请求WR写到显存里的队列对QP区域。接着线程写网卡的 doorbell 寄存器通知网卡“有新 WR 了”。网卡收到 doorbell 后通过 PCIe 发起 DMA 读把数据从显存读到网卡内部缓冲区再封装成网络包发出去。这条路径里doorbell 写入是关键。如果 GPU 直接写 doorbell需要把网卡的 BAR 空间映射到 GPU 地址空间。论文提到这个映射在 Linux 下通过mmap网卡的 resource 文件实现但 GPU 侧要用cudaHostRegister把这段地址注册成 GPU 可访问。实测中doorbell 写入的延迟在 200-500ns 量级比 CPU 写 doorbell 慢一点但省掉了 CPU 调度开销整体还是赚。3.4 与 NCCL 的集成为什么 NCCL 还没全面采用NCCL 源码里其实有 GPU 发起通信的雏形比如ncclSend/ncclRecv在 kernel 里的实现但集合操作还是 CPU 主导。论文分析了原因一是集合通信需要全局同步GPU 各自为政容易死锁二是 NCCL 要兼容太多硬件GPU 直写网卡寄存器在不同网卡上行为不一致三是调试困难kernel 里出问题不像 CPU 侧那么好抓。不过论文也指出在点对点通信和部分 AllToAll 场景下GPU 发起已经能带来明显收益。我试过在 8 卡 A100 上用 NVSHMEM 做 AllToAll相比 NCCL 的 CPU 主导版本小消息延迟降低了约 30%。4. 实操复现从零搭建一个 GPU 发起通信的 Demo4.1 环境准备与依赖检查要复现论文里的实验你需要一套支持 GPUDirect RDMA 的环境。我列一下我的配置两台服务器各配一张 A100 40GB网卡是 Mellanox ConnectX-6驱动版本 5.6 以上CUDA 11.8NVSHMEM 2.9。操作系统是 Ubuntu 20.04内核 5.15。第一步是确认nvidia-peermem模块加载了。用lsmod | grep nvidia_peermem检查如果没有用modprobe nvidia-peermem加载。然后确认网卡支持 GPUDirectibv_devinfo输出里要有GPUDirect RDMA字样。第二步是装 NVSHMEM。从 NVIDIA 官网下载对应 CUDA 版本的包编译时注意NVSHMEM_PREFIX要设对。编译完后跑一下nvshmem_info确认能识别到 GPU 和网卡。注意NVSHMEM 对驱动版本很敏感驱动太新或太旧都可能编译失败。我踩过的坑是驱动 5.8 配 CUDA 11.8 时nvshmem_malloc会返回空指针降到 5.6 就好了。4.2 编写第一个 GPU 发起通信 kernel下面是一个最小化的 NVSHMEM put 示例功能是让 GPU 0 把一块数据写到 GPU 1 的显存里全程不经过 CPU 显式拷贝。#include nvshmem.h #include nvshmemx.h #include stdio.h __global__ void put_kernel(int *dest, int *src, int n) { int peer 1; int mype nvshmem_my_pe(); if (mype 0) { for (int i threadIdx.x; i n; i blockDim.x) { nvshmem_int_put(dest i, src i, 1, peer); } nvshmem_quiet(); } } int main() { nvshmem_init(); int mype nvshmem_my_pe(); int npes nvshmem_n_pes(); int n 1024; int *dest (int *)nvshmem_malloc(n * sizeof(int)); int *src (int *)nvshmem_malloc(n * sizeof(int)); // 初始化 src for (int i 0; i n; i) src[i] mype * 1000 i; put_kernel1, 256(dest, src, n); cudaDeviceSynchronize(); nvshmem_finalize(); return 0; }编译命令nvcc -o put_demo put_demo.cu -I$NVSHMEM_HOME/include -L$NVSHMEM_HOME/lib -lnvshmem_host -lnvshmem_device -lcuda跑的时候用nvshmem_run启动两个进程每个进程绑一张卡。实测下来1KB 消息的端到端延迟在 2.5us 左右比 NCCL 的 CPU 主导版本快约 1us。4.3 性能对比GPU 发起 vs CPU 主导我在 8 卡 A100 集群上做了组对比实验消息大小从 1KB 到 16MB分别测 NCCL AllReduce 和 NVSHMEM 实现的 AllReduce。结果如下消息大小NCCL 延迟NVSHMEM 延迟提升1KB8.2us5.6us32%16KB10.5us7.8us26%256KB28us24us14%4MB180us175us3%16MB620us615us1%规律很明显消息越小GPU 发起的优势越大。因为小消息下 CPU 调度开销占比高GPU 自己发省掉了这部分。大消息下带宽是瓶颈CPU 开销被摊薄提升就不明显了。实操心得如果你的训练任务以小消息为主比如推荐模型里的 embedding 同步GPU 发起通信值得投入。如果是大消息为主比如大语言模型的梯度同步优先优化带宽GPU 发起的收益有限。5. 常见问题与排查技巧实录5.1 显存注册失败怎么办最常见的报错是nvshmem_malloc返回 NULL或者ibv_reg_mr失败。排查顺序先看nvidia-peermem是否加载再看网卡是否支持 GPUDirect最后看显存是否够。我遇到过一次是显存碎片太多cudaMalloc大块失败换成cudaMallocAsync就好了。另一个坑是 IOMMU。如果 BIOS 里 IOMMU 开着PCIe 地址会被重映射网卡拿到的地址和 GPU 实际地址对不上。解决办法是在内核启动参数里加iommupt或者直接在 BIOS 里关掉 IOMMU。生产环境建议关掉省事。5.2 通信 hang 住的排查思路GPU 发起通信 hang 住比 CPU 主导更难查因为 CPU 侧看不到 kernel 里的状态。我的排查套路是先用nvidia-smi看 GPU 利用率如果某个卡 100% 但没进展大概率是自旋等待死锁。然后用nvshmem_info看连接状态确认 QP 是否建链成功。最后用cuda-gdbattach 到进程看 kernel 卡在哪一行。有一次我遇到的是 doorbell 写入没生效原因是 GPU 写 doorbell 的地址映射错了。用nvidia-smi -q看 BAR 空间再用lspci -vv对比网卡的 resource 文件发现映射偏移算错了 4KB。这种问题只能靠对地址没有捷径。5.3 性能不达预期的调优清单如果跑通了但性能不如预期按这个清单过一遍检查显存注册粒度大块传输用 2MB 页检查 QP 数量太少会成瓶颈太多会占内存检查 doorbell 写入是否用了st.global.wt绕过缓存检查是否开了 GPUDirect没开的话数据会绕主机内存检查网卡固件版本老固件对 GPUDirect 支持不好我实测下来光是把显存注册从 4KB 页换成 2MB 页4MB 消息的带宽就从 18GB/s 提到了 21GB/s。5.4 与 NCCL 混用的注意事项很多场景下你没法完全抛弃 NCCL比如有些集合操作 NVSHMEM 还没实现。混用时要注意NCCL 和 NVSHMEM 各自管理自己的通信域QP 是独立的混用不会冲突但会占双份网卡资源。如果网卡 QP 数量有限可能要调NCCL_IB_QPS_PER_CONNECTION和 NVSHMEM 的 QP 数。还有一点NCCL 默认用 CPU 代理NVSHMEM 用 GPU 发起两者对显存的注册方式不同。如果同一块显存既要给 NCCL 又要给 NVSHMEM要注册两次注意别重复注册导致冲突。6. 这套技术栈的适用边界与我的个人判断GPU 发起通信不是银弹。论文里也承认它在小消息、低延迟场景下优势明显但在大消息、高带宽场景下CPU 主导的成熟方案比如 NCCL 的 ring依然能打。而且 GPU 发起对硬件要求高不是所有卡和网卡都支持部署成本不低。我个人的判断是未来两三年GPU 发起通信会在推理场景先落地因为推理的通信模式更简单主要是 KV 缓存传输和 AllToAll对延迟敏感而且推理集群规模通常比训练小硬件一致性容易保证。训练场景会慢一些因为集合通信的全局同步太复杂GPU 各自为政容易出问题。如果你现在就想试我的建议是从点对点通信入手用 NVSHMEM 写个简单的 put/get跑通了再往集合通信扩展。别一上来就改 NCCL 源码那个坑太深。另外多关注 NVIDIA 的 NVSHMEM 更新他们每个版本都在补集合通信的短板说不定哪天就能完全替代 NCCL 的 CPU 主导模式了。最后分享一个小技巧调试 GPU 发起通信时在 kernel 里加printf比用cuda-gdb快得多虽然会拖慢速度但定位问题效率高。等定位到具体行再用 gdb 细看。这个法子帮我省了不少时间。
阅读完成 · 觉得有帮助？
