RDMA komunikacija točka-po-točka za LLM sustave

LLM sustavi tradicionalno se oslanjaju na kolektivne komunikacije za strukturirane paralelne uzorke u treniranju i inferenciji. Međutim, nova radna opterećenja, kao što su disagregirana inferencija, usmjeravanje mješavine stručnjaka (MoE) i asinkrono fino podešavanje učenja potkrepljivanjem, zahtijevaju fleksibilnost točku-po-točku komunikacije.

Naš najnoviji istraživački rad predstavlja TransferEngine, prijenosnu komunikacijsku biblioteku temeljenu na RDMA-u koja radi na različitim mrežnim sučeljnim karticama (NIC), uključujući NVIDIA ConnectX-7 i AWS Elastic Fabric Adapter (EFA). Pokazujemo njenu učinkovitost u tri produkcijska sustava u Perplexity AI:

Disagregirano predpopunjavanje i dekodiranje: Prijenos KvCache-a sloj po sloj korištenjem straničnog pisanja i UVM promatrača, što omogućuje da se instance predpopunjavanja i dekodiranja skaliraju neovisno.

Prijenos težina za RL post-trening za manje od 2 sekunde: Nevjerojatno brza sinkronizacija težina između čvorova za treniranje i inferenciju.

Omogućavanje modela s bilijun parametara na AWS EFA: Niskolatencijske operacije usmjeravanja i kombiniranja za MoE korištenjem primitiva raspršivanja i barijere, čime se postiže latencija najsuvremenije razine na ConnectX-7 i prva održiva implementacija kompatibilna s AWS EFA.

Ovaj se članak fokusira na sam TransferEngine, komunikacijsku biblioteku koja pokreće ove sustave. Pročitajte cijeli istraživački rad na arXiv i istražite kod na GitHub.

Kolektivne komunikacije izvrsne su za strukturirane paralelne uzorke kao što su tenzorska paralelnost i paralelnost podataka, ali nameću ograničenja koja ograničavaju njihovu primjenjivost na nova radna opterećenja:

Fiksno članstvo: Kolektivi zahtijevaju uspostavljanje globalnog komunikacijskog svijeta u vrijeme inicijalizacije, uz sinkronizirano postavljanje među svim sudionicima. Za disagregiranu inferenciju to je problematično budući da se broj instanci predpopunjavanja i dekodiranja mora dinamički skalirati kao odgovor na prometne obrasce.

Ujednačeni oblici tenzora: Kolektivi pretpostavljaju da svi sudionici razmjenjuju tenzore identičnih oblika i tipova podataka. Kod MoE usmjeravanja, broj tokena poslanih svakom stručnjaku ovisi o podacima i otkriva se u vrijeme izvođenja, što primorava guste all-to-all kolektive da se dopune do veličina međuspremnika za najgori slučaj.

Redoslijed operacija: Kolektivni API-ji nameću sinkroniziranu semantiku, čak i kada aplikacije trebaju samo znati da je skup operacija dovršen. Na primjer, kod prijenosa KvCache-a, sve stranice iz instance predpopunjavanja u istancu dekodiranja mogu stići bilo kojim redoslijedom. Slično tome, kod ažuriranja RL težina, svi tenzori težina mogu se prenijeti neovisno bez ograničenja redoslijeda.

RDMA-temeljena komunikacija točka-po-točka prirodno rješava ova ograničenja, ali su tri prepreke spriječile njeno usvajanje:

Ograničenja biblioteke: Dok torch.distributed i NCCL pružaju API-je točka-po-točka kao što su Send/Recv, ove operacije i dalje nasljeđuju kolektivnu semantiku i ograničenja.

Raznolikost hardvera: NVIDIA ConnectX mrežne kartice s pouzdanim prijenosom veze (RC) najčešća su RDMA postava, koja jamči isporuku po redoslijedu. Pružaoci računalnog oblaka poput AWS-a koriste drugačiji hardver, Elastic Fabric Adapter (EFA) s prijenosom pouzdanog podatkovnog grama koji se može skalirati (SRD) koji isporučuje poruke izvan redoslijeda. EFA također zahtijeva agregiranje četiri 100 Gbps mrežne kartice (ili dvije 200 Gbps mrežne kartice) kako bi se postigla propusnost od 400 Gbps.

