Comunicació punt a punt RDMA per a sistemes LLM

Tradicionalment, els sistemes LLM s'han basat en comunicacions col·lectives per a patrons paral·lels estructurats en l'entrenament i la inferència. No obstant això, les noves càrregues de treball, com ara la inferència desagregada, l'encaminament de mescla d'experts (MoE) i l'afinament mitjançant aprenentatge per reforç asíncron, exigeixen la flexibilitat de la comunicació punt a punt.

El nostre article de recerca més recent presenta TransferEngine, una biblioteca de comunicació portàtil basada en RDMA que funciona en diferents controladors d'interfície de xarxa (NIC), inclosos NVIDIA ConnectX-7 i AWS Elastic Fabric Adapter (EFA). Demostrem la seva eficàcia en tres sistemes de producció a Perplexity AI:

Prefill i descodificació desagregats: transferència de KvCache capa per capa mitjançant escriptures paginades i observadors UVM, la qual cosa permet que les instàncies de prefill i descodificació s'escalin de manera independent.

Transferència de pesos per a la postformació d'aprenentatge per reforç en menys de 2 segons: sincronització de pesos extremadament ràpida entre nodes d'entrenament i d'inferència.

Habilitació de models de biliions de paràmetres a AWS EFA: operacions de combinació i distribució (dispatch) de MoE de baixa latència utilitzant primitives de dispersió (scatter) i barrera, assolint una latència d'avantguarda en ConnectX-7 i la primera implementació viable compatible amb AWS EFA.

Aquest article se centra en TransferEngine, la biblioteca de comunicació que impulsa aquests sistemes. Llegiu l'article de recerca complet a arXiv i exploreu el codi a GitHub.

Les comunicacions col·lectives excel·leixen en patrons paral·lels estructurats com el paral·lelisme de tensors i el paral·lelisme de dades, però inposen restriccions que limiten la seva aplicabilitat a les noves càrregues de treball:

Membresia fixa: les col·lectives requereixen establir un món de comunicació global en el moment de la inicialització, amb una configuració sincronitzada entre tots els participants. Per a la inferència desagregada, això és problemàtic, ja que el nombre d'instàncies de prefill i descodificació s'ha d'escalar dinàmicament en resposta als patrons de trànsit.

Formes de tensor uniformes: les col·lectives assumeixen que tots els participants intercanvien tensors amb formes i tipus de dades idèntics. En l'encaminament de MoE, el nombre de testimonis (tokens) enviats a cada expert depèn de les dades i es descobreix en temps d'execució, la qual cosa obliga les col·lectives tridireccionals denses (all-to-all) a farcir fins a les mides de memòria intermèdia (buffer) del pitjor escenari.

Ordre de les operacions: les API col·lectives apliquen semàntiques sincronitzades, fins i tot quan les aplicacions només necessiten saber que s'ha completat un conjunt d'operacions. Per exemple, en la transferència de KvCache, totes les pàgines de la instància de prefill a la instància de descodificació poden arribar en qualsevol ordre. De la mateixa manera, en les actualitzacions de pesos d'aprenentatge per reforç, tots els tensors de pesos es poden transferir de manera independent sense restriccions d'ordre.

La comunicació punt a punt basada en RDMA aborda naturalment aquestes limitacions, però tres barreres n'han impedit l'adopció:

Limitacions de la biblioteca: tot i que torch.distributed i NCCL proporcionen API punt a punt com Send/Recv, aquestes operacions continuen heretant semàntiques i limitacions col·lectives.

Diversitat de maquinari: els NIC NVIDIA ConnectX amb transport de connexió fiable (RC) són la configuració RDMA més comuna, la qual cosa garanteix el lliurament en ordre. Els proveïdors de núvol com AWS utilitzen un maquinari diferent; l'Elastic Fabric Adapter (EFA) amb transport de datagrames fiables escalables (SRD) lliura els missatges fora d'ordre. L'EFA també requereix agregar quatre NIC de 100 Gbps (o dos NIC de 200 Gbps) per assolir un ample de banda de 400 Gbps.

