RDMA POINT-TO-POINT COMMUNICATION FOR LLM SYSTEMS

作者/机构: Nandor Licker, Kevin Hu, Vladimir Zaytsev, Lequn Chen

A1 主要贡献

本文旨在解决新兴大语言模型（LLM）系统中点对点通信的硬件可移植性问题。

  • 核心问题: 新兴的LLM系统模式，如解耦式推理（disaggregated inference）、专家混合（MoE）路由和异步强化学习微调，需要灵活的点对点通信，而不仅仅是简单的集体通信（collectives）。现有的LLM框架主要依赖NCCL等集体通信库，这些库存在成员固定、需要同步初始化和缓冲区大小统一等限制，不适用于动态和稀疏的通信模式。虽然高性能计算领域早已使用基于RDMA的灵活通信原语，但在LLM框架中，由于云服务商（如AWS EFA）和传统硬件（如NVIDIA ConnectX）提供的RDMA实现不同（例如，EFA SRD协议是无序交付，而ConnectX RC协议是有序交付），导致缺乏统一的抽象，现有解决方案存在供应商锁定问题。

  • 研究目标与创新点: 本文的目标是开发一个可移植的RDMA通信库，以弥合不同RDMA硬件之间的功能差距。

    • 核心洞察: 作者发现，无论是NVIDIA ConnectX还是AWS EFA，它们都支持可靠但无序的交付。ConnectX的RC传输可以忽略排序，而EFA的SRD协议本身就是无序的。
    • 提出TransferEngine: 基于上述洞察，作者提出了一个名为TransferEngine的可移植RDMA通信库。它提供了一个统一的接口，支持双边SEND/RECV和单边WRITEIMM操作。
    • 创新原语IMMCOUNTER: 为了在不依赖消息排序的情况下实现完成通知，TransferEngine引入了一种新颖的IMMCOUNTER原语。
    • 关键特性: TransferEngine能够透明地管理每个GPU的多个NIC（网络接口控制器），这对于EFA（需要聚合4个100 Gbps NIC以达到400 Gbps）至关重要，并避免了供应商锁定。

  • 主要贡献与应用展示: 作者通过三个生产系统展示了TransferEngine的有效性，证明了点对点通信能够补充集体通信，同时保持高性能和可移植性。

    1. KvCache传输 (第4节): 实现了解耦式推理，允许预填充（prefiller）和解码（decoder）节点之间进行无限制的通信，支持弹性扩展，无需同步初始化或固定成员。该系统已在EFA上经过生产测试，完全支持CUDA Graph，并实现了低延迟的逐层传输。
    2. 强化学习（RL）权重更新 (第5节): 采用点对点方法，对万亿参数模型实现了1.3秒的权重更新，比现有RL框架快100倍以上。通过从每个训练GPU直接向推理GPU发起单边RDMA WRITE，充分利用了整个集群的带宽，并通过流水线执行重叠了H2D memcpy、权重准备和RDMA传输。
    3. MoE调度/合并 (第6节): 在ConnectX-7上实现了业界领先的解码延迟，性能与专门的DeepEP内核相当。同时，它是在EFA上首个可行的实现，通过并行化令牌和路由传输来隐藏设备到主机的延迟和网络延迟。

A3 背景知识与相关工作 (缩写)

2.1 网络技术

  • RDMA: 远程直接内存访问（RDMA）是现代机器学习系统的核心，提供高吞吐、低延迟的网络通信。目前部署的NIC可提供400 Gbps带宽和亚微秒级延迟。RDMA通过分离控制平面（初始化设备和缓冲区，需操作系统参与）和数据平面（数据传输和完成轮询，绕过内核）来实现高性能。

  • RDMA操作: RDMA提供双边和单边操作。双边SEND/RECV操作需要接收方先发布一个带有注册内存区域的RECV，然后发送方发出SEND，并通过完成队列通知接收方。单边WRITE操作将本地内存复制到远程内存区域，无需对端参与，但需要远程内存地址和密钥（RKEY）。WRITEIMM是WRITE的扩展，它在传输数据的同时传递一个32位立即数，并同样通过完成队列通知接收方。READ和原子操作因延迟较高，不适用于本文的场景【12, Reda et al., RDMA is turing complete, we just did not know it yet!, 2022, NSDI】。

  • RDMA传输协议: RDMA规范定义了三种传输协议：可靠连接（RC）、不可靠连接（UC）和不可靠数据报（UD）。表1总结了它们的特性，并指出本文的TransferEngine是它们之间的共同基础。
    表1. RDMA传输类型比较

  • 云RDMA适配器: 除了广泛部署的ConnectX系列，主要的云提供商也在部署自己的解决方案，如AWS弹性光纤适配器（EFA）、阿里云eRDMA和Google Falcon。虽然大多数与RC兼容，但EFA实现了专有的SRD协议，该协议通过libfabric接口暴露【29, Shalev et al., A cloudoptimized transport protocol for elastic and scalable HPC, 2020, IEEE Micro】【25, OFIWG, libfabric: Open Fabrics Interfaces (OFI), 2014】。SRD是无连接的，提供可靠但无序的交付。

