在 AWS EFA 上启用万亿参数模型

在 Perplexity，我们为我们的产品、API 以及研究团队使用最优模型。大型开源混合专家模型（Mixture-of-Experts models，如 Kimi-K2）带来了特殊的挑战，因为最大的具备 8x NVIDIA H200 GPU 的推理节点无法有效容纳它们，这使得多节点部署成为必要。我们展示了一套用于专家并行性的内核，它们在 ConnectX-7 上实现了业界领先的延迟，超越了 DeepEP 的性能。同样的内核也是首个在 AWS Elastic Fabric Adapter (EFA) 上实现可行延迟的内核，从而使万亿参数模型的部署成为可能。

请在 GitHub 上试用我们的内核，并在 arXiv 上阅读完整的研究论文。

混合专家模型（MoE）已成为扩展至数千亿或万亿参数，同时在推理时提供合理延迟的标准架构。MoE 模型将 transformer 的稠密层替换为一组专家层和一个小型线性路由层，后者决定将 token 与哪些专家相乘。通过将关联相同专家的 token 分组，该问题变得稀疏且易于并行化，因为不同的专家现在可以托管在分布于多个节点的不同 GPU 上。

与其他形式的并行性（如张量并行性 (TP) 和数据并行性 (DP)，它们可以很容易地使用诸如 torch.distributed 和 NCCL 等集合通信库来实现）不同，MoE 路由涉及稀疏的点对点通信。分发内核（dispatch kernel）必须拆分一组传入的 token 并将其分发到托管专家的等级（ranks），而其对偶问题——组合内核（combine kernel）——则必须接收它们处理后的对应项并计算其加权平均值。这些传输无法映射到集合操作，需要定制内核才能实现低延迟。

在 Perplexity，我们一直在开发可移植的节点间内核和专门的节点内内核，并在之前关于 MoE 通信的文章中更深入地介绍了我们的工作。我们还在另一篇博客文章中表明，多节点部署可以为大型 MoE 模型带来更低的延迟和更高的吞吐量。

在本文中，我们将深入探讨我们新的节点间内核的实现，这些内核极大地改进了我们之前基于 NVSHMEM 的工作的性能。除了展示在 ConnectX-7 上业界领先的结果外，我们还证明了在我们的推理运行时中，这些内核可以有效地服务于最大的开源模型。

张量并行和专家并行可以在单个节点内有效扩展，因为 NVLink 在微秒级延迟下提供了高吞吐量 (900GB/s)。虽然 NVIDIA GB200 和 GB300 系统配备了足够的内存来容纳万亿参数模型的权重和 KV 缓存，但在具有最多 8 个 H200 GPU 的 AWS p5en 实例上的部署仅限于 1120GB 的 HBM，这些必须在模型权重和 KV 缓存之间进行分配。因此，某些模型需要在 InfiniBand 上进行节点间部署。

尽管 InfiniBand 速度极快，但它仅提供 400Gbps (50GB/s) 的吞吐量，为 MoE 路由增加了数十到数百微秒的延迟。幸运的是，通过使用共享专家、微批处理（micro-batching）和计算-通信重叠，可以掩盖增加的延迟，从而以具有竞争力的成本提供可行的吞吐量。DeepSeek 通过在 H800 GPU 集群和 NVIDIA ConnectX-7 InfiniBand 适配器上部署 DeepSeek-R1，开创了这些技术。

在 AWS 集群上，由于使用了其定制的 Elastic Fabric Adapters (EFA)，情况有所不同。即使这些适配器在训练工作负载中通常使用的集合通信上提供 400Gbps 的峰值吞吐量，但在 MoE 分发和组合期间交换的消息大小方面，它们略逊于 ConnectX-7。此外，EFA 不支持 GPUDirect Async，需要一个 CPU 代理问题来架起 GPU 和 NIC 之间的桥梁以启动传输。除了增加的复杂性之外，额外的 PCIe 事务还为所有传输增加了微秒级的开销。诸如 NVSHMEM 之类的框架允许内核编写者提供独立于设备的实现，然而通用 API 会产生开销，导致部署在成本和延迟方面效率低下。

我们之前的内核通过 NVSHMEM 抽象了底层的传输实现。它们依赖 nvshmemx_putmem_signal_nbi_warp 逐个传输 token，并使用原子计数器来同步对等点。虽然这在具有 IBGDA 的 ConnectX-7 NIC 上提供了合理的性能，但基于代理的 ConnectX-7 IBRC 实现明显较慢，而 EFA 实现未能提供低于 1 毫秒的延迟。

