上周我们公司把最近对大语言模型中点对点通信的一些成果总结了一下，写了一篇论文挂在了 arXiv 上面，同时也在 GitHub 上面开源了代码。

我们构建了一套基于“无序可靠数据报传输（Unordered Reliable Datagram）”语义的 RDMA 通信库，既能跑在 AWS EFA，又能跑在 NVIDIA ConnectX。我们把这套 RDMA 通信库应用在了三个场景下面：分离式推理的 KvCache 传输、强化学习后训练中的模型参数更新、以及 MoE 通信。这个 MoE kernel 在 ConnectX-7 上面跑 decode 甚至比 DeepEP 还快一点点，在 EFA 上也是首次达到了可用的性能。

这篇文章我跟大家讲讲其中的来龙去脉，更多的是想分享一下做这些工作的动机以及背后的故事。对具体的技术细节感兴趣的读者可以在文章末尾的链接里面找到我们的论文、代码、以及相关的博客。

集合通信的痛点

首先，我们觉得集合通信虽然说在结构化并行的场景下久经考验（例如数据并行和张量并行），但是在一些新的应用场景用起来会很变扭，不如直接用 RDMA 点对点来得方便和高效。为什么说变扭呢？

第一，集合通信需要构建一个通信“世界”，而且这是一个成员固定的世界，没法增加或者删除成员，这就会给生产部署带来非常大的难题。举个例子，在分离式推理架构里面，Prefiller 和 Decoder 之间需要传输 KvCache。而我们知道，生产环境的线上流量总是会有涨落的，所以总是需要根据负载弹性地增加或者减少副本的数量。另一方面，无论是硬件故障还是软件有瑕疵，总是会有需要把某个坏掉的节点去掉的时候。而如果通信的成员只能是固定的，那既没有办法弹性伸缩，也没办法替换坏节点。

第二，通信世界的建立是阻塞的，而且需要所有成员的参与。这就意味着每次增减节点都需要让全世界停下来。虽然能用一些比较复杂的工程代码来异步地建立新的通信世界，但是这就很麻烦，而且也很容易出错。

第三，集合通信的接口保证了一个顺序执行的语义。其实这个本意是很好的，可以简化开发者的负担。然而，网络可以是乱序到达的，所以有可能就需要做一些额外的缓冲或者同步来保证顺序执行的语义。可惜的是，集合通信库辛辛苦苦为大家做了一个这么有序的抽象，有的应用场景却不领情。比方说还是 KvCache 的传输，我们只是想把所有的页面都传出去，至于哪个页面先到哪个后到则是不重要的。另外一个例子就是强化学习后训练当中训练节点往推理节点更新模型参数，我们在乎的是所有的模型参数都更新好，而不在乎哪个先哪个后。

第四，集合通信的接口要求所有参与者都有相同的张量形状和类型，这样就会带来一些编程上的不便，有些时候还会带来很严重的性能损失。比方说如果用集合通信来做 RPC，那么每次就都得传最大的消息大小。再比方说 MoE，往每个 Rank 要发多少个 token 是没法提前知道的，所以如果用集合通信来做 all-to-all 的话，那么就得按照最坏情况来算 token 数量（也就是所有 rank 都发往同一个 rank）。这种最坏情况跟平均情况可以差非常多，如果用 DP64 EP64 那就是每个 rank 要多传64倍数据，而且还是无用的数据。

RDMA 的痛点

那为什么大家还是执着于集合通信呢？我感觉一个原因是缺少简单易用的 RDMA 通信库，另一个原因是现有的 RDMA 库基本上都绑定英伟达网卡 ConnectX。众所周知，“N卡网速快”，这句话本来只是个2010年代的梗，结果在英伟达收购了 Mellanox 之后，这句话就变成了现实。我们在论文里面对比了，是真的快。而且 ConnectX 网卡就是直接可以在市场上面购买到的商品。又快，又好买，又有大量的软件支持，也不难怪大家都在 ConnectX 的生态上继续开发。

但是对于技术实力强劲的云服务商来说，他们投入资源开发自己的网卡也是对的。最开始的目的应该还是为了解决云服务器的存储和网络的虚拟化、安全性、减少网络通信的 CPU 占用。后面高性能计算以及机器学习的应用需求来了，再缝缝补补接着往上尽量补全功能，尽量能兼容已有的生态。比如 AWS 就有专门的 NCCL 插件，能让 NCCL 跑在 AWS EFA 上面。