2.2 编程接口

  • 集体通信: 机器学习框架主要依赖集体通信库（如torch.distributed, NCCL, MPI）进行GPU间数据交换【15, Li et al., PyTorch Distributed: Experiences on accelerating data parallel training, 2020, Proc. VLDB Endow.】【23, NVIDIA, NVIDIA collective communication library (NCCL), 2015】【21, MPI Forum, MPI: A Message-Passing Interface Standard Version 5.0, 2025】。虽然集体操作擅长结构化数据传输，但其点对点能力有限：(1) 固定成员：要求预先知道所有参与者，阻碍了动态扩展；(2) 同步初始化：需要全局协调来形成“通信世界”，阻塞了独立的对等连接；(3) 操作排序：要求所有参与者就操作序列达成一致，即使NCCL支持并发接收也需要额外的同步【10, Hu et al., Demystifying NCCL: An in-depth analysis of GPU communication protocols and algorithms, 2025b】；(4) 形状统一：即使是点对点传输，也限制了所有参与者的传输大小。

  • GPUDirect: GPUDirect RDMA允许RDMA NIC通过PCIe直接访问GPU内存，消除了中间的CPU内存拷贝【22, NVIDIA, GPUDirect RDMA, 2012】。传输可以由主机发起，或者在支持GPUDirect Async（IBGDA）的情况下，由GPU自己发起以绕过PCIe开销【1, Agostini et al., GPUDirect async: Exploring GPU synchronous communication techniques for InfiniBand clusters, 2018, J. Parallel Distributed Comput.】。目前，GPUDirect Async仅在ConnectX NIC上受支持。此外，GPUDirect RDMA通过GDRCopy库实现了低延迟的主机-设备memcpy【32, Shi et al., Designing efficient small message transfer mechanism for inter-node MPI communication on InfiniBand GPU clusters, 2014, HiPC】。

2.3 相关工作

  • 解耦式推理: Splitwise【26, Patel et al., Splitwise: Efficient generative LLM inference using phase splitting, 2024, ISCA】, DistServe【44, Zhong et al., Distserve: Disaggregating prefill and decoding for goodput-optimized large language model serving, 2024, OSDI】, 和Mooncake【27, Qin et al., Mooncake: Trading more storage for less computation — a KVCache-centric architecture for serving LLM chatbot, 2025, FAST】是最早展示解耦式推理架构优势的工作，它们将LLM推理的预填充和解码阶段分离到不同设备上。本文的KvCache传输用例通过TransferEngine将预填充和解码集群通过RDMA连接起来，实现了这一范式。

  • 点对点网络库: NVSHMEM同时提供了集体操作和灵活的点对点通信【14, Langer et al., Dynamic symmetric heap allocation in NVSHMEM, 2021, OpenSHMEM】。它支持GPU发起的（IBGDA）和主机代理（IBRC）通信，但在EFA上性能严重下降。NVIDIA Inference Xfer Library（NIXL）针对LLM推理的P2P通信，基于UCX构建【24, NVIDIA, NIXL: NVIDIA inference xfer library, 2025】【30, Shamis et al., UCX: an open source framework for HPC network APIs and beyond, 2015, IEEE】。本文在生产中部署的EFA实现早于NIXL（v0.6.1，2025年10月）中初步的EFA支持。Mooncake也提供了一个RDMA传输引擎，但不支持EFA。其他库如UCCL【19, Mao et al., Previewing UCCL-EP: Flexible and efficient expert parallelism for cloud and beyond, 2025】和MSCCL++【28, Shah et al., Msccl++: Rethinking gpu communication abstractions for cutting-edge ai applications, 2025】则专注于对集体操作进行网络层优化，而非点对点通信。

  • 分布式KvCache存储: Mooncake Store和DeepSeek 3FS【5, DeepSeek AI, Fire-flyer file system (3fs), 2025】提供了KV Cache的分布式存储，但目前缺乏EFA支持。本文通过提供适用于云部署的可移植RDMA原语来补充这些系统。

  • 计算与通信重叠: 关于LLM内核中计算与集体通信重叠的研究【4, Chang et al., Flux: Fast software-based communication overlap on gpus through kernel fusion, 2024】【41, Zhang et al., COMET: Fine-grained computationcommunication overlapping for mixture-of-experts, 2025】【42, Zheng et al., Sglang: Efficient execution of structured language model programs, 2024】【43, Zheng et al., Tritondistributed: Programming overlapping kernels on distributed ai systems with the triton compiler, 2025a】与本文关注的RDMA原语是正交的，尽管本文也支持后台传输。

  • LLM框架: 本文的TransferEngine和MoE内核可以集成到vLLM, SGLang, TensorRT-LLM, FlashInfer等LLM推理框架和内核库中【13, Kwon et al., Efficient memory management for large language model serving with pagedattention, 2023, SOSP】【42, Zheng et al., Sglang: Efficient execution of structured language model programs, 2024】【23, NVIDIA, TensorRT-LLM, 2023】【40, Ye et al., FlashInfer: Efficient and customizable attention engine for LLM inference serving, 2025】。本文的P2P权重更新方法可以被Slime, OpenRLHF, AReaL, veRL, LlamaRL, NVIDIA Nemo等强化学习框架采用【36, Wu et al., LlamaRL: A distributed asynchronous reinforcement learning framework for efficient large-scale LLM training, 2025】【45, Zhu et al., slime: An llm post-training framework for rl scaling, 2025】【7, Fu et al., Areal: A large-scale asynchronous reinforcement learning system for language reasoning, 2025】【20, Mei et al., Real: Efficient rlhf training of large language models with parameter reallocation, 2025】【31, Sheng et al., Hybridflow: A flexible and efficient rlhf framework, 2024】【9, Hu et al., Openrlhf: An easyto-use, scalable and high-performance rlhf framework, 2025a】。