Portabilitat del programari: les biblioteques RDMA existents pateixen de tancament de proveïdor (vendor lock-in). DeepEP requereix RDMA iniciat per GPU (IBGDA) disponible només a ConnectX. NVSHMEM mostra una degradació greu del rendiment en EFA. Biblioteques com Mooncake i NIXL no tenen suport per a EFA o es troben en una fase de desenvolupament primerenca. Com a resultat, no existeix cap solució viable entre proveïdors per a la inferència de LLM.

TransferEngine aborda aquestes barreres a través de quatre capacitats bàsiques:

Assumpcions de transport comunes: en lloc de centrar-se en un NIC específic, TransferEngine només assumeix que el transport subjacent és fiable però no ordenat, és a dir, la intersecció de les garanties proporcionades pels transports RC i SRD.

Sincronització ImmCounter: una nova primitiva per a la notificació de finalització que funciona sense dependre de l'ordre dels missatges, la qual cosa permet que el mateix codi s'executi en transports ordenats (RC) i no ordenats (SRD).

Suport per a múltiples NIC: gestió transparent de múltiples NIC per GPU, fragmentant automàticament les transferències per maximitzar l'ample de banda.

Superfície d'API mínima: un petit conjunt d'operacions suficients per construir patrons de comunicació complexos.

Disseny del sistema TransferEngine, que mostra els camins de sol·licitud i callback, treballadors de CPU, buffers UVM, GPUs i NICs a través de dos dominis GPU/NUMA.

TransferEngine genera un fil de treball per GPU. Cada treballador gestiona un DomainGroup que coordina d'1 a 4 NIC RDMA, depenent de la plataforma de maquinari (un únic ConnectX-7 proporciona 400 Gbps, mentre que l'EFA requereix quatre adaptadors de 100 Gbps o dos adaptadors de 200 Gbps). Dins d'un DomainGroup, cada Domain és específic del maquinari i gestiona un únic NIC, encarregant-se de la gestió de parells de cues, l'enviament de treballs i el sondeig de finalització.

Cada instància de TransferEngine exposa una única adreça principal per a la identificació i el descobriment de pars. Els pars remots intercanvien estructures NetAddr que codifiquen adreces de xarxa i es poden serialitzar per a la comunicació fora de banda. Tots els pars d'un sistema han d'utilitzar el mateix nombre de NIC per GPU. Aquesta configuració uniformitzada permet a TransferEngine fragmentar les transferències de manera intel·ligent: en iniciar una transferència, tant l'origen com la destinació saben exactament quins NIC estan disponibles, la qual cosa permet un equilibri de càrrega automàtic en tots els adaptadors.

TransferEngine exposa una API mínima que abstrau la complexitat de les interfícies RDMA subjacents.

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

L'API consta de diversos components clau:

Les operacions RDMA requereixen que les regions de memòria estiguin registrades amb el NIC abans de la transferència de dades. La funció reg_mr pren un punter, una longitud i un dispositiu, i retorna dos controladors (handles):

MrHandle: un controlador local utilitzat com a origen per a les transferències

MrDesc: un descriptor serialitzable que es pot intercanviar amb pars remots per permetre'ls escriure a la memòria registrada

L'estructura MrDesc conté una llista d'adreces de xarxa remotes i una clau remota (una per NIC) que atorguen accés remot a la regió de memòria. Aquest disseny gestiona de manera natural les configuracions amb múltiples NIC.

Per a missatges petits i missatges de pla de control d'estil RPC, TransferEngine proporciona submit_send i submit_recvs:

submit_send: envia el missatge a un par remot identificat per NetAddr. El missatge es copia internament, de manera que qui crida pot reutilitzar o alliberar immediatament la memòria intermèdia.

submit_recvs: publica un conjunt de memòries intermèdies de recepció. Quan arriba un missatge, es crida a una funció de devolució de trucada (callback) amb les dades rebudes. La memòria intermèdia es retorna automàticament al conjunt un cop finalitzada la devolució de trucada.

Per a transferències de dades a granel, TransferEngine proporciona operacions d'escriptura unidireccional sense còpia (zero-copy):

submit_single_write: transfereix una regió contigua de l'origen a la destinació