能跑 NCCL 并不代表就能跑其他的通信库。比方说 Mooncake、DeepEP 这样非常优秀的通信库，都没有办法在 AWS EFA 上面跑。我们3月份的时候也尝试过依托 NVSHMEM 来写一个 MoE 通信库，但是效果非常不好。在 ConnectX-7 上面，哪怕是用了 NVSHMEM 的 IBGDA 实现，速度也远比直接用 mlx5 驱动的 DeepEP 要慢。虽然 NVSHMEM 也支持 EFA，但是我们当时测下来发现要比 ConnectX-7 慢40倍，比 NCCL 还慢很多，完全不可用。

为什么 RDMA 程序难以移植到不同的网卡呢？我觉得有下面几个原因：

第一，程序的逻辑会依赖于底层通信协议的特性。大家用得最多的 RDMA 通信协议是可靠连接（Reliable Connection, RC）。RC 保证了消息是按顺序到达的。比方说如果我有一堆的消息要发，要怎么让远端知道这些消息都送到了呢？如果依赖于消息按序送达的假设，那就可以在最后一个消息上面打上一个记号，这样远端一旦收到了这个记号，它就知道了前面所有的消息都送达了。这里“打记号”的方式可以有很多种，比方说可以在最后一个消息上面带上一个32位的立即数（Immediate），再比方说对一个计数器进行原子操作。然而 EFA 用的是自研的可扩展可靠数据报协议（Scalable Reliable Datagram, SRD），虽然它保证消息能送达，但是送达的顺序是乱序的。

这个时候我们想象一下像 NVSHMEM 这样的通信库要怎么实现可移植性，让一个基于按序送达假设的程序能在一个无序送达的硬件上跑起来。显然什么都不做是会破坏正确性的。所以为了兼容这个语义，通信库就不得不多做很多额外的工作，比如可以使用缓冲区来进行消息重排，比如可以插入同步点来等待消息队列清空。显然这些权宜之计都会严重地损失性能。

第二，网卡本身的带宽也不同。ConnectX-7 网卡，一张网卡就能达到 400 Gbps 的带宽。然而 AWS p5 实例上，一张 EFA 网卡就只有 100 Gbps 带宽。p5en 实例好一点，一张 EFA 网卡有 200 Gbps 带宽。总之，需要把多张网卡聚合起来才能达到跟 ConnectX-7 一样的带宽。

第三，网卡支持的功能也不同。之前提到的 IBGDA 技术能让显卡直接向网卡发起通信，然而到目前为止这个功能只有 ConnectX 网卡支持。在缺少这个功能的情况下，又要让显卡能往网卡发起消息，那就只能让 CPU 来充当中间商了。如果这个软件实现做得比较差，那就会带来很多额外的开销。后文我还会展开讨论这个点。

第四，程序为了追求极致的性能，用的接口可能本身就是不可移植的。比方说 DeepEP 虽然说表面上用了 NVSHMEM，但其实只是把 NVSHMEM 当作一个启动器，真正跟网卡打交道的是直接用 mlx5dv，这就直接绑定了 ConnectX 网卡。

解决痛点的思路

基于上面这些背景，我解决上面这些问题的思路是取不同通信协议和硬件特性的交集，在此之上构建一套尽可能通用的 RDMA 通信库：

  • 可靠传输：网卡和网络层需要保证消息能送达。一来是省得在通信库里面实现重传等麻烦的逻辑，二来是充分利用硬件卸载的优势，尽量减少 CPU 的参与，避免 CPU 拖慢性能。
  • 无序送达：前面提到，如果硬要在无序送达的硬件上实现按序送达的语义，那就会损失性能。另一方面，应用场景也不一定需要按序送达的语义。所以不如直接放弃按序送达的语义。
  • 支持双边 RDMA 操作：RECV / SEND 非常适合用来实现 RPC。
  • 支持单边 WRITE_IMM：单边 WRITE 可以提供极致的性能，而带上一个32位的立即数又可以用来通知目标方操作已经完成。
  • 基于 ImmCounter 的同步机制：虽然应用场景不一定需要按序送达，但是还是需要一些同步机制来保证所有的消息都送达了。要在无序送达的硬件上实现这个语义，我们构建了一个 ImmCounter 抽象，即每一个 WRITE 操作都带上一个用户指定的立即数，然后目标方的通信库会对这个立即数进行计数，当计数器达到预期的值时，就认为所有的消息都送达了。
  • 针对常见的通信模式提供高性能的实现优化：比方说同样是 WRITE，除了传输一个连续的内存区域，还可以是多个不连续的内存区域（例如 KvCache 的传输），还可以是多个接收方并且每个接收方传输的数据量不同（例如 MoE 中的 scatter）。虽然理论上来说这些操作都可以在应用层通过单个 WRITE 操作来实现，但是在网络库提供这些额外的接口还是有很多好处。一是可以简化应用层的实现；二是可以减少通信库和应用程序之间跨语言及跨线程的通信开销；三是可以在内部缓存一些数据结构来优化性能。
  • 支持多网卡聚合：网络库可以提供一些默认的聚合策略，例如对于单个大消息，可以对消息字节数进行切分；对于多个页面的传输，可以对页面的列表进行切分；对于多个接收方的传输，可以对接收方的列表进行切分。当然，如果应用程序想要特殊的控制，也可以手动指定网卡聚合策略。
  • 借助 CPU 来支持显卡发起 RDMA 通信（Host-Proxy 模式）：KvCache 传输和 MoE 都需要在 Cuda Graph 里面发起 RDMA 操作，那我们就往内存和显存里面放一些特殊标记。因为 CPU 和显卡各自都可以访问自己和对方的内存，所以两者都可以通过写入和轮询这些标记来传递消息。