A2 方法细节 (缩写)

3 TRANSFERENGINE

TransferEngine是一个基础库，它通过抽象异构硬件，在一个简单的协议下实现高效的基于RDMA的点对点通信。它暴露了SEND/RECV操作来实现类似RPC的接口。对于KvCache传输，它提供分页WRITEs用于批量写入。对于RL权重传输，它暴露了低延迟高吞吐的WRITE操作。对于MoE路由，它特化了针对多个对端的WRITEs，以实现低延迟的SCATTER和BARRIER操作。所有操作之间没有任何排序保证。一个立即数可以与WRITEs关联，在接收方收到数据后递增一个计数器。

图1. TransferEngine管理跨NUMA节点的GPU，每个GPU带多个NIC。命令被转发给工作线程，工作线程响应回调处理器或IMMCOUNTER。
图1. TransferEngine管理跨NUMA节点的GPU，每个GPU带多个NIC。命令被转发给工作线程，工作线程响应回调处理器或IMMCOUNTER。
3.1 概述和设计目标

  • 核心设计：TransferEngine暴露了一个最小化的API，抽象了底层RDMA接口的复杂性。它支持多种接口，包括使用SRD协议的EFA和可通过libibverbs编程的多种NIC（包括ConnectX-7）。它能透明地检测拓扑以处理每个GPU的多个NIC：单个ConnectX-7 NIC提供400 Gbps带宽，而在AWS p5实例上实现同等带宽需要聚合四个100 Gbps的EFA NIC。单个实例会跨多个线程来管理一个节点内所有的GPU和NIC。通过零拷贝接口和硬件特定优化实现低延迟操作。

  • 解决排序差异：为了弥合RDMA RC的有序保证和EFA SRD的无序交付之间的差距，TransferEngine完全依赖一种新颖的IMMCOUNTER机制。它不依赖排序保证，所有完成通知都通过轮询完成队列并在批量接收到立即数值后传递通知来处理。通知通过回调函数或原子标志传递。

3.2 架构

  • 架构图示：图1展示了TransferEngine的架构。该引擎为每个GPU派生一个工作线程（worker），每个工作线程管理一个通用的DOMAINGROUP，该DOMAINGROUP协调所有关联的RDMA NIC（通常为1-4个，取决于硬件平台）。在一个通用的DOMAINGROUP内，每个DOMAIN都针对特定硬件进行了特化，负责单个NIC，处理队列对管理、工作提交和完成轮询。

  • 地址交换与传输分片：每个TransferEngine实例暴露一个主地址用于身份识别和远程对端发现。我们使用一个NetAddr结构体来捕获和序列化一个DOMAIN的网络地址，并在希望通信的对端之间交换该结构。一个限制是，所有对端每个GPU必须使用相同数量的NIC。因此，任何传输都完全了解源和目标域之间的NIC情况，这使得TransferEngine可以对请求进行分片或均衡。这对于E-FA至关重要，因为它使用多个适配器来实现完整的400Gbps带宽。

3.3 API 设计

TransferEngine暴露的API（如图2所示）实现了对RDMA的抽象：

#[serde] struct NetAddr(Bytes);
#[serde] struct MrDesc{ ptr: u64, rkeys: Vec<(NetAddr, u64)> }
struct MrHandle(NonNull<c_void>);
type Offset = u64;
struct Pages{ indices: Vec<u32>, stride: u64, offset: Offset }
struct PeerGroupHandle(u64);
struct ScatterDst{ len: u64, src: Offset, dst: (MrDesc,Offset)}
enum OnDone { Callback(fn () -> ()), Flag(Atomic<bool>) }