submit_paged_writes: transfereix pàgines no contigües especificades per índexs, salts (strides) i desplaçaments (offsets)

Totes dues operacions accepten un valor immediat opcional de 32 bits que es lliura al receptor de manera atòmica en completar-se.

El motor fragmenta automàticament les transferències grans en múltiples NIC per maximitzar l'ample de banda.

Per a operacions adreçades a múltiples pars (comuns en l'encaminament de MoE), TransferEngine proporciona primitives optimitzades:

add_peer_group: preregistra un conjunt d'adreces de pars per a una cerca ràpida

submit_scatter: envia diferents talls de dades a cada par del grup

submit_barrier: envia una notificació només d'inmediat a tots els pars (transferència de zero bytes)

Aquestes operacions s'optimitzen per a una baixa latència emmagatzemant en memòria cau les estructures de dades internes i creant plantilles de sol·licituds de treball RDMA, evitant així la sobrecàrrega de configuració repetitiva.

Un repte clau en els sistemes LLM és coordinar les transferències RDMA amb el progrés de la GPU. Idealment, la transferència de dades hauria de començar tan bon punt la GPU produeix sortida, però fer un sondeig (polling) de la GPU des de la CPU és ineficient.

TransferEngine aborda això amb UVM Watchers (observadors UVM): ubicacions de memòria virtual unificada que poden ser actualitzades pels nuclis (kernels) de la GPU (fins i tot dins de grafs CUDA) i analitzades de manera eficient per la CPU utilitzant GDRCopy. Quan s'assigna un observador UVM, es registra una funció de devolució de trucada que s'invoca cada vegada que la GPU actualitza el valor.

Això permet patrons com la transferència de KvCache capa per capa: un cop completada cada capa del transformador, la GPU incrementa l'observador, la qual cosa desencadena una transferència RDMA de la KvCache d'aquesta capa mentre les capes posteriors continuen calculant.

TransferEngine proporciona mecanismes de notificació de finalització tant per al costat de l'iniciador (qui crida a submit_write) com per al costat del receptor (la destinació de les operacions Write).

Notificació al costat de l'iniciador: quan s'envia una operació d'escriptura, qui crida pot triar ser notificat en completar-se mitjançant una funció de devolució de trucada o un indicador atòmic. Aquesta notificació confirma que les dades s'han transferit correctament des de la memòria intermèdia d'origen.

Notificació al costat del receptor: la innovació clau és el ImmCounter, un component que fa un seguiment de comptadors indexats per números immediats de 32 bits. Quan una operació d'escriptura transporta un valor immediat, l'ImmCounter del receptor s'incrementa de manera atòmica en lliurar-se les dades. La funció expect_imm_count permet al receptor registrar una funció de devolució de trucada que s'activa quan un valor immediat específic arriba a un compte objectiu. Els comptadors es poden sincronitzar de manera transparent amb la GPU mitjançant GDRCopy, ser objecte de sondeig directe o ser monitoritzats mitjançant devolucions de trucada.

Aquest mecanisme al costat del receptor proporciona una sincronització fiable sense requerir l'ordenació dels missatges. Funciona igualment bé en NIC ConnectX (que garanteixen el lliurament en ordre mitjançant el transport RC) i en NIC EFA (que lliuren els missatges fora d'ordre mitjançant el transport SRD).

TransferEngine està implementat en Rust, optimitzant acuradament les assignacions, els fils i la sincronització per minimitzar la latència i assolir un alt rendiment. El motor genera un fil de treball per DomainGroup, cadascú fixat (pinned) a un nucli de CPU al node NUMA al qual estan connectats els dispositius, la qual cosa minimitza tant la latència de planificació com la d'accés a la memòria. Les estructures de dades específiques d'un domini s'assignen després de fixar el fil per assegurar que la memòria estigui reservada al node NUMA correcte. Un fil de treball gestiona fins a 4 Domains (cadascú gestionant un únic NIC), mentre que un altre fil dedicat fa un sondeig de la GPU per actualitzar els observadors UVM. La comunicació entre fils es realitza a través de cues sense bloqueig (lock-free).

