RDMA komunikácia bod-bod pre systémy LLM

Systémy LLM sa tradične spoliehali na kolektívnu komunikáciu pre štruktúrované paralelné vzory pri trénovaní a inferencii. Vznikajúce pracovné záťaže, ako napríklad disagregovaná inferencia, smerovanie Mixture-of-Experts (MoE) a asynchrónne jemné dolaďovanie pomocou posilňovania (reinforcement learning fine-tuning), však vyžadujú flexibilitu bod-bod (point-to-point) komunikácie.

Náš najnovší výskumný príspevok predstavuje TransferEngine, prenosnú komunikačnú knižnicu založenú na technológii RDMA, ktorá funguje naprieč rôznymi sieťovými radičmi (NIC), vrátane NVIDIA ConnectX-7 a AWS Elastic Fabric Adapter (EFA). Jej účinnosť demonštrujeme v troch produkčných systémoch v spoločnosti Perplexity AI:

Disagregovaný Prefill a Decode: Prenos KvCache po vrstvách pomocou stránkovaných zápisov (paged writes) a sledovačov UVM, čo umožňuje nezávislé škálovanie inštancií prefillu a decodu.

Prenos váh pre RL post-trénovanie za menej ako 2 sekundy: Neuveriteľne rýchla synchronizácia váh medzi tréningovými a inferenčnými uzlami.

Podpora modelov s biliónmi parametrov na AWS EFA: Nízko latenčné operácie rozdeľovania (dispatch) a kombinovania MoE pomocou primitív scatter a barrier, ktoré dosahujú špičkovú latenciu na ConnectX-7 a predstavujú prvú životaschopnú implementáciu kompatibilnú s AWS EFA.

Tento článok sa zameriava na samotný TransferEngine, komunikačnú knižnicu, ktorá poháňa tieto systémy. Prečítajte si kompletný výskumný príspevok na arXiv a preskúmajte kód na GitHub.

Kolektívna komunikácia vyniká pri štruktúrovaných paralelných vzoroch, ako sú tenzorový paralelizmus (Tensor Parallelism) a dátový paralelizmus (Data Parallelism), avšak kládie obmedzenia, ktoré znižujú jej využiteľnosť pre nové typy pracovnej záťaže:

Pevné členstvo: Kolektívne operácie vyžadujú vytvorenie globálneho komunikačného sveta v čase inicializácie so synchronizovaným nastavením naprieč všetkými účastníkmi. Pri disagregovanej inferencii je to problematické, pretože počet inštancií prefillu a decodu sa musí dynamicky škálovať v reakcii na vzorce prenosov.

Jednotné tvary tenzorov: Kolektívne operácie predpokladajú, že všetci účastníci vymieňajú tenzory s identickými tvarmi a typmi údajov. Pri smerovaní MoE závisí počet tokenov odoslaných každému expertovi od dát a zistí sa až počas behu, čo núti husté kolektívne operácie typu all-to-all vypĺňať medzery (pad) na najhoršie veľkosti vyrovnávacej pamäte.

Poradie operácií: Kolektívne rozhrania API vynucujú synchronizovanú sémantiku, a to aj vtedy, keď aplikácie potrebujú iba vedieť, že množina operácií bola dokončená. Napríklad pri prenose KvCache môžu všetky stránky z inštancie prefillu do inštancie decodu doraziť v ľubovoľnom poradí. Podobne pri aktualizáciách váh v RL je možné všetky tenzory váh prenášať nezávisle bez obmedzení poradia.

Komunikácia typu bod-bod založená na RDMA prirodzene rieši tieto obmedzenia, avšak jej adoptovaniu bránili tri prekážky:

Obmedzenia knižníc: Zatiaľ čo torch.distributed a NCCL poskytujú rozhrania API pre komunikáciu bod-bod, ako sú Send/Recv, tieto operácie stále dedia kolektívnu sémantiku a obmedzenia.

