Comunicare punct la punct RDMA pentru sisteme LLM

Sistemele LLM s-au bazat în mod tradițional pe comunicații colective pentru tipare paralele structurate în antrenament și inferență. Cu toate acestea, sarcinile de lucru emergente, cum ar fi inferența dezagregată, rutarea Mixture-of-Experts (MoE) și reglarea fină asincronă prin învățare prin consolidare, necesită flexibilitatea comunicării punct la punct.

Cea mai recentă lucrare a noastră de cercetare introduce TransferEngine, o bibliotecă de comunicare portabilă bazată pe RDMA, care funcționează pe diferite controlere de interfață de rețea (NIC), inclusiv NVIDIA ConnectX-7 și AWS Elastic Fabric Adapter (EFA). Demonstrăm eficacitatea acesteia în trei sisteme de producție de la Perplexity AI:

Prefill și Decode dezagregate: Transfer KvCache strat cu strat utilizând scrieri paginate și observatori UVM, permițând instanțelor de prefill și decode să se scaleze independent.

Transfer de ponderi pentru post-antrenarea RL în sub 2 secunde: Sincronizare extrem de rapidă a ponderilor între nodurile de antrenament și cele de inferență.

Activarea modelelor cu milioane de miliarde de parametri pe AWS EFA: Operațiuni de expediere și combinare MoE cu latență redusă utilizând primitive de tip scatter și barieră, obținând o latență de ultimă generație pe ConnectX-7 și prima implementare viabilă compatibilă cu AWS EFA.

Acest articol se concentrează pe TransferEngine în sine, biblioteca de comunicare care alimentează aceste sisteme. Citiți lucrarea completă de cercetare pe arXiv și explorați codul pe GitHub.

Comunicațiile colective excelează în tipare paralele structurate, cum ar fi Paralelismul Tenzorial și Paralelismul de Date, dar impun constrângeri care le limitează aplicabilitatea la sarcinile de lucru emergente:

Membri ficși: Colectivele, la momentul inițializării, necesită stabilirea unei lumi de comunicare globale, cu o configurare sincronizată între toți participanții. Pentru inferența dezagregată, acest lucru este problematic deoarece numărul de instanțe de prefill și decode trebuie să se scaleze dinamic ca răspuns la tiparele de trafic.

Forme de tensori uniforme: Colectivele presupun că toți participanții schimbă tensori cu forme și tipuri de date identice. În rutarea MoE, numărul de tokenuri trimise către fiecare expert depinde de date și este descoperit în timpul executării (runtime), forțând colectivele dense all-to-all să facă padding la dimensiunile de tampon pentru cel mai defavorabil caz.

Ordonarea operațiunilor: API-urile colective impun semantică sincronizată, chiar și atunci când aplicațiile trebuie doar să știe că un set de operațiuni s-a finalizat. De exemplu, în transferul KvCache, toate paginile de la instanța de prefill la instanța de decode pot sosi în orice ordine. În mod similar, în actualizările de ponderi RL, toți tensorii de ponderi pot fi transferați independent, fără constrângeri de ordonare.

Comunicarea punct la punct bazată pe RDMA abordează în mod natural aceste limitări, dar trei bariere i-au împiedicat adoptarea:

Limitări ale bibliotecilor: În timp ce torch.distributed și NCCL oferă API-uri punct la punct precum Send/Recv, aceste operațiuni moștenesc în continuare semantică și limitări colective.

Diversitate hardware: NIC-urile NVIDIA ConnectX cu transport Reliable Connection (RC) sunt cea mai comună configurație RDMA, garantând livrarea în ordine. Furnizorii de servicii cloud, cum ar fi AWS, utilizează hardware diferit; Elastic Fabric Adapter (EFA) cu transport Scalable Reliable Datagram (SRD) livrează mesaje în dezordine. De asemenea, EFA necesită agregarea a patru NIC-uri de 100 Gbps (sau două NIC-uri de 200 Gbps) pentru a obține o lățime de bandă de 400 Gbps.