我们将这些内核用作底层硬件的测试平台，同时确定向我们的模型公开的最合适的接口。我们推测，内核与 NIC 或代理问题之间频繁的传输和交互阻碍了性能。此外，我们的实现不依赖于除了可靠性之外的任何来自硬件的保证，这表明我们可以通过减少同步原语的开销来加快速度。

根据这些观察，我们设计了基于 EFA 的内核，这些内核公开了相同的接口，并在底层优化了设备内核与主机代理的交互，以实现峰值性能。

基于主代理的 RDMA 和 NVLink 混合通信，展示发起和接收角色、代理工作线程以及 NIC。

新的内核遵循混合 CPU-GPU 架构，其中 GPU 内核与设备上的模型交互，而在主机 CPU 上运行的代理问题处理与 NIC 的交互。新的内核重用了我们为 KV 缓存传输构建的相同 TransferEngine。

MoE 路由由分发内核执行，该内核接收 token 以及模型中用于路由它们的每个 token 的专家，并通过网络将它们运送到托管专家的等级。在分组 GEMM 内核将 token 与专家权重相乘后，组合内核将 token 传输到其源等级，并在将其返回模型之前计算其加权平均值。

分发内核和组合内核都拆分为发送端和接收端，以促进重叠和微批处理。发送者将 token 洗牌（shuffled）到发送缓冲区，以便可以通过单个 RDMA 写入操作将其写入对等点。执行发送者后，GPU 可以执行其他工作（如共享专家乘法），而网络传输则在后台挂起。接收器内核阻塞直到所有传输都被接收。在后台，主机侧的工作问题处理传输：当发送缓冲区被填充时，它由 GPU 发出信号，之后它会分发 RDMA 写入。一旦 NIC 发出完成信号，它就会解除接收器的阻塞，接收器将 token 返回给模型。

内核与主机紧密耦合，通过统一内存和 GDRCopy 进行通信。CUDA 统一内存允许将设备内存或主机内存映射到两者的地址空间中，通过 PCIe 透明地处理通信。

虽然统一内存在批量传输方面很适用，但低延迟轮询是通过 GDRCopy 完成的。GDRCopy 使用 GPUDirect RDMA 以微秒级延迟读取和写入 GPU 内存。

除了 RDMA 之外，在节点内，内核还利用 NVLink 来交换 token。这可以卸载很大一部分流量，对于 EP32 高达四分之一，对于 EP64 高达八分之一，从而实现整体更低的延迟。

分发内核的发送缓冲区设计对于实现涉及主机代理的可行延迟至关重要。如果有 N 个等级，每个等级分发 T 个 token，每个 token 路由到 R 个专家（共 E 个），则在最坏的情况下，如果所有 token 都路由到它，一个等级最多可以接收 N * T * max(R, E / N) 个 token。此外，只有当发送者连续布置其写入时，才能实现此最小界限，因为私有接收缓冲区会消耗过多的内存。为了使发送者能够连续且无冲突地布置其写入，路由信息的完整知识必须分发到每个等级。因此，我们的分发内核首先交换每个专家的 token 计数，之后它们可以各自确定其写入的位置以及它们从每个等级接收 token 的偏移量。

仅仅交换路由信息会导致不必要的高延迟，因为每个等级在分发其自己的写入之前必须等待写入到达。相反，我们在每个接收者等级上保留少量私有的每个发送者空间，并随着路由信息的写入一起分发一些 token。这确保了在整个分发过程中保持完全的带宽利用率。交换路由后，剩余的 token 通过每个远程对等点的一次写入进行传输。路由信息的交换和固定大小的缓冲区，以及填充它们的传输，如下图所示。

MoE 派发的缓冲区布局，展示 token 数据和路由元数据如何在两个角色之间排列和交换。

组合内核要简单得多，因为路由信息在分发阶段对它们来说已经是可用的了。它们可以以每个对等点一次写入的方式复制 token，之后读取接收缓冲区以将 token 的加权平均值计算到输出张量中。

对于 NVLink，节点内的发送和接收缓冲区被映射到对等点中，以便它们能够直接从相邻的等级推送或拉取 token。

下图详细说明了 GPU、CPU 和 NIC 之间的时序和交互，展示了与模型内核交错的分发和组合的一次迭代。