Rôznorodosť hardvéru: Sieťové radiče NVIDIA ConnectX s transportom Reliable Connection (RC) predstavujú najbežnejšie nastavenie RDMA, ktoré zaručuje doručenie v správnom poradí. Poskytovatelia cloudových služieb, ako napríklad AWS, používajú iný hardvér – Elastic Fabric Adapter (EFA) s transportom Scalable Reliable Datagram (SRD), ktorý doručuje správy v nesprávnom poradí (out-of-order). EFA tiež vyžaduje agregáciu štyroch 100 Gbps sieťových radičov (alebo dvoch 200 Gbps sieťových radičov), aby sa dosiahla šírka pásma 400 Gbps.

Softvérová prenosnosť: Existujúce knižnice RDMA trpia viazanosťou na dodávateľa (vendor lock-in). DeepEP vyžaduje RDMA iniciované GPU (IBGDA), ktoré je dostupné iba na kartách ConnectX. NVSHMEM vykazuje vážny pokles výkonu na EFA. Knižnice ako Mooncake a NIXL nemajú podporu pre EFA alebo sú v počiatočnom štádiu vývoja. V dôsledku toho neexistuje žiadne životaschopné riešenie pre rôznych poskytovateľov pre inferenciu LLM.

TransferEngine prekonáva tieto prekážky pomocou štyroch kľúčových schopností:

Spoločné predpoklady o transporte: Namiesto zamerania sa na konkrétny sieťový radič predpokladá TransferEngine iba to, že podkladový transport je spoľahlivý, ale neusporiadaný, t. j. predstavuje prienik záruk poskytovaných transportmi RC aj SRD.

Synchronizácia ImmCounter: Nová primitíva na notifikáciu dokončenia, ktorá funguje bez spoliehania sa na usporiadanie správ, čo umožňuje spúšťať rovnaký kód na usporiadaných (RC) aj neusporiadaných (SRD) transportoch.

Podpora viacerých sieťových radičov (Multi-NIC): Transparentná správa viacerých sieťových radičov na jedno GPU, ktorá automaticky rozdeľuje prenosy na shard-y s cieľom maximalizovať šírku pásma.

Minimálny povrch API: Malá množina operácií postačujúca na vytvorenie komplexných komunikačných vzorov.

Návrh systému TransferEngine zobrazujúci cesty požiadaviek a callbackov, CPU workerov, UVM buffre, GPU a NIC v dvoch GPU/NUMA doménach.

TransferEngine spúšťa jedno pracovné vlákno na jedno GPU. Každý pracovný proces spravuje štruktúru DomainGroup, ktorá koordinuje 1 až 4 RDMA sieťové radiče v závislosti od hardvérovej platformy (jediný ConnectX-7 poskytuje 400 Gbps, zatiaľ čo EFA vyžaduje štyri 100 Gbps adaptéry alebo dva 200 Gbps adaptéry). V rámci štruktúry DomainGroup je každá Domain špecifická pre daný hardvér a spravuje jeden sieťový radič, pričom zabezpečuje správu párov radov (queue pairs), odosielanie úloh a zisťovanie dokončenia (polling).

Každá inštancia TransferEngine sprístupňuje jednu hlavnú adresu na identifikáciu a objavovanie rovesníkov (peers). Vzdialení rovesníci si vymieňajú štruktúry NetAddr, ktoré kódujú sieťové adresy a môžu byť serializované pre komunikáciu mimo pásma (out-of-band). Všetci rovesníci v systéme musia používať rovnaký počet sieťových radičov na jedno GPU. Táto jednotná konfigurácia umožňuje knižnici TransferEngine inteligentne rozdeľovať prenosy na shardy: pri iniciovaní prenosu vedia zdroj aj cieľ presne, ktoré sieťové radiče sú k dispozícii, čo umožňuje automatické vyvažovanie záťaže naprieč všetkými adaptérmi.