trait TransferEngine {
    fn main_address() -> NetAddr;
    // 内存区域管理
    fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc);
    // 双边 Send/Recv
    fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ());
    fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ());
    // 单边 Write
    fn expect_imm_count(imm: u32, count: u32, cb: fn () -> ());
    fn submit_single_write(len: u64, imm: Option<u32>,
                           src: (MrHandle, Offset),
                           dst: (MrDesc, Offset), OnDone);
    fn submit_paged_writes(page_len: u64, imm: Option<u32>,
                           src: (MrHandle, Pages),
                           dst: (MrDesc, Pages), OnDone);
    // 对一组对端的单边 Write
    fn add_peer_group(addrs: Vec<NetAddr>) -> PeerGroupHandle;
    fn submit_scatter(h: Option<PeerGroupHandle>, OnDone,
                      imm: Option<u32>, src: MrHandle,
                      dst: Vec<ScatterDst>);
    fn submit_barrier(h: Option<PeerGroupHandle>, OnDone,
                      imm: u32, dst: Vec<MrDesc>);
    // 用于 CPU-GPU 同步的观察者
    fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>;
}

  • 内存注册: 内存区域必须在引擎中注册，返回一个可序列化的MrDesc（可与对端交换以提交WRITEs）和一个本地MrHandle（用作传输源）。该API可以注册主机端缓冲区和GPU内存。在这些不透明类型背后，句柄携带了它们关联的所有NIC的地址，以及附加的特定于域的远程密钥列表，形式为(NetAddr, RKEY)对。

  • 点对点传输: submit_send及其远程对应的submit_recvs封装了SEND/RECV操作，以暴露RPC风格的通信，用于交换小载荷。提交时会进行一次拷贝，以便调用者可以立即重用或释放缓冲区。接收操作会发布一个循环的缓冲区池来接收数据。每收到一条消息，一个缓冲区会暂时从池中取出，让回调函数无需拷贝即可处理。回调完成后，它会自动被重新发布。必须分配足够的缓冲区以避免拒绝消息。这些操作只使用域组中的第一个NIC。

  • 单边写操作: submit_single_write和submit_paged_writes将数据从源缓冲区传输到目标缓冲区，分别写入由间接索引、步长和偏移决定的连续或分页区域。引擎将这些操作转化为可能多个零拷贝的单边RDMA WRITE操作，并在可用的NIC上进行分片和轮转。每次传输可以携带一个可选的32位立即数，以便在完成时通知接收方。传输完成通知会异步地传递给调用者。

  • 组播与屏障: submit_scatter将源缓冲区的一个切片发送到对等组中的每个对端，并写入它们接收缓冲区的不同偏移处；而submit_barrier是一个仅含立即数的操作，用于对端通知。这些是围绕WRITE的优化封装，允许应用程序预先注册一个PeerGroup以进行低延迟的批量传输。传输在两个设备之间执行，用户需要协调系统中多个设备的操作。

  • UVM观察者: 为了让主机能够在GPU取得进展时发起传输，alloc_uvm_watcher注册一个回调函数，当内存中某个字发生变化时调用。它分配一个统一虚拟内存（UVM）位置，设备（包括CUDA图中的内核）可以更新该位置，然后一个CPU线程使用GDRCopy持续轮询它。由于不保证所有变化都能立即被观察到，回调函数被调用时会传入旧值和新值，使其能够响应GPU端的进展。

  • 完成通知: 通知会为发送确认和接收完成两种情况传递，通过回调函数或原子标志。IMMCOUNTER是一个专用组件，负责跟踪每个立即数对应的计数器，这些计数器在从底层设备完成队列中检索到的事件上递增。事件在发送方传输完成后，以及在接收方完全交付带有立即数的载荷后生成，保证了原子性。计数器分配在与域工作线程相同的NUMA节点上。计数器可以透明地与GPU通过GDRCopy同步，可以直接通过轮询观察，也可以通过expect_imm_count注册后，在引擎内一个专用的独立线程上交给回调函数处理。所有同步都使用此类计数器实现，因为没有其他排序保证。

3.4 实现

  • 实现细节: 我们的TransferEngine用Rust实现，精心优化了内存分配、线程和同步以最小化延迟，从而实现高吞吐量。引擎为每个DOMAINGROUP派生一个工作线程，每个线程都固定到设备所连接的NUMA节点上的一个CPU核心，以最小化调度和内存访问延迟。特定于域的数据结构在核心固定后分配，以确保内存在正确的NUMA节点上保留。一个工作线程负责处理多达4个DOMAIN，每个DOMAIN管理一个NIC，而另一个专用线程负责轮询GPU以更新UVM观察者。跨线程通信通过无锁队列完成。

  • 工作循环与分片: API将请求转发给为相应设备服务的DOMAINGROUP，其中第一个也服务于主机。在一个紧凑的循环中，域工作线程轮询新工作，优先提交新请求。根据硬件和配置，请求会在可用的DOMAIN之间进行分片和负载均衡。复合请求的第一个WRITE会立即发布到DOMAIN中NIC的发送队列。当新请求处理完后，工作线程会继续处理待处理的工作，发布writes以填满硬件流水线。随后，轮询完成队列以查询已完成的传输和立即数计数器增量。事件被聚合以传递每次传输的通知，并将传输交给所有组共享的一个专用回调线程。

  • 分片灵活性: DOMAINGROUP内的分片是灵活的。传输可以通过索引指定特定的NIC。单个WRITE可以被拆分。分页传输、scatter和barrier操作都会转化为多个WRITEs，可以在所有NIC上进行分片。

3.5 硬件特定优化

