Comunicazione RDMA punto a punto per sistemi LLM

I sistemi LLM si sono tradizionalmente affidati alle comunicazioni collettive per schemi paralleli strutturati durante la formazione e l'inferenza. Tuttavia, i carichi di lavoro emergenti, come l'inferenza disaggregata, il routing Mixture-of-Experts (MoE) e il fine-tuning tramite apprendimento per rinforzo asincrono, richiedono la flessibilità della comunicazione punto a punto.

Il nostro ultimo articolo di ricerca introduce TransferEngine, una libreria di comunicazione basata su RDMA portabile che funziona con diversi Network Interface Controller (NIC), inclusi NVIDIA ConnectX-7 e AWS Elastic Fabric Adapter (EFA). Ne dimostriamo l'efficacia in tre sistemi di produzione presso Perplexity AI:

Prefill e decodifica disaggregati: Trasferimento della KvCache strato per strato utilizzando scritture paginate e UVM watcher, che consente alle istanze di prefill e decodifica di scalare indipendentemente.

Trasferimento dei pesi per il post-addestramento RL in meno di 2 secondi: Sincronizzazione dei pesi incredibilmente veloce tra i nodi di addestramento e di inferenza.

Abilitazione di modelli con trilioni di parametri su AWS EFA: Operazioni di dispatch e combinazione MoE a bassa latenza che utilizzano primitive scatter e barrier, raggiungendo una latenza allo stato dell'arte su ConnectX-7 e la prima implementazione valida compatibile con AWS EFA.

Questo articolo si concentra su TransferEngine, la libreria di comunicazione che alimenta questi sistemi. Leggi l'intero documento di ricerca su arXiv ed esplora il codice su GitHub.

Le comunicazioni collettive eccellono in schemi paralleli strutturati come il parallelismo dei tensori e il parallelismo dei dati, ma impongono vincoli che ne limitano l'applicabilità ai carichi di lavoro emergenti:

Appartenenza fissa: Le collettive richiedono la creazione di un mondo di comunicazione globale al momento dell'inizializzazione, con una configurazione sincronizzata tra tutti i partecipanti. Per l'inferenza disaggregata, questo è problematico poiché il numero di istanze di prefill e decodifica deve scalare dinamicamente in risposta ai modelli di traffico.

Forme tensoriali uniformi: Le collettive presuppongono che tutti i partecipanti scambino tensori con forme e tipi di dati identici. Nel routing MoE, il numero di token inviati a ciascun esperto dipende dai dati e viene rilevato in fase di esecuzione, costringendo le collettive all-to-all dense a eseguire il padding alle dimensioni del buffer peggiori.

Ordinamento delle operazioni: Le API collettive impongono semantiche sincronizzate, anche quando le applicazioni devono solo sapere che una serie di operazioni è stata completata. Ad esempio, nel trasferimento KvCache, tutte le pagine dall'istanza di prefill all'istanza di decodifica possono arrivare in qualsiasi ordine. Allo stesso modo, negli aggiornamenti dei pesi RL, tutti i tensori dei pesi possono essere trasferiti indipendentemente senza vincoli di ordinamento.

La comunicazione punto a punto basata su RDMA affronta naturalmente queste limitazioni, ma tre barriere ne hanno impedito l'adozione:

Limitazioni della libreria: Mentre torch.distributed e NCCL forniscono API punto a punto come Send/Recv, queste operazioni ereditano ancora semantiche e limitazioni collettive.

Diversità dell'hardware: I NIC NVIDIA ConnectX con trasporto Reliable Connection (RC) sono la configurazione RDMA più comune, garantendo la consegna in ordine. I provider cloud come AWS utilizzano hardware diverso; l'Elastic Fabric Adapter (EFA) con trasporto Scalable Reliable Datagram (SRD) consegna i messaggi fuori ordine. EFA richiede inoltre l'aggregazione di quattro NIC da 100 Gbps (o due da 200 Gbps) per ottenere una larghezza di banda di 400 Gbps.