TransferEngine sprístupňuje minimálne API, ktoré abstrahuje zložitosť podkladových rozhraní RDMA.

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 pozostáva z niekoľkých kľúčových komponentov:

Operácie RDMA vyžadujú, aby boli pamäťové oblasti zaregistrované v sieťovom radiči pred prenosom dát. Funkcia reg_mr prijíma ukazovateľ, dĺžku a zariadenie a vracajú dva ovládače (handles):

MrHandle: Lokálny ovládač používaný ako zdroj pre prenosy

MrDesc: Serializovateľný deskriptor, ktorý možno vymeniť so vzdialenými rovesníkmi, aby mohli zapisovať do zaregistrovanej pamäte

Štruktúra MrDesc obsahuje zoznam vzdialených sieťových adries a vzdialených kľúčov (jeden pre každý sieťový radič), ktoré udeľujú vzdialený prístup k pamäťovej oblasti. Tento dizajn prirodzene zvláda konfigurácie s viacerými sieťovými radičmi.

Pre malé správy a správy riadiacej roviny v štýle RPC poskytuje TransferEngine funkcie submit_send a submit_recvs:

submit_send: Odosiela správu vzdialenému rovesníkovi identifikovanému pomocou NetAddr. Správa sa skopíruje interne, takže volajúci môže okamžite znovu použiť alebo uvoľniť vyrovnávaciu pamäť.

submit_recvs: Zverejňuje bazén (pool) vyrovnávacích pamätí na prijímanie. Keď dorazí správa, zavolá sa spätné volanie (callback) s prijatými údajmi. Vyrovnávacia pamäť sa po dokončení spätného volania automaticky vráti do bazéna.

Pre hromadné prenosy dát poskytuje TransferEngine operácie jednostranného zápisu s nulovou kopírovaním (zero-copy):

submit_single_write: Prenáša súvislú oblasť zo zdroja do cieľa

submit_paged_writes: Prenáša nesúvislé stránky určené indexmi, krokmi (strides) a posunmi (offsets)

Obidve operácie akceptujú voliteľnú 32-bitovú okamžitú hodnotu, ktorá sa príjemcovi doručí atomicky po dokončení.

Motor (engine) automaticky rozdeľuje veľké prenosy na shardy naprieč viacerými sieťovými radičmi s cieľom maximalizovať šírku pásma.

Pre operácie zamerané na viacerých rovesníkov (bežné pri smerovaní MoE) poskytuje TransferEngine optimalizované primitíva:

add_peer_group: Predbežne zaregistruje množinu adries rovesníkov pre rýchle vyhľadávanie

submit_scatter: Odosiela rôzne dátové časti každému rovesníkovi v skupine

submit_barrier: Odosiela notifikáciu obsahujúcu iba okamžitú hodnotu všetkým rovesníkom (prenos s nulovým počtom bajtov)

Tieto operácie sú optimalizované pre nízku latenciu prostredníctvom ukladania vnútorných dátových štruktúr do vyrovnávacej pamäte a vytvárania šablón pre požiadavky na prácu s RDMA, čím sa vyhýbajú opakovanej réžii nastavenia.

Kľúčovou výzvou v systémoch LLM je koordinácia prenosov RDMA s priebehom spracovania na GPU. V ideálnom prípade by sa mal prenos dát začat hneď, ako GPU vyprodukuje výstup, ale neustále sledovanie (polling) GPU zo strany CPU nie je efektivne.

TransferEngine to rieši pomocou UVM Watchers: umiestnení v jednotnej virtuálnej pamäti (UVM), ktoré môžu byť aktualizované jadrami GPU (dokonca aj v rámci grafov CUDA) a efektívne kontrolované (polled) zo strany CPU pomocou knižnice GDRCopy. Po alokácii sledovača UVM sa zaregistruje spätné volanie, ktoré sa vyvolá vždy, keď GPU aktualizuje túto hodnotu.