Prijenosnost softvera: Postojeće RDMA biblioteke pate od zaključavanja od strane dobavljača. DeepEP zahtijeva RDMA iniciran od strane GPU-a (IBGDA) dostupan samo na ConnectX-u. NVSHMEM pokazuje ozbiljno smanjenje performansi na EFA-i. Biblioteke poput Mooncake i NIXL nemaju podršku za EFA ili su u ranoj fazi razvoja. Kao rezultat toga, ne postoji održivo rješenje za više pružatelja usluga za LLM inferenciju.

TransferEngine rješava ove prepreke kroz četiri temeljne sposobnosti:

Uobičajene pretpostavke prijenosa: Umjesto ciljanja određenog mrežnog sučelja, TransferEngine pretpostavlja samo da je temeljni prijenos pouzdan, ali neuređen, tj. presjek jamstava koja pružaju i RC i SRD prijenos.

ImmCounter sinkronizacija: Nova primitiva za obavijest o dovršetku koja radi bez oslanjanja na redoslijed poruka, omogućujući da se isti kod izvodi i na uređenim (RC) i na neuređenim (SRD) prijenosima.

Podrška za više mrežnih kartica: Transparentno upravljanje s više mrežnih kartica po GPU-u, automatskim dijeljenjem prijenosa radi maksimiziranja propusnosti.

Minimalna površina API-ja: Mali skup operacija dovoljan za izgradnju složenih komunikacijskih uzoraka.

Dizajn sustava TransferEngine, koji prikazuje putanje zahtjeva i povratnih poziva, CPU radnike, UVM međuspremnike, GPU-ove i NIC-ove kroz dva GPU/NUMA domena.

TransferEngine pokreće jednu radnu nit po GPU-u. Svaki radnik upravlja DomainGroup-om koja koordinira 1 do 4 RDMA mrežne kartice, ovisno o hardverskoj platformi (jedan ConnectX-7 pruža 400 Gbps, dok EFA zahtijeva četiri adaptera od 100 Gbps ili dva adaptera od 200 Gbps). Unutar DomainGroup-a, svaka Domain je specifična za hardver i upravlja jednom mrežnom karticom, rješavajući upravljanje parom redova čekanja, podnošenje radova i anketiranje dovršetaka.

Svaka instanca TransferEngine-a izlaže jednu glavnu adresu za identifikaciju i otkrivanje vršnjaka. Udaljeni vršnjaci razmjenjuju NetAddr strukture koje kodiraju mrežne adrese i mogu se serijalizirati za izvanpojasnu komunikaciju. Svi vršnjaci u sustavu moraju koristiti isti broj mrežnih kartica po GPU-u. Ova ujednačena konfiguracija omogućuje TransferEngine-u inteligentno dijeljenje prijenosa: prilikom pokretanja prijenosa, i izvor i odredište točno znaju koje su mrežne kartice dostupne, čime se omogućuje automatsko balansiranje opterećenja na svim adapterima.

TransferEngine izlaže minimalan API koji apstrahira složenost temeljnih RDMA sučelja.

rust

#[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>;
}

API se sastoji od nekoliko ključnih komponenti:

RDMA operacije zahtijevaju da memorijska područja budu registrirana kod mrežne kartice prije prijenosa podataka. Funkcija reg_mr uzima pokazivač, duljinu i uređaj, vraćajući dvije ručke:

MrHandle: Lokalna ručka koja se koristi kao izvor za prijenose

MrDesc: Serijalizabilni opisnik koji se može razmijeniti s udaljenim vršnjacima kako bi im se omogućilo pisanje u registriranu memoriju

Struktura MrDesc sadrži popis udaljenih mrežnih adresa i udaljenih ključeva (jedan po mrežnoj kartici) koji odobravaju udaljeni pristup memorijskom području. Ovaj dizajn prirodno rješava postave s više mrežnih kartica.

Za male poruke i poruke upravljačke ravnine u stilu RPC-a, TransferEngine pruža submit_send i submit_recvs:

submit_send: Šalje poruku udaljenom vršnjaku identificiranom s NetAddr. Poruka se kopira interno, tako da pozivatelj može odmah ponovno upotrijebiti ili otpustiti međuspremnik.

submit_recvs: Objavljuje skup međuspremnika za primanje. Kada poruka stigne, poziva se povratni poziv s primljenim podacima. Međuspremnik se automatski vraća u skup nakon što povratni poziv završi.