Portabilitate software: Bibliotecile RDMA existente suferă de blocarea furnizorului (vendor lock-in). DeepEP necesită RDMA inițiat de GPU (IBGDA) disponibil doar pe ConnectX. NVSHMEM prezintă o degradare severă a performanței pe EFA. Biblioteci precum Mooncake și NIXL nu au suport pentru EFA sau rămân în stadiu incipient de dezvoltare. Ca urmare, nu există nicio soluție viabilă cross-provider pentru inferența LLM.

TransferEngine abordează aceste bariere prin patru capabilități principale:

Ipoteze de transport comune: În loc să vizeze un anumit NIC, TransferEngine presupune doar că transportul subiacent este fiabil, dar neordonat, adică intersecția garanțiilor furnizate atât de transporturile RC, cât și de cele SRD.

Sincronizare ImmCounter: O primitivă inovatoare pentru notificarea finalizării care funcționează fără a se baza pe ordonarea mesajelor, permițând ca același cod să ruleze atât pe transporturi ordonate (RC), cât și neordonate (SRD).

Suport pentru mai multe NIC-uri: Gestionarea transparentă a mai multor NIC-uri per GPU, împărțind automat transferurile în fragmente (sharding) pentru a maximiza lățimea de bandă.

Suprafață API minimală: Un set restrâns de operațiuni suficiente pentru a construi tiparecomplexe de comunicare.

Designul sistemului TransferEngine, ilustrând traseele de solicitare și callback, workerii CPU, buffer-ele UVM, GPU-urile și NIC-urile pe două domenii GPU/NUMA.

TransferEngine generează un fir de execuție worker per GPU. Fiecare worker gestionează un DomainGroup care coordonează între 1 și 4 NIC-uri RDMA, în funcție de platforma hardware (un singur ConnectX-7 oferă 400 Gbps, în timp ce EFA necesită patru adaptoare de 100 Gbps sau două de 200 Gbps). În cadrul unui DomainGroup, fiecare Domain este specific hardware-ului și gestionează un singur NIC, ocupându-se de gestionarea perechilor de cozi, trimiterea lucrărilor și interogarea finalizărilor.

Fiecare instanță TransferEngine expune o singură adresă principală pentru identificarea și descoperirea egalilor (peers). Egalii de la distanță schimbă structuri NetAddr care codifică adresele de rețea și pot fi serializate pentru comunicarea în bandă laterală (out-of-band). Toți egalii dintr-un sistem trebuie să utilizeze același număr de NIC-uri per GPU. Această configurație uniformă permite TransferEngine să fragmenteze inteligent transferurile: la inițierea unui transfer, atât sursa, cât și destinația știu exact ce NIC-uri sunt disponibile, permițând echilibrarea automată a încărcării pe toate adaptoarele.

TransferEngine expune un API minimal care abstractizează complexitatea interfețelor RDMA subiacente.

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-ul constă din mai multe componente cheie:

Operațiunile RDMA necesită ca regiunile de memorie să fie înregistrate la NIC înainte de transferul de date. Funcția reg_mr ia ca argumente un pointer, o lungime și un dispozitiv, returnând două handlere:

MrHandle: Un handler local utilizat ca sursă pentru transferuri

MrDesc: Un descriptor serializabil care poate fi schimbat cu egali de la distanță pentru a le permite să scrie în memoria înregistrată

Structura MrDesc conține o listă de adrese de rețea de la distanță și chei de la distanță (câte una per NIC) care acordă acces de la distanță la regiunea de memorie. Acest design gestionează în mod natural configurațiile cu mai multe NIC-uri.

Pentru mesaje mici și mesaje de plan de control de tip RPC, TransferEngine oferă submit_send și submit_recvs:

submit_send: Trimite mesajul către un egal de la distanță identificat prin NetAddr. Mesajul este copiat intern, astfel încât apelantul poate reutiliza sau elibera imediat tamponul.

submit_recvs: Postează un grup (pool) de tampoane de recepție. Când soseste un mesaj, este apelat un callback cu datele primite. Tamponul este returnat automat în grup după finalizarea callback-ului.

Pentru transferuri de date în masă, TransferEngine oferă operațiuni de scriere unidirecțională fără copiere (zero-copy):

submit_single_write: Transferă o regiune contiguă de la sursă la destinație