Portabilità del software: Le librerie RDMA esistenti soffrono di vendor lock-in. DeepEP richiede RDMA avviato da GPU (IBGDA) disponibile solo su ConnectX. NVSHMEM mostra un grave degrado delle prestazioni su EFA. Librerie come Mooncake e NIXL mancano di supporto EFA o rimangono in fase di sviluppo iniziale. Di conseguenza, non esiste una soluzione valida tra i diversi provider per l'inferenza LLM.

TransferEngine affronta queste barriere attraverso quattro capacità fondamentali:

Ipotesi di trasporto comuni: Piuttosto che puntare a un NIC specifico, TransferEngine presuppone solo che il trasporto sottostante sia affidabile ma non ordinato, ovvero l'intersezione delle garanzie fornite dai trasporti RC e SRD.

Sincronizzazione ImmCounter: Una nuova primitiva per la notifica di completamento che funziona senza fare affidamento sull'ordinamento dei messaggi, consentendo l'esecuzione dello stesso codice su trasporti sia ordinati (RC) che non ordinati (SRD).

Supporto multi-NIC: Gestione trasparente di più NIC per GPU, con sharding automatico dei trasferimenti per massimizzare la larghezza di banda.

Superficie API minima: Un piccolo set di operazioni sufficiente per costruire complessi schemi di comunicazione.

Progettazione del sistema TransferEngine, che mostra i percorsi delle richieste e dei callback, i worker CPU, i buffer UVM, le GPU e le NIC in due domini GPU/NUMA.

TransferEngine genera un thread di lavoro per GPU. Ogni lavoratore gestisce un DomainGroup che coordina da 1 a 4 NIC RDMA, a seconda della piattaforma hardware (un singolo ConnectX-7 fornisce 400 Gbps, mentre EFA richiede quattro adattatori da 100 Gbps o due da 200 Gbps). All'interno di un DomainGroup, ogni Domain è specifico per l'hardware e gestisce un singolo NIC, occupandosi della gestione delle coppie di code, dell'invio del lavoro e del polling di completamento.

Ogni istanza di TransferEngine espone un singolo indirizzo principale per l'identificazione e la scoperta dei peer. I peer remoti scambiano strutture NetAddr che codificano gli indirizzi di rete e possono essere serializzate per la comunicazione fuori banda. Tutti i peer in un sistema devono utilizzare lo stesso numero di NIC per GPU. Questa configurazione uniforme consente a TransferEngine di eseguire lo shard dei trasferimenti in modo intelligente: quando si avvia un trasferimento, sia l'origine che la destinazione sanno esattamente quali NIC sono disponibili, consentendo il bilanciamento automatico del carico su tutti gli adattatori.

TransferEngine espone un'API minima che astrae la complessità delle interfacce RDMA sottostanti.

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 consiste di diversi componenti chiave:

Le operazioni RDMA richiedono che le regioni di memoria siano registrate con il NIC prima del trasferimento dei dati. La funzione reg_mr accetta un puntatore, una lunghezza e un dispositivo, restituendo due handle:

MrHandle: Un handle locale utilizzato come origine per i trasferimenti

MrDesc: Un descrittore serializzabile che può essere scambiato con i peer remoti per consentire loro di scrivere nella memoria registrata

La struttura MrDesc contiene un elenco di indirizzi di rete remoti e una chiave remota (una per NIC) che concedono l'accesso remoto alla regione di memoria. Questa progettazione gestisce naturalmente le configurazioni multi-NIC.

Per messaggi piccoli e messaggi di piano di controllo in stile RPC, TransferEngine fornisce submit_send e submit_recvs:

submit_send: Invia il messaggio a un peer remoto identificato da NetAddr. Il messaggio viene copiato internamente, quindi il chiamante può riutilizzare o rilasciare immediatamente il buffer.

submit_recvs: Pubblica un pool di buffer di ricezione. Quando arriva un messaggio, viene invocato un callback con i dati ricevuti. Il buffer viene restituito automaticamente al pool dopo il completamento del callback.

Per trasferimenti di dati in blocco, TransferEngine fornisce operazioni di scrittura a una via senza copia (zero-copy):