Za prijenose masovnih podataka, TransferEngine nudi jednostrane operacije pisanja bez kopiranja:

submit_single_write: Prenosi uzastopno područje s izvora na odredište

submit_paged_writes: Prenosi neuzastopne stranice određene indeksima, koracima i pomacima

Obje operacije prihvaćaju opcionalnu 32-bitnu neposrednu vrijednost koja se isporučuje primatelju atomski po dovršetku.

Pogon automatski dijeli velike prijenose na više mrežnih kartica kako bi maksimizirao propusnost.

Za operacije usmjerene na više vršnjaka (uobičajeno u MoE usmjeravanju), TransferEngine nudi optimizirane primitive:

add_peer_group: Unaprijed registrira skup adresa vršnjaka za brzo pretraživanje

submit_scatter: Šalje različite odsječke podataka svakom vršnjaku u grupi

submit_barrier: Šalje obavijest koja sadrži samo neposredne podatke svim vršnjacima (prijenos od nula bajtova)

Ove su operacije optimizirane za nisku latenciju predmemoriranjem unutarnjih struktura podataka i izradom predložaka RDMA radnih zahtjeva, čime se izbjegava ponavljajući trošak postavljanja.

Ključni izazov u LLM sustavima je usklađivanje RDMA prijenosa s napretkom GPU-a. U idealnom slučaju, prijenos podataka trebao bi početi čim GPU proizvede izlaz, ali anketiranje GPU-a s CPU-a je neučinkovito.

TransferEngine to rješava pomoću UVM promatrača: unificiranih lokacija virtualne memorije koje GPU jezgre mogu ažurirati (čak i unutar CUDA grafova) i koje CPU može učinkovito anketirati pomoću GDRCopy-ja. Kada se dodijeli UVM promatrač, registrira se povratni poziv koji se pokreće kad god GPU ažurira vrijednost.

To omogućuje uzorke poput prijenosa KvCache-a sloj po sloj: nakon što svaki sloj transformatora završi, GPU inkrementira promatrač, pokrećući RDMA prijenos KvCache-a tog sloja dok naknadni slojevi nastavljaju s računanjem.

TransferEngine pruža mehanizme obavijesti o dovršetku i za stranu pokretača (pozivatelja submit_write) i za stranu primatelja (cilj operacija pisanja).

Obavijest na strani pokretača: Kada se podnese operacija pisanja, pozivatelj može odabrati primanje obavijesti po dovršetku putem funkcije povratnog poziva ili atomske zastavice. Ova obavijest potvrđuje da su podaci uspješno preneseni iz izvornog međuspremnika.

Obavijest na strani primatelja: Ključna inovacija je ImmCounter, komponenta koja prati brojače indeksirane 32-bitnim neposrednim brojevima. Kada operacija pisanja nosi neposrednu vrijednost, ImmCounter primatelja se atomski inkrementira po isporuci podataka. Funkcija expect_imm_count omogućuje primatelju registraciju povratnog poziva koji se aktivira kada određena neposredna vrijednost dostigne ciljani broj. Brojači se mogu transparentno sinkronizirati s GPU-om putem GDRCopy-ja, izravno anketirati ili nadzirati putem povratnih poziva.

Ovaj mehanizam na strani primatelja pruža pouzdanu sinkronizaciju bez zahtijevanja redoslijeda poruka. Jednako dobro funkcionira na ConnectX mrežnim karticama (koje jamče isporuku po redoslijedu preko RC prijenosa) i EFA mrežnim karticama (koje isporučuju poruke izvan redoslijeda preko SRD prijenosa).

TransferEngine je implementiran u Rustu, pažljivo optimizirajući alokacije, niti i sinkronizaciju kako bi se smanjila latencija i postigao visok protok. Pogon pokreće jednu radnu nit po DomainGroup-u, od kojih je svaka prikvačena za CPU jezgru na NUMA čvoru na koji su uređaji spojeni, minimizirajući latenciju raspoređivanja i pristupa memoriji. Strukture podataka specifične za domenu alociraju se nakon prikvačivanja kako bi se osiguralo da je memorija rezervirana u ispravnom NUMA čvoru. Jedna radna nit rukuje s do 4 Domain-a (od kojih svaka upravlja jednom mrežnom karticom), dok druga namjenska nit anketira GPU radi ažuriranja UVM promatrača. Komunikacija između niti obavlja se kroz redove čekanja bez zaključavanja.