submit_paged_writes: Transferă pagini necontigue specificate prin indici, pași (strides) și deplasări (offsets)

Ambele operațiuni acceptă o valoare imediată opțională pe 32 de biți care este livrată receptorului în mod atomic la finalizare.

Motorul fragmentează automat transferurile mari pe mai multe NIC-uri pentru a maximiza lățimea de bandă.

Pentru operațiuni care vizează mai mulți egali (comune în rutarea MoE), TransferEngine oferă primitive optimizate:

add_peer_group: Pre-înregistrează un set de adrese ale egalilor pentru căutare rapidă

submit_scatter: Trimite diferite segmente de date către fiecare egal din grup

submit_barrier: Trimite o notificare exclusiv imediată către toți egalii (transfer de zero octeți)

Aceste operațiuni sunt optimizate pentru latență redusă prin stocarea în cache a structurilor de date interne și șablonoane (templating) pentru cererile de lucru RDMA, evitând suprasolicitarea repetitivă de configurare.

O provocare cheie în sistemele LLM este coordonarea transferurilor RDMA cu progresul GPU. În mod ideal, transferul de date ar trebui să înceapă imediat ce GPU-ul produce un rezultat, dar interogarea (polling) GPU-ului de la CPU este ineficientă.

TransferEngine abordează acest lucru prin Observatori UVM (UVM Watchers): locații de memorie virtuală unificată care pot fi actualizate de kernel-urile GPU (chiar și în cadrul graficelor CUDA) și interogate eficient de CPU utilizând GDRCopy. Când este alocat un observator UVM, este înregistrat un callback care este invocat ori de câte ori GPU-ul actualizează valoarea.

Acest lucru permite tipare precum transferul KvCache strat cu strat: după finalizarea fiecărui strat de transformer, GPU-ul incrementează observatorul, declanșând un transfer RDMA al KvCache-ului acelui strat în timp ce straturile următoare își continuă calculul.

TransferEngine oferă mecanisme de notificare a finalizării atât pentru partea inițiatoare (apelantul lui submit_write), cât și pentru partea receptoare (ținta operațiunilor de scriere).

Notificare pe partea inițiatoare: Când este trimisă o operațiune de scriere, apelantul poate alege să fie notificat la finalizare fie prin intermediul unei funcții callback, fie printr-un semnalizator (flag) atomic. Această notificare confirmă că datele au fost transferate cu succes din tamponul sursă.

Notificare pe partea receptoare: Inovația cheie este ImmCounter, o componentă care urmărește contoare indexate prin numere imediate pe 32 de biți. Când o operațiune de scriere poartă o valoare imediată, ImmCounter-ul receptorului este incrementat atomic la livrarea datelor. Funcția expect_imm_count permite receptorului să înregistreze un callback care se declanșează atunci când o anumită valoare imediată atinge un număr țintă. Contoarele pot fi sincronizate în mod transparent cu GPU-ul prin intermediul GDRCopy, interogate direct sau monitorizate prin callback-uri.

Acest mecanism pe partea receptoare oferă o sincronizare fiabilă fără a necesita ordonarea mesajelor. Funcționează la fel de bine pe NIC-urile ConnectX (care garantează livrarea în ordine prin transportul RC) și pe NIC-urile EFA (care livrează mesaje în dezordine prin transportul SRD).

TransferEngine este implementat în Rust, optimizând cu atenție alocările, firele de execuție și sincronizarea pentru a minimiza latența și a obține un debit ridicat. Motorul generează câte un fir de execuție worker per DomainGroup, fiecare fixat (pinned) pe un nucleu de CPU de pe nodul NUMA la care sunt atașate dispozitivele, minimizând atât programarea (scheduling), cât și latența accesului la memorie. Structurile de date specifice unui domeniu sunt alocate după fixare pentru a se asigura că memoria este rezervată în nodul NUMA corect. Un fir de execuție worker gestionează până la 4 Domain-uri (fiecare gestionând un singur NIC), în timp ce un alt fir dedicat interoghează GPU-ul pentru a actualiza observatorii UVM. Comunicarea între firele de execuție se face prin cozi fără blocare (lock-free).

