跳到正文
IFAI

开源项目

Perplexity 开源 MoE 通信库:GPU 间传输提速 10 倍,大幅缩短推理延迟

首次发布 传播指数 41

Perplexity 开源了 MoE 模型通信库 pplx-kernels,通信速度比标准 All-to-All 提升最高达 10 倍。它通过 NVSHMEM 绕过 CPU 搞 GPU 直连传输,并让计算与通信重叠进行。最关键的是兼顾了跨网络环境的可移植性,帮大规模 MoE 模型推理显著省下延迟。

Perplexity 开源 MoE 通信库:GPU 间传输提速 10 倍,大幅缩短推理延迟
图:perplexity.ai

pplx-kernels 是 Perplexity 开源的 MoE 通信内核库,专门用来解决多 GPU 跑大体量混合专家模型时的通信延迟瓶颈。它把跨卡分发和汇总 token 的通信速度,拉高到了标准 All-to-All 原语的最高 10 倍。

如果你在部署 DeepSeek R1 或 Mixtral 这类 MoE 架构模型,就会发现算力往往不是最大瓶颈,卡在 GPU 之间的等待才最磨人。DeepSeek R1 总参数量高达 671B,每个 token 虽只激活 37B,但这些被激活的专家散落在不同节点的不同 GPU 上,跨设备传输不可避免。

它能解决什么问题

在传统的 PyTorch 方案中,多卡分发 token 基本依赖 torch 的 all_to_all。这种原生算子面对稀疏通信很不顺手,往往要求数据填充对齐,还经常需要 GPU 和 CPU 反复同步,元数据广播也占满了带宽。

pplx-kernels 就是为了干掉这部分冗余开销。它专为 MoE 的调度和汇总阶段设计,负责把 token 精准送给持有对应专家的显卡,计算完成后再通过加权求和拉回结果。

在单节点配置下,它的通信延迟比此前业内最快的实现还要低 2.5 倍。换到跨节点的多机环境下,它的传输表现最高能达到标准通信原语的 10 倍。

它是怎么做到的

这个库的核心思路很明确:绕过 CPU,并且把网络传输藏进计算里。

首先是 GPU 直通网卡。借助 ConnectX 网卡的 IBGDA 特性,GPU 可以直接向网卡提交 RDMA 操作,自己去轮询完成队列,全程不需要 CPU 插手干预。即便在不支持这项特性的硬件上,它也能退回到代理线程模式运行。

其次是拆分内核架构。它把发送和接收阶段完全拆开,形成 split kernel 结构,让网络传输的同时,GPU 还能继续跑本地计算,两者重叠进行。

更有意思的是它对 NVSHMEM 的精简用法。官方 NVSHMEM 提供了庞大的接口体系,但频繁的栅栏同步极耗时间。pplx-kernels 基本上只依靠三个底层函数来构建数据搬运与信号同步,全部采用非阻塞调用,把跨卡依赖降到了最低。

兼容性与获取方式

相比那些针对特定超算集群深度魔改的方案,这个库最大的卖点是兼顾了通用性。它虽然比某些极致特化的高手方案慢大约 2 倍,但不需要你绑定死特定的硬件环境。

它同时支持单机单卡间高达 900Gbps 的 NVLink 通信,以及跨节点多机之间的 InfiniBand、ConnectX-7 和 AWS EFA 网络,对不同 NVSHMEM 版本的兼容性也很扎实。

目前该项目已经完全开源,代码托管在 GitHub 的 ppl-ai/pplx-kernels 仓库中,开发者可以直接拉取编译集成。

为什么重要

我们认为这个项目的价值在于切中了当前开源大模型落地最难啃的工程细节。DeepSeek R1 这类超大 MoE 模型的普及,让很多团队第一次面对几十甚至上百张 GPU 的分布式推理,而通信延迟往往是吞掉显卡算力的罪魁祸首。

对于自建集群推理服务的工程团队来说,pplx-kernels 提供了一个现成且不挑网络拓扑的高性能解法。你不需要推翻已有的基础设施去专门定制,就能直接拿它替换掉效率低下的原生算子,大幅削减跨机推理的等待时间。

来源 (1)

  1. [1]
    perplexity.ai(在新标签页打开)原文perplexity.ai

每日 AI 内容精选