API proslijeđuje zahtjeve DomainGroup-u koji poslužuje odgovarajući uređaj, pri čemu prvi također poslužuje host. U tijesnoj petlji, radni domenski proces anketira novi rad, dajući prioritet podnošenju novih zahtjeva. Ovisno o hardveru i konfiguraciji, zahtjevi se dijele i opterećenje se balansira na dostupnim Domain-ima. Prvo pisanje složenog zahtjeva odmah se objavljuje u red čekanja za slanje mrežne kartice. Kada se novi zahtjevi iscrpe, radnik prelazi na napredovanje operacija na čekanju, objavljujući pisanja kako bi popunio hardverski cjevovod. Nakon toga se anketiraju redovi čekanja dovršetaka kako bi se dohvatili završeni prijenosi i povećanja neposrednih brojača. Događaji se agregiraju radi isporuke obavijesti po prijenosu, prepuštajući prijenos namjenskoj niti povratnog poziva koju dijele sve grupe.

Dijeljenje unutar DomainGroup-a je fleksibilno. Prijenosi mogu ciljati određene mrežne kartice po indeksu. Jedno pisanje može se podijeliti. Prijenosi stranica, raspršivanje i operacije barijere, koji se svi prevode u višestruka pisanja, mogu se dijeliti na svim mrežnim karticama.

Domain-i unutar TransferEngine-a specijalizirani su i optimizirani za hardver pod njihovom kontrolom:

AWS EFA: EFA implementacija koristi libfabric, upravljajući mrežnom domenom po mrežnoj kartici unutar DomainGroup-a. Ključne odluke implementacije i optimizacije uključuju:

Pisanja nulte veličine: EFA odstupljuje od RDMA specifikacije zahtijevajući valjane ciljne opisnike čak i za pisanja nulte veličine. Stoga TransferEngine provodi valjane opisnike za sve prijenose.

Izrada predložaka radnih zahtjeva: Unaprijed popunjava i zadržava uobičajena polja libfabric struktura prije objavljivanja, smanjujući opterećenje po operaciji za masovne prijenose i operacije grupa vršnjaka

NVIDIA ConnectX-7: ConnectX-7 implementacija koristi libibverbs s RC prijenosom. Ključne odluke implementacije i optimizacije uključuju:

UD parovi redova čekanja za rukovanja: Svaki Domain koristi par redova čekanja nepouzdanog datagrama (UD) za razmjenu rukovanja RC vezom.

Odvojeni RC parovi redova čekanja: Stvara 2 RC para redova čekanja po vršnjaku, jedan za dvostrane Send/Recv operacije, a drugi za jednostrane Write i WriteImm operacije. Ovo odvajanje je neophodno jer i Recv i WriteImm dovršetci troše radne zahtjeve redoslijedom objavljivanja, dopuštajući visokorazinsku Recv semantiku uz istovremenu podršku za WriteImm bez smetnji.

Izrada predložaka radnih zahtjeva: Unaprijed popunjava uobičajena polja struktura radnih zahtjeva za masovne prijenose i grupe vršnjaka.

Ulančavanje radnih zahtjeva: Povezuje do 4 radna zahtjeva kroz next pokazivač ibv_send_wr-a, smanjujući broj zvonjenja zvona (doorbell rings) prema mrežnoj kartici.

Opušteni redoslijed: Omogućuje IBV_ACCESS_RELAXED_ORDERING kako bi se dopustile PCIe transakcije izvan redoslijeda između mrežne kartice i GPU memorije.

Ciljajući presjek jamstava u različitim RDMA prijenosima i uvođenjem sinkronizacije neovisne o redoslijedu putem ImmCounter-a, TransferEngine pruža prenosivu osnovu za komunikaciju točka-po-točka u LLM sustavima. Biblioteka je otvorena koda i dostupna na GitHub-u, te potičemo zajednicu da je istraži i nadograđuje za nove komunikacijske uzorke u sustavima strojni učenja velikih razmjera.

BibTeX

@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}
}

Motivacija

TransferEngine

Dizajn API-ja

Registracija memorije

Dvostrano slanje/primanje (Send/Recv)

Jednostrano pisanje

Grupe vršnjaka i masovne operacije

UVM promatrač: Signalizacija s GPU-a na CPU

Obavijest o dovršetku

Implementacija

Optimizacije specifične za hardver

Zaključak

Citiranje