TransferEngine内的DOMAIN针对其控制的硬件进行了特化和优化：

  • AWS EFA: 我们使用libfabric实现EFA支持，在DOMAINGROUP内为每个NIC管理一个fabric domain。由于EFA与RDMA规范不同，它对仅含立即数的零大小写入不要求有效的目标描述符，我们强制所有传输都使用有效的描述符。对于批量传输和对等组，我们采用工作请求（WR）模板化，在发布前预先填充并保留libfabric描述符的公共字段。

  • NVIDIA ConnectX-7: 我们通过libibverbs实现ConnectX支持。对于每个对端，我们使用一个UD队列对来交换RC握手信息。我们为每个对端创建2个RC队列对：一个用于双边SEND/RECV操作，另一个用于单边WRITE和WRITEIMM操作。这种分离是必要的，因为RECV和WRITEIMM完成都会按发布顺序消耗工作请求。这使我们能够在支持WRITEIMM的同时提供高层级的RECV语义而不产生干扰。除了WR模板化，我们还通过ibv_send_wr的next指针链接多达4个工作请求，从而采用WR链式操作，减少了对NIC的doorbell rings次数。此外，我们启用了IBV_ACCESS_RELAXED_ORDERING，允许NIC和GPU内存之间进行无序的PCIe事务，从而降低延迟。

4 KVCACHE TRANSFER

本节概述了一个依赖TransferEngine、经过生产测试的解耦式推理实现。在解耦模式下，一个预填充（prefiller）节点对输入令牌运行预填充，并将生成的KV页以及任何额外上下文（如用于推测解码的最后一个令牌的隐藏状态和logits）传输到解码器（decoder）节点，解码器节点随后逐个解码令牌。

图3. 预填充器和解码器之间的KV传输
图3. 预填充器和解码器之间的KV传输

  • 工作流程: 如图3所示，收到请求后，全局调度器选择一个预填充节点和一个解码器节点来处理它，并将请求转发给解码器。解码器预先分配KV页和任何上下文所需的存储空间，然后使用submit_send将请求分派到指定的预填充器，指示内容应传输到的KV页索引。

  • 分块预填充与传输触发: 在分块预填充期间，我们在每层的注意力输出投影后递增UVM观察者的值，这与CUDA Graph兼容。一旦TransferEngine检测到UVM观察者值的变化，该层的传输就会启动，通过submit_paged_writes从当前块复制页面。当最后一个块完成并且上下文已填充后，它会通过submit_single_write被复制过去。预填充器使用submit_recvs等待命令。

  • 处理分片差异: KV传输必须考虑预填充器和解码器之间KV缓存分片或复制的差异。MLA【18, Liu et al., Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024】在使用张量并行时会复制压缩的KV缓存条目：在这种方案下，预填充器的rank会与解码器的rank随机匹配，以平衡副本的传输。在GQA【2, Ainslie et al., Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023】下，我们依靠页级偏移和步长从源KV缓存中选择切片，复制到目标KV缓存的相应偏移处。为最小化写入次数并确保单个写入足够大，KV缓存的布局是头（heads）在页（pages）之前，以确保障连续头内的连续性。

  • 完成与错误处理: 预填充器不发送明确的完成消息：解码器预先知道它期望的传输数量，并使用expect_imm_count来由TransferEngine通知传输完成并开始解码。生产级解耦解码实现的复杂性在于错误和取消的处理。由解码器触发的取消必须由预填充器明确确认，因为只要存在远程写入覆盖它们的可能性，KV页就不能被重用。我们依靠预填充器和解码器之间的心跳消息来检测传输层故障。如果预填充节点无响应，请求在解码器上超时后被取消，因为传输无法再到达它。每个请求的取消令牌可以停止所有未来的传输，同时在发送取消确认之前等待所有挂起的操作。

图4. 不同方法的权重传输数据路径。
图4. 不同方法的权重传输数据路径。
图5. 流水线化的权重传输执行。
图5. 流水线化的权重传输执行。

5 RL ROLLOUT WEIGHT TRANSFER

在异步强化学习微调中，训练和推理在不同的GPU上运行。每次训练步骤后，新权重必须被推送到推理节点，对于万亿参数模型，使用现有框架可能需要几十到几百秒。我们的解决方案为Kimi-K2（1T参数）、DeepSeek V3（671B）和Qwen3（235B）等规模的模型实现了1.3秒的跨机器参数更新【12, Kimi Team et al., Kimi k2: Open agentic intelligence, 2025】【6, DeepSeek-AI et al., Deepseek-v3 technical report, 2025】【37, Yang et al., Qwen3 technical report, 2025】，将权重从256个训练GPU（bf16）传输到128个推理GPU（fp8）。

5.1 点对点权重传输

  • 方法对比：现有框架倾向于为所有训练和推理GPU形成一个全局集体通信世界。权重被收集到训练子组的Rank0，然后广播到每个推理子组的Rank0，这受到训练Rank0的NIC瓶颈限制。相比之下，在我们的P2P方法中，每个训练GPU通过单边RDMA WRITE直接向推理GPU发送权重，从而利用了所有NIC的全部集群带宽。图4展示了这两种方法的区别。

  • 初始化步骤：在初始化时，控制器脚本执行三个步骤：首先，它从所有训练和推理GPU收集参数元数据，包括权重名称、形状、数据类型和DTensor分片信息。其次，它计算一个静态的权重传输计划，映射哪个训练GPU将哪个参数发送到哪个推理GPU，以及发送顺序。最后，它将该计划广播给所有训练GPU。在每个训练步骤，控制器向训练GPU发信号开始发送权重。由于使用单边操作，推理节点对传输过程无感知。