注意到，这里我舍弃掉了对已有 RDMA 程序的兼容性。其一是因为对网络底层假设发生了变化，为了保证性能，我们要求应用程序必须按照我们这套更松弛的假设来编写。其二是因为我和我同事们都喜欢造轮子，所以上层的应用迟早都会被我们造一遍，也没有去兼容已有程序的需求。

我造的这套通信库一开始只是为了 KvCache 传输一个用处，后面发现用在强化学习后训练的模型参数更新上面也非常方便。让我比较吃惊的是，这套通信库用在 MoE 上面效果也出奇的好。

MoE Kernel

MoE 的故事还得回到今年5月份去 MLSys 开会的时候，在跟字节的朋友们交流的时候得到了一个非常重要的洞见：CPU 和显卡之间的 PCIe 延迟只有2微秒。我跑了一下 GDRCopy 自带的性能测试，确实如此。所以开完会之后我就想开一个坑，用这套通信库实现类似 NVSHMEM 的功能，来弥补没有 IBGDA 的遗憾，毕竟一次 MoE 通信也得好几百微秒，就算这套 host-proxy 带来15微秒的额外开销，也完全能够接受。然后在考虑 EFA 的乱序送达特性的时候，我意识到了直接移植现有的 DeepEP 代码或者我们之前的 NVSHMEM 实现都可能带来性能损失，要达到最好的性能最好还是按照乱序送达的语义来重新设计一套 MoE kernel。正好我同事也对他上次用 NVSHMEM 实现的 MoE kernel 不太满意，所以这个艰巨的任务就交给他了。我的同事也是非常给力，单枪匹马很快就实现好了一个新的 MoE kernel。

看到新的 MoE kernel 能在 EFA 上面跑起来，我们都很兴奋。但是我们不知道它跑得慢到底是我们优化得不够，还是 EFA 本身的性能不好。所以我们就想着要在 ConnectX-7 上跑起来，然后跟 DeepEP 比一比。算上 CPU 和显卡之间的 PCIe 延迟，以及 CPU 上面的各种开销，我们觉得优化的终点应该就是比 DeepEP 慢个15微秒就差不多可以收工了。

要让我们的 MoE kernel 能在 ConnectX-7 上面跑起来，首先我们得有 ConnectX-7 机器，于是就到处化缘。在这里我要非常感谢英伟达，他们非常慷慨地给我们提供了一个64卡的开发集群。

有了开发机，我就开始往我们的 RDMA 通信库里面加 ConnectX-7 的支持。我本来以为 libfabric 本身就能支持 ConnectX-7，结果发现怎么都跑不起来，问了一下 libfabric 社区和英伟达，他们都说 ConnectX 没有作为开发的重点，所以我猜就算能跑性能也不好。于是我就用 libibverbs 自己写了一套。好在我之前造的这套 RDMA 通信库的抽象做得还行，所以只需要最底下硬件相关的代码换一下就好了；中间的抽象不需要做大的改动；暴露给应用的接口更是不需要改。

在加 verbs 支持的时候，也算是对比了一下 SRD 和 RC 两种协议的编程接口。虽然说 EFA 跑得慢，但从编程接口的角度来说，我更喜欢 SRD 一点：

  • SRD 是基于数据报的协议，知道地址就能直接发消息，很方便。RC 是基于连接的协议，所以跟每一个远端节点都需要建立一个连接。然而要建立 RC 连接，需要用另一个信道来交换彼此的 QPN 和 PSN，这个握手的过程就变成一个需要双方参与的异步操作。为了不引入额外的通信信道，所以我就用 UD 实现了这个握手过程。
  • 我有考虑过维护连接这个事情到底是在网络库里面做比较好还是让应用程序来做比较好。最后我觉得，还是让网络库把这个复杂度吃掉比较好。毕竟我在写网络库的时候都已经觉得维护连接是一个很麻烦的事情了，更不用说应用的开发者了。（结合上前面的乱序送达特性，我觉得可以把我们这个网络库对网络的假设称作 Relaxed Reliable Datagram，RRD。）
  • libfabric 接收 WRITE_IMM 的立即数是不需要消耗 RECV 操作的，然而 verbs 会消耗，而且 verbs 的 RECV 操作是按照提交的顺序消耗的。这就意味着没法把 WRITE_IMM 和正常的双边 RECV 操作混合在一起跑。好在我后面想了一个挺简单的解决方法，开两个 RC QP，把单边操作和双边操作分开跑。