submit_single_write: Trasferisce una regione contigua dall'origine alla destinazione

submit_paged_writes: Trasferisce pagine non contigue specificate da indici, passi e offset

Entrambe le operazioni accettano un valore immediato a 32 bit facoltativo che viene consegnato al ricevitore atomicamente al completamento.

Il motore esegue automaticamente lo shard di grandi trasferimenti su più NIC per massimizzare la larghezza di banda.

Per le operazioni destinate a più peer (comuni nel routing MoE), TransferEngine fornisce primitive ottimizzate:

add_peer_group: Pre-registra un set di indirizzi peer per una ricerca rapida

submit_scatter: Invia diverse porzioni di dati a ciascun peer nel gruppo

submit_barrier: Invia una notifica solo immediata a tutti i peer (trasferimento a zero byte)

Queste operazioni sono ottimizzate per una bassa latenza memorizzando nella cache le strutture dati interne e creando modelli per le richieste di lavoro RDMA, evitando il sovraccarico della configurazione ripetitiva.

Una sfida chiave nei sistemi LLM è il coordinamento dei trasferimenti RDMA con l'avanzamento della GPU. Idealmente, il trasferimento dei dati dovrebbe iniziare non appena la GPU produce un output, ma eseguire il polling della GPU dalla CPU è inefficiente.