5.2 流水线化权重传输执行

  • 并行化策略：我们的训练任务使用FSDP【41, Zhao et al., PyTorch FSDP: Experiences on scaling fully sharded data parallel, 2023】对模型权重进行分片。不同类型的参数（例如MoE与非MoE）需要不同的FSDP分片策略。每种分片策略将全局DeviceMesh划分为不相交的子网格（sub-mesh）。我们将每个子网格称为一个MeshGroup。MeshGroup内的参数并行传输，而MeshGroup之间则顺序处理。

  • 四阶段流水线: 我们将每个参数张量的传输视为一个任务。为同时利用不同硬件资源，我们将每个任务分为四个时间上重叠的流水线阶段，如图5所示：(1) 主机到设备memcpy，如果FSDP将权重卸载到CPU；(2) 参数准备：使用full_tensor()重建完整权重，应用投影融合，如果需要则进行量化；(3) RDMA传输：零拷贝WRITE到远程推理GPU内存；(4) 全局屏障：在所有full_tensor()调用完成后，通过以太网使用GLOO在网格组之间同步。

  • 内存管理：full_tensor()和其他GPU操作会引入额外的GPU内存使用。为避免内存不足错误，我们仅在当前正在进行的任务占用的临时GPU内存少于一个可配置的水印时才启动新任务。

6 MOE DISPATCH/COMBINE

我们介绍了一套围绕TransferEngine构建的用于MoE调度和合并的低延迟内核，它依赖一个主机代理线程来协调GPU和NIC。在节点内，我们还利用NVLink来减少网络负载。尽管代理线程增加了延迟，我们在保持预填充性能有竞争力且无需任何调整的情况下，实现了业界领先的解码性能。这些内核展示了基于代理的MoE调度在支持更广泛网络卡（如EFA）上的可行性。因此，我们专注于解码性能（每个rank 128个token），因为它受延迟限制，并且更容易受到跨设备增加的PCIe、驱动和固件开销的影响。

图6. Dispatch和Combine的GPU-CPU-NIC协调
图6. Dispatch和Combine的GPU-CPU-NIC协调
6.1 架构

  • 设计概览: 路由是使用分离的内核实现的，用于调度和合并。发送方的一半将数据准备到发送缓冲区以写入对端，接收方的一半则将来自接收缓冲区的令牌重排到其他内核使用的张量中。单个内核充分利用所有SM和GPU的内存带宽以减少延迟，但它们的运行时间足够短，可以在其间插入其他工作，从而实现重叠和微批处理。主机代理使用GDRCopy轮询GPU的进度，在源缓冲区就绪时调用TransferEngine。

  • 减少开销：我们的设计（如图6所示）通过减少发出的写入次数来最小化代理开销。为了减少发送和接收缓冲区所需的GPU内存，所有对端首先交换路由信息（即每个专家的令牌计数），以确定一个连续接收缓冲区中的唯一范围进行写入。由于这些载荷很小，它们的延迟可以通过投机性地将少量令牌分派到私有的每个源的缓冲区中来隐藏。合并操作仅发出一个scatter，因为它重用了路由信息。路由信息总是由TransferEngine处理，但载荷可以在同一节点内通过NVLink复制。因此，对于每个节点间的对端，每个rank在调度时最多发出2次WRITE，在合并时发出1次WRITE。

  • 缓冲区大小：接收缓冲区的大小必须能容纳发送到当前rank的所有令牌。假设有$N$个rank，每个rank上有$E$个专家，每个rank将$T$个令牌分派给$R$个专家，那么上限是$N \cdot T \cdot \max(R, E)$。发送方将写入打包到这样一个连续的缓冲区中，而不是依赖于更大的每个rank的接收缓冲区。为了最小化开销，较大的调度接收缓冲区可以被合并发送重用。