在 ConnectX-7 上面跑通了我们 RDMA 通信库的点对点测试之后，我就开始胆战心惊地尝试跑 MoE kernel。本来以为肯定还要在折腾各种奇怪的问题的，结果没想到竟然第一次就丝滑地跑通了！这也是对我们这套 RDMA 通信库的可移植性的一个很好的证明。

跑通了之后同事和我就开始做优化。优化的维度有很多：调整算法增加并行度；调整 RDMA 和 NVLink 的同步；调整 Cuda kernel 里面的 PTX 指令、内存屏障；优化 load/store；减少 CPU 的内存分配次数；在 RDMA 通信库里面引入多个接收方的优化接口；ConnectX 和 EFA 通用的优化以及各自单独的优化；等等。具体的细节就不展开了，有兴趣的读者可以看我们的论文、其他几篇博客、以及具体的代码。

最终优化完了之后，没想到在 ConnectX-7 上面 Decode 竟然跑的比 DeepEP 还要快。为了确保结果的可比性，我还特地让 DeepEP 用跟我们相同的性能测试代码来跑，测出来的结果也跟 DeepEP 官方的性能测试结果差不多。另外一方面，虽然我们一直是在 ConnectX-7 上面做优化，但这些优化大多也同时作用于 EFA 上，所以最终我们 EFA 上面的性能也提升了不少。

在这里我也要非常感谢 AWS 尤其是 EFA 团队。他们回复我们问题的速度非常快，给了我们很多关于 EFA 的建议。我们之后也会接着和 AWS 合作，继续优化我们的通信库以及 MoE kernel 在 EFA 上的性能。

在对 Decode 结果满意之后，我们又试着跑了跑 Prefill。跟 DeepEP 不同，我们 Prefill 和 Decode 用的是同一个实现，不过我们还没专门针对 Prefill 做优化。测出来我们能打满 RDMA 带宽，不过还是比 DeepEP 慢了不少。最让我们觉得奇怪的是，我们用 RDMA 带宽以及传输的字节总数算了一下理论上的最低延迟，结果发现 DeepEP 的延迟比理论的更低。在仔细检查了性能测试程序，确认一些无误之后，唯一的解释就是 DeepEP 传的字节总数比我们的少。后面又请教了一下 DeepEP 的朋友，证实了这个猜测。在 dispatch 的时候，如果同一个 token 发给同一台机器的不同 rank，DeepEP 可以只用 RDMA 发一次，然后用 NVLink 在机内进行转发；在 combine 的时候，如果同一台机器的不同 rank 会发同一个 token 给同一个远程的 rank，DeepEP 可以先在 NVLink 上做一个部分和，然后在 RDMA 上就只需要发送一个 token 就行了。不得不说 DeepEP 真的是太精妙了！这也解释了为什么 DeepSeek-V3/R1 里面选专家的时候会限定最多分到4个节点上。

感慨万千

到这里，我们的优化工作也暂告一个段落，于是我们整理了一下思路，写了一篇论文，也把开源代码放出来了。感兴趣的读者可以点击下面的链接查看：

这一路上我从同行朋友们的交流以及开源项目中学到了太多东西。我自己写博客以及在公司里面推动技术博客和代码开源，也算是向社区做出的一点微小的贡献吧。

文章的最后我想感慨一下，在不到一年的时间里面，在 RDMA 这方面我真的学习到了很多。去年12月初的时候我还不懂怎么在 AWS 上跑 RDMA，趁着圣诞节放假摸索了怎么用 EFA，到现在有了一套比较完整的解决方案。我们小小的4人 Inference Team （这两个月加到了7个人）在一年半里面也做了很多事。去年4月份我加入的时候，我们还完全依赖 TensorRT-LLM，后面我们搞了自己的推理引擎、搞了上层的智能路由、搞了下层的算子和网络。一开始我只是把我已经知道的推理优化认真的做工程落地，做着做着发现我们有了一些工程上的创新，再到现在感觉有了学术研究贡献的创新，真的是收获的远远比我预想的多。

（如果大家对我为什么选择了我司以及对我司的工作环境感兴趣，我后面也可以专门讲讲。）

如果在美国的读者有对我司的系统优化感兴趣的，欢迎直接联系我。对我司后训练、AI 产品、或者其他职位感兴趣的，我也可以帮忙内推。