TransferEngine affronta questo problema con UVM Watcher: posizioni di memoria virtuale unificate che possono essere aggiornate dai kernel GPU (anche all'interno dei grafici CUDA) ed efficientemente interrogate dalla CPU utilizzando GDRCopy. Quando viene allocato un UVM watcher, viene registrato un callback che viene invocato ogni volta che la GPU aggiorna il valore.

Ciò abilita schemi come il trasferimento KvCache strato per strato: dopo il completamento di ogni strato transformer, la GPU incrementa il watcher, attivando un trasferimento RDMA della KvCache di quello strato mentre gli strati successivi continuano a calcolare.

TransferEngine fornisce meccanismi di notifica di completamento sia per il lato iniziatore (il chiamante di submit_write) che per il lato ricevitore (la destinazione delle operazioni di Scrittura).

Notifica lato iniziatore: Quando viene inviata un'operazione di scrittura, il chiamante può scegliere di essere notificato al completamento tramite una funzione di callback o un flag atomico. Questa notifica conferma che i dati sono stati trasferiti con successo dal buffer di origine.

Notifica lato ricevitore: L'innovazione chiave è l'ImmCounter, un componente che tiene traccia dei contatori indicizzati da numeri immediati a 32 bit. Quando un'operazione di scrittura porta un valore immediato, l'ImmCounter del ricevitore viene incrementato atomicamente alla consegna dei dati. La funzione expect_imm_count consente al ricevitore di registrare un callback che si attiva quando uno specifico valore immediato raggiunge un conteggio target. I contatori possono essere sincronizzati in modo trasparente con la GPU tramite GDRCopy, interrogati direttamente o monitorati tramite callback.

Questo meccanismo lato ricevitore fornisce una sincronizzazione affidabile senza richiedere l'ordinamento dei messaggi. Funziona ugualmente bene su NIC ConnectX (che garantiscono la consegna in ordine su trasporto RC) e NIC EFA (che consegnano i messaggi fuori ordine su trasporto SRD).

TransferEngine è implementato in Rust, ottimizzando attentamente le allocazioni, il threading e la sincronizzazione per ridurre al minimo la latenza e ottenere un throughput elevato. Il motore genera un thread di lavoro per DomainGroup, ciascuno collegato a un core della CPU sul nodo NUMA a cui sono collegati i dispositivi, riducendo al minimo sia la pianificazione che la latenza di accesso alla memoria. Le strutture dati specifiche di un dominio vengono allocate dopo il pinning per garantire che la memoria sia riservata nel nodo NUMA corretto. Un thread di lavoro gestisce fino a 4 Domain (ciascuno gestisce un singolo NIC), mentre un altro thread dedicato interroga la GPU per aggiornare gli UVM watcher. La comunicazione tra thread avviene tramite code senza lock (lock-free).

L'API inoltra le richieste al DomainGroup che serve il dispositivo appropriato, con il primo che serve anche l'host. In un ciclo serrato, il lavoratore di dominio interroga per trovare nuovo lavoro, dando priorità all'invio di nuove richieste. A seconda dell'hardware e della configurazione, le richieste vengono suddivise (sharded) e bilanciate tra i Domain disponibili. La prima Scrittura di una richiesta composta viene immediatamente pubblicata nella coda di invio del NIC. Una volta esaurite le nuove richieste, il lavoratore procede all'avanzamento delle operazioni in sospeso, pubblicando scritture per riempire la pipeline hardware. Successivamente, le code di completamento vengono interrogate per verificare i trasferimenti terminati e gli incrementi del contatore immediato. Gli eventi vengono aggregati per fornire notifiche per trasferimento, consegnando il trasferimento a un thread di callback dedicato condiviso da tutti i gruppi.

Lo sharding all'interno di un DomainGroup è flessibile. I trasferimenti possono puntare a NIC specifici per indice. Una singola Scrittura può essere suddivisa. I trasferimenti paginati, lo scatter e le operazioni barrier, che si traducono tutti in Scritture multiple, possono essere suddivisi su tutti i NIC.

I Domain all'interno di TransferEngine sono specializzati e ottimizzati per l'hardware sotto il loro controllo:

AWS EFA: L'implementazione EFA utilizza libfabric, gestendo un dominio di fabric per NIC all'interno del DomainGroup. Le decisioni e le ottimizzazioni chiave includono:

Scritture a dimensione zero: EFA diverge dalla specifica RDMA richiedendo descrittori target validi anche per scritture a dimensione zero. Pertanto, TransferEngine impone descrittori validi per tutti i trasferimenti.

Creazione di modelli per richieste di lavoro: Pre-popola e mantiene i campi comuni delle strutture libfabric prima della pubblicazione, riducendo il sovraccarico per operazione per trasferimenti in blocco e operazioni di gruppo di peer

NVIDIA ConnectX-7: L'implementazione ConnectX-7 utilizza libibverbs con trasporto RC. Le decisioni e le ottimizzazioni chiave includono:

Coppie di code UD per handshake: Ogni Domain utilizza una coppia di code Unreliable Datagram (UD) per scambiare handshake di connessione RC.

Coppie di code RC separate: Crea 2 coppie di code RC per peer, una per operazioni di Invio/Ricezione a due vie e un'altra per operazioni di Scrittura e WriteImm a una via. Questa separazione è necessaria perché sia i completamenti Recv che WriteImm consumano richieste di lavoro nell'ordine di pubblicazione, consentendo semantiche Recv di alto livello pur supportando WriteImm senza interferenze.

Creazione di modelli per richieste di lavoro: Pre-popola i campi comuni delle strutture di richiesta di lavoro per trasferimenti in blocco e gruppi di peer.

Concatenazione di richieste di lavoro: Collega fino a 4 richieste di lavoro tramite il puntatore next di ibv_send_wr, riducendo il numero di squilli del campanello (doorbell rings) al NIC.

Ordinamento rilassato: Abilita IBV_ACCESS_RELAXED_ORDERING per consentire transazioni PCIe fuori ordine tra la memoria NIC e GPU.

Mirando all'intersezione delle garanzie tra diversi trasporti RDMA e introducendo la sincronizzazione indipendente dall'ordine tramite ImmCounter, TransferEngine fornisce una base portabile per la comunicazione punto a punto nei sistemi LLM. La libreria è open source e disponibile su GitHub, e incoraggiamo la comunità a esplorarla e costruirci sopra per i modelli di comunicazione emergenti nei sistemi di apprendimento automatico su larga scala.

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

Motivazione

TransferEngine

Progettazione API

Registrazione della memoria

Invio/Ricezione a due vie

Scrittura a una via

Gruppi di peer e operazioni in blocco

UVM Watcher: Segnalazione da GPU a CPU

Notifica di completamento

Implementazione

Ottimizzazioni specifiche per l'hardware

Conclusione

Citazione