To umožňuje vzory, ako je prenos KvCache po vrstvách: po dokončení každej vrstvy transformátora GPU inkrementuje sledovač, čím spustí prenos KvCache danej vrstvy cez RDMA, zatiaľ čo nasledujúce vrstvy pokračujú vo výpočtoch.

TransferEngine poskytuje mechanizmy notifikácie dokončenia pre stranu iniciátora (volajúci funkcie submit_write) aj stranu príjemcu (cieľ operácií Write).

Notifikácia na strane iniciátora: Keď je odoslaná operácia zápisu, volajúci si môže zvoliť, že bude informovaný o dokončení buď prostredníctvom funkcie spätnej väzby (callback), alebo pomocou atomického príznaku. Táto notifikácia potvrďuje, že údaje boli úspešne prenesené zo zdrojového vyrovnávacieho pamäte.

Notifikácia na strane príjemcu: Kľúčovou inováciou je ImmCounter, komponent, ktorý sleduje počítadlá indexované 32-bitovými okamžitými číslami. Keď operácia zápisu nesie okamžitú hodnotu, počítadlo ImmCounter príjemcu sa po doručení údajov atomicky inkrementuje. Funkcia expect_imm_count umožňuje príjemcovi zaregistrovať spätné volanie, ktoré sa spustí, keď špecifická okamžitá hodnota dosiahne cieľový počet. Počítadlá možno transparentne synchronizovať s GPU via GDRCopy, kontrolovať ich priamo alebo ich monitorovať pomocou spätných volaní.

Tento mechanizmus na strane príjemcu poskytuje spoľahlivú synchronizáciu bez nutnosti usporiadania správ. Funguje rovnako dobre na sieťových radičoch ConnectX (ktoré zaručujú doručenie v správnom poradí cez transport RC) aj na sieťových radičoch EFA (ktoré doručujú správy v nesprávnom poradí cez transport SRD).

TransferEngine je implementovaný v jazyku Rust, pričom starostlivo optimalizuje alokácie, prácu s vláknami a synchronizáciu s cieľom minimalizovať latenciu a dosiahnuť vysokú priepustnosť. Motor spúšťa jedno pracovné vlákno na každú štruktúru DomainGroup, pričom každé vlákno je pripnuté na jadro CPU v uzle NUMA, ku ktorému sú pripojené zariadenia, čo minimalizuje latenciu plánovania aj prístupu do pamäte. Dátové štruktúry špecifické pre doménu sa alokujú až po pripnutí, aby sa zabezpečilo, že pamäť je rezervovaná v správnom uzle NUMA. Jedno pracovné vlákno spravuje až 4 Domainy (každá spravuje jeden sieťový radič), zatiaľ čo ďalšie vyhradené vlákno vykonáva polling GPU na aktualizáciu sledovačov UVM. Komunikácia medzi vláknami prebieha prostredníctvom frontov bez zámkov (lock-free queues).

API presmeruje požiadavky do štruktúry DomainGroup obsluhujúcej príslušné zariadenie, pričom prvá skupina obsluhuje aj hostiteľský systém. V tesnej slučke doménový pracovný proces vyhľadáva novú prácu, pričom uprednostňuje odosielanie nových požiadaviek. V závislosti od hardvéru a konfigurácie sa požiadavky rozdeľujú na shardy a záťaž sa vyvažuje naprieč dostupnými Domainmi. Prvý zápis zloženého požiadavku sa okamžite odošle do odosielacieho frontu sieťového radiča. Po vyčerpaní nových požiadaviek pracovný proces pokračuje v spracovaní prebiehajúcich operácií a odosiela zápisy na naplnenie hardvérovej rúry (pipeline). Následne sa kontrolujú fronty dokončenia, aby sa zistili dokončené prenosy a prírastky okamžitých počítadladiel. Udalosti sa agregujú s cieľom doručiť notifikácie pre jednotlivé prenosy a odovzdať prenos vyhradenému vláknu pre spätné volania, ktoré zdieľajú všetky skupiny.

