跳到正文
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 內容精選