RDMA-Punkt-zu-Punkt-Kommunikation für LLM-Systeme
Elegantes Tool für LLM-Kommunikationsmuster
LLM-Systeme stützten sich bisher auf kollektive Kommunikationen für strukturierte parallele Muster beim Training und bei der Inferenz. Neue Arbeitslasten wie die disaggregierte Inferenz, das Mixture-of-Experts (MoE)-Routing und die asynchrone Feinabstimmung durch Bestärkendes Lernen erfordern jedoch die Flexibilität der Punkt-zu-Punkt-Kommunikation.
Unser neuestes Forschungspapier stellt TransferEngine vor, eine portable RDMA-basierte Kommunikationsbibliothek, die über verschiedene Netzwerkschnittstellen-Controller (NICs) hinweg funktioniert, einschließlich NVIDIA ConnectX-7 und AWS Elastic Fabric Adapter (EFA). Wir demonstrieren ihre Effektivität anhand von drei Produktionssystemen bei Perplexity AI:
Disaggregierte Prefill- und Decode-Phase: Layer-für-Layer KvCache-Transfer unter Verwendung von Paged-Writes und UVM-Watchern, wodurch Prefill- und Decode-Instanzen unabhängig voneinander skaliert werden können.
Gewichtstransfer für RL-Post-Training in unter 2 Sekunden: Blitzschnelle Gewichtssynchronisierung zwischen Trainings- und Inferenz-Knoten.
Ermöglichung von Billion-Parameter-Modellen auf AWS EFA: MoE-Dispatch- und Combine-Operationen mit niedriger Latenz unter Verwendung von Scatter- und Barrier-Primitiven, wodurch eine hochmoderne Latenz auf ConnectX-7 erreicht wird und die erste tragfähige Implementierung vorliegt, die mit AWS EFA kompatibel ist.
Dieser Artikel konzentriert sich auf TransferEngine selbst, die Kommunikationsbibliothek, die diese Systeme betreibt. Lesen Sie das vollständige Forschungspapier auf arXiv und erkunden Sie den Code auf GitHub.
Motivation
Kollektive Kommunikationen zeichnen sich durch strukturierte parallele Muster wie Tensor-Parallelität und Datenparallelität aus, bringen jedoch Einschränkungen mit sich, die ihre Anwendbarkeit auf neue Arbeitslasten begrenzen:
Feste Mitgliedschaft: Kollektive erfordern die Einrichtung einer globalen Kommunikationsumgebung zum Zeitpunkt der Initialisierung, mit synchronisierter Einrichtung über alle Teilnehmer hinweg. Bei der disaggregierten Inferenz ist dies problematisch, da die Anzahl der Prefill- und Decode-Instanzen dynamisch als Reaktion auf Verkehrsmuster skaliert werden muss.
Einheitliche Tensor-Formen: Kollektive setzen voraus, dass alle Teilnehmer Tensoren mit identischen Formen und Datentypen austauschen. Beim MoE-Routing ist die Anzahl der an jeden Experten gesendeten Token datenabhängig und wird zur Laufzeit ermittelt, was dichte All-to-All-Kollektive dazu zwingt, auf Worst-Case-Puffergrößen aufzufüllen.
Reihenfolge der Operationen: Kollektive APIs erzwingen synchronisierte Semantiken, auch wenn Anwendungen nur wissen müssen, dass eine Reihe von Operationen abgeschlossen ist. Zum Beispiel können beim KvCache-Transfer alle Seiten von der Prefill-Instanz zur Decode-Instanz in beliebiger Reihenfolge eintreffen. Ähnlich können bei RL-Gewichtsaktualisierungen alle Gewichtstensoren unabhängig und ohne Reihenfolgebeschränkungen übertragen werden.
RDMA-basierte Punkt-zu-Punkt-Kommunikation adressiert diese Einschränkungen auf natürliche Weise, aber drei Hindernisse haben ihre Einführung bisher verhindert:
Bibliotheksbeschränkungen: Während torch.distributed und NCCL Punkt-zu-Punkt-APIs wie Send/Recv bereitstellen, erben diese Operationen dennoch kollektive Semantiken und Einschränkungen.
Hardware-Vielfalt: NVIDIA ConnectX NICs mit Reliable Connection (RC)-Transport sind das gebräuchlichste RDMA-Setup und garantieren eine geordnete Zustellung. Cloud-Anbieter wie AWS verwenden andere Hardware; der Elastic Fabric Adapter (EFA) mit Scalable Reliable Datagram (SRD)-Transport liefert Nachrichten ungeordnet aus. EFA erfordert zudem die Aggregation von vier 100-Gbps-NICs (oder zwei 200-Gbps-NICs), um eine Bandbreite von 400 Gbps zu erreichen.
Software-Portabilität: Bestehende RDMA-Bibliotheken leiden unter Herstellerbindung. DeepEP erfordert GPU-initiiertes RDMA (IBGDA), das nur auf ConnectX verfügbar ist. NVSHMEM weist auf EFA eine starke Leistungsverschlechterung auf. Bibliotheken wie Mooncake und NIXL unterstützen EFA nicht oder befinden sich noch in der frühen Entwicklung. Infolgedessen existiert keine tragfähige, anbieterübergreifende Lösung für die LLM-Inferenz.
TransferEngine
TransferEngine begegnet diesen Hindernissen durch vier Kernfunktionen:
Gemeinsame Transportannahmen: Anstatt auf eine spezifische NIC abzuzielen, geht TransferEngine lediglich davon aus, dass der zugrunde liegende Transport zuverlässig, aber ungeordnet ist, d. h. die Schnittmenge der Garantien, die sowohl von RC- als auch von SRD-Transporten bereitgestellt werden.
ImmCounter-Synchronisierung: Ein neuartiges Primitiv für die Benachrichtigung bei Abschluss, das funktioniert, ohne sich auf die Nachrichtenreihenfolge zu verlassen, und es ermöglicht, denselben Code sowohl auf geordneten (RC) als auch auf ungeordneten (SRD) Transporten auszuführen.
Multi-NIC-Unterstützung: Transparente Verwaltung mehrerer NICs pro GPU, wobei Transfers automatisch aufgeteilt werden, um die Bandbreite zu maximieren.
Minimale API-Oberfläche: Eine kleine Menge an Operationen, die ausreicht, um komplexe Kommunikationsmuster aufzubauen.
TransferEngine startet einen Worker-Thread pro GPU. Jeder Worker verwaltet eine DomainGroup, die 1 bis 4 RDMA-NICs koordiniert, abhängig von der Hardware-Plattform (eine einzelne ConnectX-7 bietet 400 Gbps, während EFA vier 100-Gbps-Adapter oder zwei 200-Gbps-Adapter erfordert). Innerhalb einer DomainGroup ist jede Domain hardwarespezifisch, verwaltet eine einzelne NIC und kümmert sich um das Queue-Pair-Management, das Einreichen von Aufgaben und das Completion-Polling.
Jede TransferEngine-Instanz stellt eine einzige Hauptadresse für die Peer-Identifizierung und -Erkennung bereit. Remote-Peers tauschen NetAddr-Strukturen aus, die Netzwerkadressen kodieren und für die Out-of-Band-Kommunikation serialisiert werden können. Alle Peers in einem System müssen die gleiche Anzahl an NICs pro GPU verwenden. Diese einheitliche Konfiguration ermöglicht es TransferEngine, Transfers intelligent aufzuteilen: Beim Initiieren eines Transfers wissen sowohl Quelle als auch Ziel genau, welche NICs verfügbar sind, was einen automatischen Lastausgleich über alle Adapter hinweg ermöglicht.
API-Design
TransferEngine stellt eine minimale API bereit, die die Komplexität der zugrunde liegenden RDMA-Schnittstellen abstrahiert.
#[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;
// Memory Region Management
fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc);
// Two-sided Send/Recv
fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ());
fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ());
// One-sided 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);
// One-sided Write to a group of peers
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>);
// Watcher for CPU-GPU synchronization
fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>;
}Die API besteht aus mehreren Schlüsselkomponenten:
Speicherregistrierung
RDMA-Operationen erfordern, dass Speicherbereiche vor dem Datentransfer bei der NIC registriert werden. Die Funktion reg_mr nimmt einen Pointer, eine Länge und ein Gerät entgegen und gibt zwei Handles zurück:
MrHandle: Ein lokales Handle, das als Quelle für Transfers verwendet wird
MrDesc: Ein serialisierbarer Deskriptor, der mit Remote-Peers ausgetauscht werden kann, damit diese in den registrierten Speicher schreiben können
Die Struktur MrDesc enthält eine Liste von Remote-Netzwerkadressen und Remote-Keys (einen pro NIC), die den Remote-Zugriff auf den Speicherbereich gewähren. Dieses Design handhabt Multi-NIC-Setups auf natürliche Weise.
Zweiseitiges Senden/Empfangen
Für kleine Nachrichten und RPC-artige Control-Plane-Nachrichten stellt TransferEngine submit_send und submit_recvs bereit:
submit_send: Sendet die Nachricht an einen Remote-Peer, der durch NetAddr identifiziert wird. Die Nachricht wird intern kopiert, sodass der Aufrufer den Puffer sofort wiederverwenden oder freigeben kann.
submit_recvs: Stellt einen Pool von Empfangspuffern bereit. Wenn eine Nachricht eintrifft, wird ein Callback mit den empfangenen Daten aufgerufen. Der Puffer wird nach Abschluss des Callbacks automatisch in den Pool zurückgegeben.
Einseitiges Schreiben
Für den Transfer großer Datenmengen bietet TransferEngine Zero-Copy One-Sided-Write-Operationen:
submit_single_write: Überträgt einen zusammenhängenden Bereich von der Quelle zum Ziel
submit_paged_writes: Überträgt nicht zusammenhängende Seiten, die durch Indizes, Strides und Offsets angegeben sind
Beide Operationen akzeptieren einen optionalen 32-Bit-Sofortwert (Immediate Value), der dem Empfänger atomar bei Abschluss zugestellt wird.
Die Engine verteilt große Transfers automatisch auf mehrere NICs, um die Bandbreite zu maximieren.
Peer-Gruppen und Bulk-Operationen
Für Operationen, die auf mehrere Peers abzielen (üblich beim MoE-Routing), bietet TransferEngine optimierte Primitive:
add_peer_group: Registriert vorab eine Menge von Peer-Adressen für eine schnelle Suche
submit_scatter: Sendet unterschiedliche Datensegmente an jeden Peer in der Gruppe
submit_barrier: Sendet eine Benachrichtigung an alle Peers (0-Byte-Transfer)
Diese Operationen sind auf niedrige Latenz optimiert, indem interne Datenstrukturen zwischengespeichert und RDMA-Arbeitsanfragen als Vorlagen verwendet werden, wodurch redundanter Overhead bei der Einrichtung vermieden wird.
UVM-Watcher: GPU-zu-CPU-Signalisierung
Eine zentrale Herausforderung in LLM-Systemen ist die Koordination von RDMA-Transfers mit dem Fortschritt der GPU. Idealerweise sollte der Datentransfer beginnen, sobald die GPU Ausgaben erzeugt, aber das Polling der GPU durch die CPU ist ineffizient.
TransferEngine begegnet dem mit UVM-Watchern: einheitliche virtuelle Speicherorte, die von GPU-Kerneln (sogar innerhalb von CUDA-Graphen) aktualisiert und von der CPU effizient mittels GDRCopy abgefragt werden können. Wenn ein UVM-Watcher zugewiesen wird, wird ein Callback registriert, das aufgerufen wird, sobald die GPU den Wert aktualisiert.
Dies ermöglicht Muster wie den Layer-für-Layer KvCache-Transfer: Nach Abschluss jedes Transformer-Layers inkrementiert die GPU den Watcher, was einen RDMA-Transfer des KvCaches dieses Layers auslöst, während nachfolgende Layer weiter berechnet werden.
Benachrichtigung bei Abschluss
TransferEngine bietet Mechanismen zur Benachrichtigung bei Abschluss sowohl für die Initiatorseite (der Aufrufer von submit_write) als auch für die Empfängerseite (das Ziel von Write-Operationen).
Benachrichtigung auf Initiatorseite: Wenn eine Schreiboperation eingereicht wird, kann der Aufrufer wählen, bei Abschluss über eine Callback-Funktion oder ein atomares Flag benachrichtigt zu werden. Diese Benachrichtigung bestätigt, dass die Daten erfolgreich aus dem Quellpuffer übertragen wurden.
Benachrichtigung auf Empfängerseite: Die wichtigste Neuerung ist der ImmCounter, eine Komponente, die Zähler verfolgt, welche durch 32-Bit-Sofortwerte indexiert sind. Wenn eine Schreiboperation einen Sofortwert enthält, wird der ImmCounter des Empfängers bei der Datenzustellung atomar inkrementiert. Die Funktion expect_imm_count ermöglicht es dem Empfänger, einen Callback zu registrieren, der ausgelöst wird, wenn ein bestimmter Sofortwert einen Zielzähler erreicht. Die Zähler können transparent mit der GPU via GDRCopy synchronisiert, direkt abgefragt oder über Callbacks überwacht werden.
Dieser mechanismus auf Empfängerseite ermöglicht eine zuverlässige Synchronisierung, ohne dass eine Nachrichtenreihenfolge erforderlich ist. Er funktioniert gleichermaßen gut auf ConnectX-NICs (die eine geordnete Zustellung über RC-Transport garantieren) und EFA-NICs (die Nachrichten ungeordnet über SRD-Transport zustellen).
Implementierung
TransferEngine ist in Rust implementiert, wobei Zuweisungen, Threading und Synchronisierung sorgfältig optimiert wurden, um die Latenz zu minimieren und einen hohen Durchsatz zu erreichen. Die Engine startet einen Worker-Thread pro DomainGroup, die jeweils an einen CPU-Kern auf dem NUMA-Knoten gebunden sind, an den die Geräte angeschlossen sind, was sowohl das Scheduling als auch die Speicherzugriffslatenz minimiert. Domainspezifische Datenstrukturen werden nach dem Binden zugewiesen, um sicherzustellen, dass der Speicher im richtigen NUMA-Knoten reserviert ist. Ein Worker-Thread verarbeitet bis zu 4 Domains (jede verwaltet eine einzelne NIC), während ein weiterer dedizierter Thread die GPU abfragt, um UVM-Watcher zu aktualisieren. Die Kommunikation zwischen den Threads erfolgt über lock-freie Warteschlangen.
Die API leitet Anfragen an die DomainGroup weiter, die das entsprechende Gerät bedient, wobei die erste auch den Host bedient. In einer engen Schleife fragt der Domain-Worker nach neuen Aufgaben und priorisiert das Einreichen neuer Anfragen. Abhängig von Hardware und Konfiguration werden Anfragen aufgeteilt und über die verfügbaren Domains hinweg lastverteilt. Der erste Write einer zusammengesetzten Anfrage wird sofort in die Sendewarteschlange der NIC eingereiht. Sobald neue Anfragen erschöpft sind, fährt der Worker fort, ausstehende Operationen voranzutreiben und Schreibvorgänge einzureihen, um die Hardware-Pipeline zu füllen. Anschließend werden die Completion-Queues abgefragt, um nach abgeschlossenen Transfers und Inkrementen von Sofortzählern zu suchen. Ereignisse werden aggregiert, um Benachrichtigungen pro Transfer zuzustellen, wobei der Transfer an einen dedizierten Callback-Thread übergeben wird, der von allen Gruppen gemeinsam genutzt wird.
Das Sharding innerhalb einer DomainGroup ist flexibel. Transfers können auf spezifische NICs nach Index abzielen. Ein einzelner Write kann aufgeteilt werden. Paged-Transfers, Scatter- und Barrier-Operationen, die alle in mehrere Writes übersetzt werden, können über alle NICs hinweg aufgeteilt werden.
Hardwarespezifische Optimierungen
Die Domains innerhalb von TransferEngine sind spezialisiert und auf die von ihnen gesteuerte Hardware optimiert:
AWS EFA: Die EFA-Implementierung verwendet libfabric und verwaltet eine Fabric-Domain pro NIC innerhalb der DomainGroup. Zu den wichtigsten Implementierungsentscheidungen und Optimierungen gehören:
Schreibvorgänge mit Nullgröße: EFA weicht von der RDMA-Spezifikation ab, indem gültige Zieldeskriptoren selbst für Schreibvorgänge mit Nullgröße erforderlich sind. Daher erzwingt TransferEngine gültige Deskriptoren für alle Transfers.
Arbeitsanfrage-Templating: Vorab-Bestückung und Beibehaltung der gemeinsamen Felder von libfabric-Strukturen vor dem Einreichen, wodurch der Overhead pro Operation für Bulk-Transfers und Peer-Gruppen-Operationen reduziert wird
NVIDIA ConnectX-7: Die ConnectX-7-Implementierung verwendet libibverbs mit RC-Transport. Zu den wichtigsten Implementierungsentscheidungen und Optimierungen gehören:
UD-Queue-Pairs für Handshakes: Jede Domain verwendet ein Unreliable Datagram (UD)-Queue-Pair, um RC-Verbindungshandshakes auszutauschen.
Separate RC-Queue-Pairs: Erstellt 2 RC-Queue-Pairs pro Peer, eines für zweiseitige Send/Recv-Operationen und eines für einseitige Write- und WriteImm-Operationen. Diese Trennung ist notwendig, da sowohl Recv- als auch WriteImm-Abschlüsse Arbeitsanfragen in der Reihenfolge des Einreichens verbrauchen, was übergeordnete Recv-Semantiken ermöglicht und gleichzeitig WriteImm ohne Interferenzen unterstützt.
Arbeitsanfrage-Templating: Vorab-Bestückung gemeinsamer Felder von Arbeitsanfragestrukturen für Bulk-Transfers und Peer-Gruppen.
Verkettung von Arbeitsanfragen: Verknüpft bis zu 4 Arbeitsanfragen über den next-Pointer von ibv_send_wr, wodurch die Anzahl der Doorbell-Aufrufe an die NIC reduziert wird.
Aufgeweichte Reihenfolge: Ermöglicht IBV_ACCESS_RELAXED_ORDERING, um ungeordnete PCIe-Transaktionen zwischen der NIC und dem GPU-Speicher zuzulassen.
Fazit
Durch das Anzielen der Schnittmenge von Garantien über verschiedene RDMA-Transporte hinweg und die Einführung einer reihenfolgeunabhängigen Synchronisierung mittels ImmCounter bietet TransferEngine ein portables Fundament für die Punkt-zu-Punkt-Kommunikation in LLM-Systemen. Die Bibliothek ist Open Source und auf GitHub verfügbar. Wir ermutigen die Community dazu, diese zu erkunden und für aufkommende Kommunikationsmuster in Machine-Learning-Systemen im großen Maßstab weiterzuentwickeln.
Zitieren
@online{perplexity2025RdmaPointToPointCommunicationForLlmSystems,
author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
title = {{RDMA point-to-point communication for LLM systems}},
date = {2025-11-05},
url = {https://www.perplexity.ai/hub/blog/rdma-point-to-point-communication-for-llm-systems}
}