API-ul redirecționează cererile către DomainGroup care deservește dispozitivul corespunzător, primul deservind și gazda. Într-o buclă strânsă (tight loop), workerul de domeniu interoghează pentru lucrări noi, acordând prioritate trimiterii de cereri noi. În funcție de hardware și configurație, cererile sunt fragmentate și echilibrate ca sarcină pe Domain-urile disponibile. Prima scriere a unei cereri compuse este postată imediat în coada de trimitere a NIC-ului. Odată epuizate noile cereri, workerul trece la procesarea operațiunilor în așteptare, postând scrieri pentru a umple conducta hardware (pipeline). Ulterior, cozile de finalizare sunt interogate pentru a afla transferurile finalizate și incrementările de contoare imediate. Evenimentele sunt agregate pentru a livra notificări per transfer, predând transferul unui fir de execuție callback dedicat, partajat de toate grupurile.

Fragmentarea (sharding) în cadrul unui DomainGroup este flexibilă. Transferurile pot viza anumite NIC-uri după index. O singură scriere poate fi împărțită. Transferurile paginate, operațiunile scatter și barrier, care se traduc toate în scrieri multiple, se pot fragmenta pe toate NIC-urile.

Domain-urile din cadrul TransferEngine sunt specializate și optimizate pentru hardware-ul aflat sub controlul lor:

AWS EFA: Implementarea EFA utilizează libfabric, gestionând un domeniu fabric per NIC în cadrul DomainGroup. Deciziile cheie de implementare și optimizările includ:

Scrieri de dimensiune zero: EFA se abate de la specificația RDMA prin necesitatea unor descriptori țintă validați chiar și pentru scrieri de dimensiune zero. Prin urmare, TransferEngine impune descriptori valizi pentru toate transferurile.

Șablonoane pentru cereri de lucru (Work request templating): Pre-populează și păstrează câmpurile comune ale structurilor libfabric înainte de postare, reducând suprasolicitarea per operațiune pentru transferurile în masă și operațiunile pe grupuri de egali

NVIDIA ConnectX-7: Implementarea ConnectX-7 utilizează libibverbs cu transport RC. Deciziile cheie de implementare și optimizările includ:

Perechi de cozi UD pentru handshake-uri: Fiecare Domain utilizează o pereche de cozi Unreliable Datagram (UD) pentru a face schimb de handshake-uri de conexiune RC.

Perechi de cozi RC separate: Creează 2 perechi de cozi RC per egal, una pentru operațiuni Send/Recv bidirecționale și alta pentru operațiuni Write și WriteImm unidirecționale. Această separare este necesară deoarece atât finalizările Recv, cât și cele WriteImm consumă cereri de lucru în ordinea postării, permițând semantica de nivel înalt Recv în timp ce acceptă WriteImm fără interferențe.

Șablonoane pentru cereri de lucru (Work request templating): Pre-populează câmpurile comune ale structurilor de cereri de lucru pentru transferuri în masă și grupuri de egali.

Înlănțuirea cererilor de lucru (Work request chaining): Conectează până la 4 cereri de lucru prin pointerul next al ibv_send_wr, reducând numărul de apeluri doorbell către NIC.

Ordonare relaxată: Activează IBV_ACCESS_RELAXED_ORDERING pentru a permite tranzacții PCIe dezordonate (out-of-order) între NIC și memoria GPU.

Vizând intersecția garanțiilor între diferite transporturi RDMA și introducând sincronizarea independentă de ordonare prin ImmCounter, TransferEngine oferă o bază portabilă pentru comunicarea punct la punct în sistemele LLM. Biblioteca este open-source și disponibilă pe GitHub, iar noi încurajăm comunitatea să o exploreze și să construiască pe baza ei pentru modelele emergente de comunicare în sisteme de învățare automată la scară largă.

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

Motivație

TransferEngine

Design API

Înregistrare memorie

Send/Recv bidirecțional (Two-Sided)

Scriere unidirecțională (One-Sided Write)

Grupuri de egali și operațiuni în masă

Observator UVM: Semnalizare GPU-CPU

Notificare de finalizare

Implementare

Optimizări specifice hardware-ului

Concluzie

Citare