6.2 Dispatch

  • 调度流程: dispatch内核接收令牌和它们必须被路由到的$R$个专家的索引，返回一个张量，该张量打包了分配给本地专家的所有rank的令牌。该内核首先在共享内存中计算发送到每个专家的令牌数量，并通过统一内存将计数传输到主机。然后它通知代理，代理使用TransferEngine启动向所有对端rank的路由scatter。输入令牌随后被复制到发送缓冲区，创建一个连续的源，以便scatter到每个对端，如图7所示。
    图7. Dispatch到私有和连续缓冲区

  • 延迟隐藏: 代理被通知将最多固定数量的令牌scatter到每个接收端对端的私有缓冲区中。对于解码大小的批次，从dispatch内核启动到第一次传输的延迟大约为15 µs（假设EP=64）。此时，路由和令牌的传输使NIC带宽饱和。一旦接收到路由并且可以在每个目标rank上确定来自每个源rank的令牌位置，剩余的令牌就被scatter到对端，将它们连续地放入共享的接收缓冲区。主机端处理路由和分派第二轮传输的工作需要几十微秒。初始传输中的令牌数量被选择来隐藏这个开销。虽然这个延迟不在关键路径上，但减少它可以进一步减小私有缓冲区的大小。一旦所有传输被确认为已完成并且所有传入的写入都已记账，一个屏障通过TransferEngine同步代理。

  • NVLink传输与同步: 在通过RDMA向节点间对端分派传输后，发送内核继续通过NVLink发出写入，以在同一节点内传输令牌，而RDMA传输在后台挂起。NVLink对端通过它们自己的一组屏障进行同步：在写入对端之前，每个rank必须确保数据已被读取并且可以被覆盖。这是通过使用松散语义写入和读取的标志来完成的。在载荷写入后，rank之间通过release-acquire标志进行同步。

  • 写操作排序优化: 由于内存屏障的粒度不足，发送内核中的写入顺序对延迟至关重要。NVLink通过虚拟内存暴露，透明地将对映射到当前地址空间的对端设备的读写操作转换为互连上的事务。加载操作通常开销很大，因为它们会阻塞执行流水线直到满足。相比之下，存储操作是“即发即忘”的，直到遇到内存屏障，该屏障会阻塞直到其范围内所有先前的存储完成。由于主机系统和NVLink对端都在同一范围内，确保与主机排序的屏障可能会被先前通过NVLink发出的写入减慢。为避免这种情况，我们首先向主机发信号，然后在网格屏障后发出NVLink写入。这个策略增加了发送内核的总执行时间，但减少了到第一次RDMA传输的关键路径上的延迟。

  • NVLink读写策略: 使用NVLink时，通常最好从源设备向目标推送数据，以节省一次往返时间。此外，在当前设备上确认存储后，可以在传输到远程设备的过程中执行有用的工作。在dispatch阶段，我们只将令牌推送到私有接收缓冲区，因为此时中心化的路由信息尚不可用，无法确定其余令牌在对端上的确切位置。内核的接收部分通过在屏障上同步并读取剩余的令牌来启动。这些加载很可能在RDMA操作之前完成。

  • 接收端处理: 在dispatch发送后，传输在后台运行。当接收内核被调用时，它等待TransferEngine通过IMMCOUNTER和GDRCopy报告所有传输的完成。依靠基于路由信息计算的索引，令牌被重新排序，并根据适用于分组GEMM内核的布局在专家之间进行可选的填充。由于接收内核必须在关键路径上处理大约$T \cdot R$个令牌，工作被分配到所有可用的SM上，并最大化流水线以充分利用可用的HBM带宽。

6.3 Combine

  • 合并流程: 由于路由信息在dispatch阶段是中心化的，combine操作通过一次scatter传输所有载荷。命令在发送和合并之间的空闲时间内准备好，在此期间GPU通常会执行一个分组GEMM。发送方与dispatch类似，准备发送缓冲区并通过NVLink向节点内对端推送令牌。主机代理被通知向TransferEngine发送SCATTER请求，之后内核完成执行。接收方在等待所有令牌接收完毕之前，将所有从路由信息派生的相关偏移缓存在共享内存中。然后，在每个rank上本地计算令牌的加权平均值。

  • 资源重用与同步: combine阶段重用了与dispatch阶段相同的缓冲区，因此它需要一个NVLink屏障和一个RDMA屏障，以确保在覆盖发送缓冲区之前所有先前的操作都已完成。此时，主机代理也等待TransferEngine确认所有写入都已发送出去。特别是对于EFA，它还需要等待接收确认，因此最大化发布写入和检查其状态之间的时间间隔非常重要。

6.4 与DeepEP的比较

  • 对比分析: DeepEP提供了业界领先的延迟，但由于依赖IBGDA和mlx5驱动，它们被绑定在ConnectX上。我们的内核更具可移植性，依赖于主机代理，同时支持EFA和ConnectX。尽管有额外的开销，我们的延迟超过了DeepEP。

  • 实现差异: DeepEP内核依赖于RC QP的强排序保证。令牌在可用的SM之间平衡，并通过一个QP逐个传输。这确保了到第一次传输的延迟较低，但也意味着每个令牌需要更多的工作，并导致通过网络发送更多的包。令牌计数和完成通过ATOMICs进行信号。相比之下，我们的内核在第一次传输发起前花费更多时间，因为有额外的GPU-CPU-NIC通信开销（通过PCIe）。然而，批量传输实现了更好的网络利用率。

  • 预填充性能: 对于预填充，我们无需调整即可扩展我们的单次传输策略。相比之下，DeepEP通过在发送节点上使用NVLink预累积令牌来获得更好的延迟，减少了传输的数据量。此外，DeepEP内核使用的缓冲区内存更少，因为令牌的子集是分批传输的。虽然这种方法更快，但它也对准确性和确定性有影响，因为累积并不完全在fp32上以固定顺序完成。

A4 实验环境 (总结)

  • 硬件配置:

    • GPU: 每个节点配备8块NVIDIA H200 GPU，通过NVLink连接。
    • 网络: 每个GPU搭配一个400 Gbps的ConnectX-7网卡或两个200 Gbps的EFA网卡。
    • CPU: 双路Intel Sapphire Rapids CPU。

  • 软件配置:

    • 评测对象: TransferEngine。
    • 对比库: NIXL v0.6.1, DeepEP, pplx-kernels (基于NVSHMEM v3.4.5)。
    • 底层工具: rdma-core (用于ConnectX基准测试), libfabric (用于EFA基准测试)。
    • 模型与参数: MoE实验的设置参考了DeepSeek-V3模型。

  • 实验负载:

    • 点对点通信: 测试不同消息大小下的单次写入（Single Write）和分页写入（Paged Write）吞吐量。
    • MoE调度/合并: 在8、16、32和64个GPU上评估解码（每批次128个token）和预填充（每块4096个token）的延迟。每个token被分派到8个随机专家。