L'API redirigeix les sol·licituds al DomainGroup que dona servei al dispositiu adequat, i el primer també dona servei a l'amfitrió (host). En un bucle cenyit, el treballador del domini busca nou treball, prioritzant l'enviament de noves sol·licituds. Depenent del maquinari i de la configuració, les sol·licituds es fragmenten i s'equilibren en termes de càrrega entre els Domains disponibles. La primera escriptura d'una sol·licitud composta es publica immediatament a la cua d'enviament del NIC. Unసారి que s'esgoten les noves sol·licituds, el treballador procedeix a fer avançar les operacions pendents, publicant escriptures per omplir la canonada (pipeline) del maquinari. Posteriorment, es fan sondeigs de les cues de finalització per consultar les transferències acabades i els increments de comptador immediats. Els esdeveniments s'agreguen per lliurar notificacions per transferència, delegant la transferència a un fil de devolució de trucada dedicat compartit per tots els grups.

La fragmentació dins d'un DomainGroup és flexible. Les transferències poden adreçar-se a NICs específics per índex. Una única escriptura es pot dividir. Les transferències paginades, la dispersió i les operacions de barrera, que es tradueixen totes en múltiples escriptures, es poden fragmentar en tots els NICs.

Els Domains de TransferEngine estan especialitzats i optimitzats per al maquinari sota el seu control:

AWS EFA: la implementació d'EFA utilitza libfabric, gestionant un domini de fabric per NIC dins del DomainGroup. Les decisions clau d'implementació i les optimitzacions inclouen:

Escriptures de mida zero: l'EFA divergeix de l'especificació RDMA en exigir descriptors de destinació vàlids fins i tot per a escriptures de mida zero. Per tant, TransferEngine aplica descriptors vàlids per a totes les transferències.

Plantilles de sol·licitud de treball: preemplena i reté els camps comuns de les estructures de libfabric abans de publicar-les, la qual cosa redueix la sobrecàrrega per operació per a transferències a granel i operacions de grups de pars

NVIDIA ConnectX-7: la implementació de ConnectX-7 utilitza libibverbs amb transport RC. Les decisions clau d'implementació i les optimitzacions inclouen:

Parells de cues UD per a establiments de connexió (handshakes): cada Domain utilitza un parell de cues de datagrames no fiables (UD) per intercanviar establiments de connexió RC.

Parells de cues RC separats: crea 2 parells de cues RC per par, un per a operacions Send/Recv bidireccionals i un altre per a operacions Write i WriteImm unidireccionals. Aquesta separació és necessària perquè tant les finalitzacions de Recv com de WriteImm consumeixen sol·licituds de treball en l'ordre de publicació, la qual cosa permet una semàntica de Recv d'alt nivell alhora que admet WriteImm sense interferències.

Plantilles de sol·licitud de treball: preemplena els camps comuns de les estructures de sol·licitud de treball per a transferències a granel i grups de pars.

Encadenament de sol·licituds de treball: enllaça fins a 4 sol·licituds de treball a través del punter next d'ibv_send_wr, la qual cosa redueix el nombre de trucades de timbre (doorbell rings) al NIC.

Ordre relaxat: habilita IBV_ACCESS_RELAXED_ORDERING per permetre transaccions PCIe fora d'ordre entre el NIC i la memòria de la GPU.

En centrar-se en la intersecció de les garanties entre diferents transports RDMA i introduir la sincronització independent de l'ordre mitjançant ImmCounter, TransferEngine proporciona una base portàtil per a la comunicació punt a punt en sistemes LLM. La biblioteca és de codi obert i està disponible a GitHub, i animem la comunitat a explorar-la i construir a partir d'ella per a nous patrons de comunicació en sistemes d'aprenentatge automàtic a gran escala.

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

Motivació

TransferEngine

Disseny d'API

Registre de memòria

Enviament/Recepció bidireccional (Two-Sided Send/Recv)

Escriptura unidireccional (One-Sided Write)

Grups de pars i operacions a granel

Observador UVM: senyalització de GPU a CPU

Notificació de finalització

Implementació

Optimitzacions específiques de maquinari

Conclusió

Citació