Sharding v rámci DomainGroup je flexibilný. Prenosy môžu cieliť na konkrétne sieťové radiče podľa indexu. Jeden zápis je možné rozdeliť. Stránkované prenosy, operácie scatter a barrier, ktoré sa všetky prekladajú na viaceré zápisy, sa môžu rozdeliť (shardovať) na všetky sieťové radiče.

Domény (Domain) v rámci TransferEngine sú špecializované a optimalizované pre hardvér, ktorý majú pod kontrolou:

AWS EFA: Implementácia EFA používa knižnicu libfabric a spravuje jednu doménu fabric na jeden sieťový radič v rámci DomainGroup. Medzi kľúčové implementačné rozhodnutia a optimalizácie patria:

Zápisy s nulovou veľkosťou: EFA sa odchýli od špecifikácie RDMA tým, že vyžaduje platné cieľové deskriptory aj pre zápisy s nulovou veľkosťou. Preto TransferEngine vynucuje platné deskriptory pre všetky prenosy.

Tvorba šablón pre požiadavky na prácu (Work request templating): Predbežne vyplní a uchováva spoločné polia štruktúr libfabric pred odoslaním, čím znižuje režijné náklady na jednotlivé operácie pri hromadných prenosoch a operáciách skupín rovesníkov.

NVIDIA ConnectX-7: Implementácia ConnectX-7 používa knižnicu libibverbs s transportom RC. Medzi kľúčové implementačné rozhodnutia a optimalizácie patria:

Rady UD pre nadviazanie spojenia (handshakes): Každá Domain používa pár radov nespätného datagramu (UD queue pair) na výmenu nadviazaní spojenia RC.

Samostatné rady RC: Vytvára 2 páry radov RC na jedného rovesníka – jeden pre obojstranné operácie Send/Recv a druhý pre jednostranné operácie Write a WriteImm. Toto oddelenie je nevyhnutné, pretože dokončenia operácií Recv aj WriteImm spotrebúvajú požiadavky na prácu v poradí odoslania, čo umožňuje vysokú sémantiku Recv a zároveň podporuje WriteImm bez vzájomného rušenia.

Tvorba šablón pre požiadavky na prácu: Predbežne vypĺňa spoločné polia štruktúr požiadaviek na prácu pre hromadné prenosy a skupiny rovesníkov.

Zreťazenie požiadaviek na prácu (Work request chaining): Prepojí až 4 požiadavky na prácu prostredníctvom ukazovateľa next v štruktúre ibv_send_wr, čím sa zníži počet zvonení zvončeka (doorbell rings) do sieťového radiča.

Uvoľnené usporiadanie (Relaxed ordering): Povoľuje IBV_ACCESS_RELAXED_ORDERING na umožnenie transakcií PCIe v nesprávnom poradí medzi sieťovým radičom a pamäťou GPU.

Zameraním sa na prienik záruk naprieč rôznymi transportmi RDMA a zavedením synchronizácie nezávislej od poradia prostredníctvom komponentu ImmCounter poskytuje TransferEngine prenosný základ pre komunikáciu bod-bod v systémoch LLM. Knižnica je open-source a je dostupná na GitHub. Budeme radi, ak ju komunita preskúma a využije na budovanie nových komunikačných vzorov vo veľkorozmerných systémoch strojového učenia.

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

Motivácia

TransferEngine

Návrh API

Registrácia pamäte

Odosielanie/Prijímanie s obojstrannou komunikáciou (Two-Sided Send/Recv)

Jednostranný zápis (One-Sided Write)

Skupiny rovesníkov a hromadné operácie

UVM Watcher: Signalizácia z GPU na CPU

Notifikácia dokončenia

Implementácia

Hardvérovo špecifické optimalizácie

Záver

Citácia