工作从分发发送开始，内核在单个块内通过共享内存中的原子累积来汇总路由信息并计算每个专家的 token 计数。一旦 token 计数可用，它们就会被复制到主机内存，代理会立即收到通知，开始将路由信息分发给所有其他对等等级。并行地，分发继续计算 token 计数的累积和，确定将 token 打包到连续的每个对等发送缓冲区所需的偏移量。一旦偏移量可用，所有要通过 EFA 发送的 token 都会被打包到发送缓冲区中。代理会收到通知，将写入发布到节点间对等点的发送缓冲区中。在节点内，内核还会通过 NVLink 立即将 token 推送到同一节点内的私有缓冲区。

MoE 派发的时序与协调，展示分阶段的 GPU、CPU 和 NIC 操作及其同步。

一旦一个等级从所有对等点接收到路由信息，它就会对其进行集中处理，并计算识别分发和组合的 token 所需的各种偏移量。这项工作是在 CPU 上完成的，而传输已经在挂起，因此它对延迟不敏感。随后，一旦所有偏移量已知，超过私有缓冲区大小的其余 token 的写入就会被发布，并且分发的发送阶段完成。

在接收端，内核首先通过 NVLink 拉取同一节点中超出私有缓冲区容量的 token。接下来，它阻塞直到主机代理发出所有传输完成的信号。一旦收到信号，token 就会从接收缓冲区洗牌到传递给分组 GEMM 的输出张量中。由于缓冲区在分发和组合之间被重用，NVLink 和 RDMA 屏障（barriers）都确保在允许下一个内核覆盖它们之前已经读取了接收器缓冲区。

组合内核遵循大致相似的模式，洗牌 token，向代理发出信号并等待完成。然而，用于发布其写入的元数据对它们来说是现成的，并且所有远程偏移量都是已知的，因此所有节点内传输的 token 都通过 NVLink 推送到接收者，而节点间 token 通过单次写入传输。在返回结果之前，内核会累积来自接收缓冲区的 token，将它们写入输出张量。

为了建立我们的实现原型，我们最初重用了为分离式预填充（disaggregate prefill）和解码构建的 Transfer Engine，以便发布和跟踪 RDMA 写入。然而，与每个节点向同一对等点发布大量写入的分离式解码不同，MoE 分发在一个操作中向所有对等点发布少量传输。现成的性能并不理想，我们的实现落后于 ConnectX-7 上的 DeepEP 约 20us。因此，我们专门化了 TransferEngine 以支持两种操作：scatter 和 barrier，以优化与大量对等点的交互。

scatter 和 barrier 操作都公开了对等组，可以预先注册这些组以预填充底层传输实现使用的数据结构（工作请求 (WR) 模板化）。scatter 将来自同一源缓冲区的不同切片复制到对等组的所有成员上的不同偏移量中。barrier 通过立即数（immediate）向所有对等点发出信号。批量排队和处理这些操作可以在关键路径上节省宝贵的微秒。在 EFA 上，两个 200Gbps NIC 聚合以提供 400Gbps 带宽，我们将对等组分片（shard）在两个 NIC 上，而不是分片每次传输的字节。

虽然我们最初仅使用 libfabric 在 EFA 上开发 TransferEngine，但在使用 libibverbs 在 TransferEngine 中添加 ConnectX-7 支持后，MoE 内核无需更改任何一行代码即可在 ConnectX-7 上工作。与基于可扩展可靠数据报 (SRD) 的 EFA 相比，ConnectX-7 需要额外的工作来进行连接设置和对等管理。我们使用不可靠数据报 (UD) 队列对来引导可靠连接 (RC)。我们为每个远程对等点创建两个队列对，一个用于双向 SEND/RECV，一个用于单向 WRITE_IMM，因为 RECV 和 WRITE_IMM 都按顺序消耗 RECV 操作，而这两个操作需要不同的缓冲区大小。此外，我们放宽了 ConnectX-7 和 GPU 内存之间的 PCIe 事务排序。与 EFA 类似，我们实现了 WR 模板化优化。此外，当向同一对等点发送多条小消息时，我们使用 WR 链接来提高消息速率。

由于 GPU 侧的工作在两个适配器上是相同的，我们将底层的传输实现推向了极限，以超越 DeepEP 的性能，从而突出 EFA 和 ConnectX-7 MoE 路由之间的差距是由底层硬件的能力决定的。