A5 实验结果 (总结)

7.1 点对点通信

  • 实验内容: 评估了TransferEngine和NIXL v0.6.1在单次写入（Single Write）和分页写入（Paged Write）场景下的带宽性能，并与底层硬件基准进行了比较。
  • 实验结果 (图8, 表2):
    • 带宽饱和: 对于单次写入，需要至少16MiB的消息才能饱和带宽。而对于分页写入，TransferEngine和NIXL分别只需32KiB和64KiB的消息即可饱和。
    • 性能对比: TransferEngine的性能与NIXL相当，略微更快。EFA需要比ConnectX-7更大的消息才能达到饱和。
    • 具体数据: 对于256 KiB的单次写入（MoE路由的典型大小），TransferEngine在EFA上达到54 Gbps，在ConnectX-7上达到116 Gbps。对于64 KiB的分页写入（KV Cache页的典型大小），EFA和ConnectX-7都能饱和可用带宽。
      图8. 点对点通信性能
      表2. EFA和ConnectX-7性能比较

7.2 MoE调度/合并

  • 实验内容: 在不同GPU规模（8, 16, 32, 64）下，将本文的MoE内核与DeepEP和pplx-kernels（基于NVSHMEM）进行解码和预填充延迟的比较。
7.2.1 私有缓冲区大小的影响 (图9)
  • 实验内容: 改变用于隐藏路由信息交换延迟的私有缓冲区的大小，观察总解码分派时间的延迟变化。
  • 结论: 私有缓冲区对于隐藏延迟至关重要。对于节点内通信，至少需要约32个令牌来隐藏路由交换延迟。对于节点间通信，ConnectX-7可以使用少至24个令牌，而EFA在少于32个令牌时性能会下降，因为其路由交换更慢。
    图9. 私有缓冲区大小对p50解码延迟的影响
7.2.2 发送和接收延迟 (图10)
  • 实验内容: 分别测量发送（send）和接收（receive）内核的执行延迟。
  • 结论: 本文的发送内核（仅内存复制）和合并接收内核（因更快的累加）均比DeepEP快。内核总执行时间不到传输时间的15%，表明大部分时间花在了通信上。主机代理在发送内核执行中途就开始了RDMA工作，效率很高。
    图10. EP=64时独立的发送和接收延迟
7.2.3 解码延迟 (图11)
  • 实验内容: 比较不同实现方案在解码大小批次下的端到端延迟。
  • 结论:
    • ConnectX-7: 在16和32个GPU的节点间场景下，本文内核的性能超过了DeepEP，成为新的业界最佳，这得益于批量传输和高效的流水线。在64个GPU时，主机代理的CPU开销变得明显，调度延迟略高。
    • EFA: 本文的实现使EFA在MoE场景下变得可行，延迟仅比ConnectX-7高约30%。
    • 与NVSHMEM对比: 本文内核比基于NVSHMEM的pplx-kernels快一个数量级。
    • 节点内 (8 GPU): 略慢于DeepEP（约2µs），这主要因为使用NIC交换路由信息。
      图11. MoE解码延迟。条形图高度为均值。误差棒显示p01, p25, p50, p75, p95, p99。pplx的误差棒表示标准差。
7.2.4 预填充延迟 (图12)
  • 实验内容: 比较不同实现方案在预填充大小批次下的端到端延迟。
  • 结论: DeepEP在预填充场景下表现更好。在调度阶段，DeepEP通过NVLink在节点内复制令牌，减少了RDMA传输量。在合并阶段，DeepEP的发送端部分求和策略极大地减少了RDMA数据量，从而延迟更低（尽管这会牺牲一些累积精度）。本文为解码优化的内核在预填充时存在内存开销过高的问题。
    图12. MoE预填充延迟。条形图高度为均值。误差棒显示p01, p25, p50, p75, p95, p99。

A6 结论 (总结)

现有的用于LLM系统的RDMA解决方案存在供应商锁定问题，尤其是在AWS EFA等定制云硬件上缺乏可行的实现。本文提出的TransferEngine通过识别异构RDMA硬件间的共同功能（可靠但无序的交付）来解决这一问题。通过在底层协议之上构建一个不依赖排序保证的可靠抽象层，我们透明地将支持扩展到多种RDMA NIC，特别是EFA和ConnectX。

我们通过三个生产系统展示了这种方法的有效性：用于解耦式推理的KvCache传输、为万亿参数模型实现1.3秒更新的RL权重更新，以及在ConnectX-7上实现业界领先延迟并在AWS EFA上首次实现可行性能的MoE调度/合并。TransferEngine为现代LLM架构提供了可移植的点对点通信能力，在补充集体通信库的同时，避免了供应商锁定，尤其适用于云原生部署。