我们在配备 ConnectX-7 和 EFA 适配器的 H200 系统上评估了分发和组合延迟，采用标准的 DeepSeek-V3 和 Kimi-K2 配置，具有 7168 的隐藏维度，块缩放 fp8 分发和 bf16 组合。我们考虑分发批量大小为 128，组合批量大小为 4096。在 ConnectX-7 上，我们与 DeepEP 和我们之前基于 NVSHMEM 的 pplx-内核进行了比较，同时使用了 GPUDirect Async (IBGDA) 和主机代理 (IBRC)。我们在 EP16、EP32 和 EP64 上分别实现了 459us、582us 和 692us 的组合延迟。我们优于其他依赖通用代理的基于 EFA 的内核（如 UCCL-EP），后者实现了 519us、966us 和 1159us。

在解码时，尽管我们的内核比 DeepEP 慢了几微秒，但我们实现了整体更低的延迟，因为我们的组合内核明显更快。虽然我们使用了主机代理，但由于我们的批量传输，我们超过了所有基于 IBGDA 的无代理实现的性能。我们还表明代理实现很重要，因为我们高度专门化的代理比通用的 NVSHMEM 实现快一个数量级。尽管 EFA 在传输 256KB 数据包（接近 MoE 路由交换的典型数据包大小）时比 ConnectX-7 慢两倍，但两个适配器之间的差距总体上并不显著。

EP64 MoE 预填充的派发和合并延迟，对比 pplx 和 DeepEP 通信路径。

即使我们的解码优化内核可以在没有任何额外调整的情况下扩展到预填充批量大小，性能仍落后于 DeepEP。这主要是因为我们仅在组合的接收端累积 token，而 DeepEP 在发送端累积以减少传输的数据量。然而，这种优化可能会对准确性产生影响。此外，由于内核在一次写入中传输数据，我们为发送和接收缓冲区使用了更多的内存。

为了证明我们内核的可行性，我们基准测试了我们内部推理引擎的端到端性能，该引擎在配备 Elastic Fabric Adapters (EFA) 和 H200 GPU 的多个 p5en 实例上服务 DeepSeek-V3（671B 参数）和 Kimi-K2（1T 参数）。我们以草稿验证模式运行模型，并假设使用单个 MTP 层时的接受率为 80% 来报告延迟。

在 DeepSeek-V3 上，我们与运行高效仅 NVLink 分发和组合内核的 EP=8 节点内基准进行了比较。节点间部署总体上要么匹配，要么超过了单节点部署的性能，特别是在中等大小的批次上，这证明了跨节点扩展确实带来了整体吞吐量的改进。

在 Kimi-K2 上，这些内核实现了服务。由于其大小，模型无法使用数据并行 (DP) 容纳在单个 H200 节点内，因此需要节点间部署来服务该模型。由于模型具有较少的注意力头，我们实现了比 DeepSeek-V3 更好的延迟，并且我们能够以可行的延迟在 EFA 上服务模型。

DeepSeek-V3 在 EP32、EP16 和 EP8 配置下的单 GPU 吞吐量与输出速度对比。

Kimi-K2 在 EP32 和 EP16 配置下的单 GPU 吞吐量与输出速度对比。

我们还考虑了微批处理，尽管它仅在非常高的批量大小或使用多于一个草稿 token 时才改进了延迟。为了从将分组 GEMM 问题拆分为两个中受益，token 批量大小必须足够密集。大部分时间仍然花费在 MoE 分发和组合上，这表明对内核的任何进一步改进都将相应地改善延迟并促进微批处理。

我们正在与 AWS 工程师密切合作，以提高 EFA 上的性能。在我们的端，我们将关注 libfabric 的更新，这将减少数据平面开销。此外，我们将尝试 efa-direct，以进一步减少用户空间开销并直接与硬件接口。

使用 NVIDIA Magnum IO NVSHMEM 和 GPUDirect Async 提高 HPC 系统的网络性能

BibTeX

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}

简介

基于 InfiniBand 的推理

基于 NVSHMEM 的 pplx-内核

基于 EFA 的 MoE 全对全通信

发送和接收缓冲区

分发与组合

用于 MoE 路由的 TransferEngine

评估

分发/组合延迟

Kimi-K2 / DeepSeek-V3 吞吐量

未来工作

参考资料

引用