RDMA komunikace bod-bod pro systémy LLM

Systémy LLM se tradičně spoléhaly na kolektivní komunikaci pro strukturované paralelní vzorce při trénování a inferenci. Nově vznikající pracovní zátěže, jako je disagregovaná inference, routování typu Mixture-of-Experts (MoE) a asynchronní doladění pomocí posilovaného učení (reinforcement learning fine-tuning), však vyžadují flexibilitu komunikace typu bod-bod.

Náš nejnovější výzkumný článek představuje TransferEngine, přenositelnou komunikační knihovnu založenou na RDMA, která funguje napříč různými síťovými řadiči (NIC), včetně NVIDIA ConnectX-7 a AWS Elastic Fabric Adapter (EFA). Její efektivitu demonstrujeme na třech produkčních systémech v Perplexity AI:

Disagregovaná předvolba (prefill) a dekódování: Přenos KvCache vrstvu po vrstvě pomocí stránkovaných zápisů a sledovačů UVM, což umožňuje nezávislé škálování instancí pro předvolbu a dekódování.

Přenos vah pro RL po-trénink za méně než 2 sekundy: Bleskově rychlá synchronizace vah mezi trénovacími a inferenčními uzly.

Zprovoznění modelů s bilionem parametrů na AWS EFA: Nízko latenční operace distribuce a kombinace MoE pomocí primitiv scatter a barrier, které dosahují špičkové latence na ConnectX-7 a představují první životaschopnou implementaci kompatibilní s AWS EFA.

Tento článek se zaměřuje přímo na TransferEngine, komunikační knihovnu, která tyto systémy pohání. Přečtěte si celý výzkumný článek na arXiv a prozkoumejte kód na GitHub.

Kolektivní komunikace vyniká u strukturovaných paralelních vzorců, jako je tenzorový paralelismus a datový paralelismus, ale klade omezení, která omezují její použitelnost pro nově vznikající pracovní zátěže:

Pevné členství: Kolektivní komunikace vyžaduje vytvoření globálního komunikačního světa v době inicializace se synchronizovaným nastavením napříč všemi účastníky. U disagregované inference je to problém, protože počet instancí pro předvolbu a dekódování se musí dynamicky škálovat v reakci na vzorce provozu.

Jednotné tvary tenzorů: Kolektivní komunikace předpokládá, že všichni účastníci vyměňují tenzory s identickými tvary a datovými typy. Při routování MoE závisí počet tokenů odeslaných každému expertu na datech a je zjištěn za běhu, což nutí husté kolektivní operace typu all-to-all doplňovat data na velikost vyrovnávací paměti pro nejhorší případ.

Způsob řazení operací: Kolektivní API vynucují synchronizovanou sémantiku, i když aplikace potřebují pouze vědět, že sada operací byla dokončena. Například při přenosu KvCache mohou všechny stránky z instance předvolby do instance dekódování dorazit v libovolném pořadí. Podobně u aktualizací vah RL lze všechny tenzory vah přenášet nezávisle bez omezení pořadí.

Komunikace typu bod-bod založená na RDMA tato omezení přirozeně řeší, ale jejímu přijetí bránily tři překážky:

Omezení knihoven: Zatímco torch.distributed a NCCL poskytují rozhraní API typu bod-bod, jako jsou Send/Recv, tyto operace přesto dědí sémantiku a omezení kolektivní komunikace.

Různorodost hardwaru: Síťové řadiče NVIDIA ConnectX s přenosem Reliable Connection (RC) jsou nejběžnějším nastavením RDMA a zaručují doručení v pořadí. Poskytovatelé cloudu, jako je AWS, používají jiný hardware – Elastic Fabric Adapter (EFA) s přenosem Scalable Reliable Datagram (SRD), který doručuje zprávy mimo pořadí. EFA také vyžaduje agregaci čtyř 100 Gbps síťových řadičů (nebo dvou 200 Gbps síťových řadičů), aby se dosáhlo propustnosti 400 Gbps.

Přenositelnost softwaru: Stávající knihovny RDMA trpí závislostí na dodavateli (vendor lock-in). DeepEP vyžaduje RDMA iniciované GPU (IBGDA) dostupné pouze na ConnectX. NVSHMEM vykazuje vážný pokles výkonu na EFA. Knihovny jako Mooncake a NIXL postrádají podporu EFA nebo jsou v počáteční fázi vývoje. Výsledkem je, že pro inferenci LLM neexistuje žádné životaschopné řešení nezávislé na poskytovateli.

TransferEngine řeší tyto překážky pomocí čtyř klíčových schopností:

Společné předpoklady o přenosu: Namísto cílení na konkrétní síťový řadič předpokládá TransferEngine pouze to, že základní přenos je spolehlivý, ale neseřazený – tj. představuje průnik záruk poskytovaných přenosy RC i SRD.

Synchronizace ImmCounter: Nové primitivum pro notifikaci o dokončení, které funguje bez spoléhání se na řazení zpráv, což umožňuje spuštění stejného kódu na seřazených (RC) i neseřazených (SRD) přenosech.

Podpora více síťových řadičů: Transparentní správa několika síťových řadičů na jedno GPU, která automaticky rozděluje přenosy za účelem maximalizace propustnosti.

Minimální povrch API: Malá sada operací postačující k vytvoření složitých komunikačních vzorců.

Návrh systému TransferEngine, znázorňující cesty požadavků a callbacků, CPU workery, UVM buffery, GPU a NIC napříč dvěma GPU/NUMA doménami.

TransferEngine spouští jedno pracovní vlákno na jedno GPU. Každé vlákno spravuje skupinu domén (DomainGroup), která koordinuje 1 až 4 síťové řadiče RDMA v závislosti na hardwarové platformě (jeden ConnectX-7 poskytuje 400 Gbps, zatímco EFA vyžaduje čtyři 100 Gbps adaptéry nebo dva 200 Gbps adaptéry). V rámci skupiny domén (DomainGroup) je každá doména (Domain) specifická pro daný hardware a spravuje jediný síťový řadič, přičemž zajišťuje správu frontových párů, odesílání úloh a dotazování na dokončení.

Každá instance TransferEngine zpřístupňuje jedinou hlavní adresu pro identifikaci a zjišťování partnerů (peerů). Vzdálení partneři si vyměňují struktury NetAddr, které kódují síťové adresy a lze je serializovat pro komunikaci mimo pásmo (out-of-band). Všichni partneři v systému musí používat stejný počet síťových řadičů na GPU. Tato jednotná konfigurace umožňuje TransferEngine inteligentně rozdělovat přenosy: při zahájení přenosu vědí zdroj i cíl přesně, které síťové řadiče jsou k dispozici, což umožňuje automatické vyrovnávání zátěže napříč všemi adaptéry.

TransferEngine poskytuje minimální API, které abstrahuje složitost základní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 se skládá z několika klíčových komponent:

Operace RDMA vyžadují, aby byly paměťové oblasti registrovány u síťového řadiče před přenosem dat. Funkce reg_mr přebírá ukazatel, délku a zařízení a vrací dva handle:

MrHandle: Lokální handle používaný jako zdroj pro přenosy

MrDesc: Serializovatelný deskriptor, který lze vyměnit se vzdálenými partnery, aby mohli zapisovat do registrované paměti

Struktura MrDesc obsahuje seznam vzdálených síťových adres a vzdálených klíčů (jeden pro každý síťový řadič), které udělují vzdálený přístup k paměťové oblasti. Tento návrh přirozeně zvládá nastavení s více síťovými řadiči.

Pro malé zprávy a zprávy řídicí roviny ve stylu RPC poskytuje TransferEngine funkce submit_send a submit_recvs:

submit_send: Odešle zprávu vzdálenému partnerovi identifikovanému pomocí NetAddr. Zpráva se zkopíruje interně, takže volající může vyrovnávací paměť okamžitě znovu použít nebo uvolnit.

submit_recvs: Vloží pool přijímacích vyrovnávacích pamětí. Když dorazí zpráva, je vyvolán zpětný volání (callback) s přijatými daty. Vyrovnávací paměť se po dokončení zpětného volání automaticky vrátí do poolu.

Pro hromadné přenosy dat poskytuje TransferEngine jednostranné operace zápisu bez kopírování (zero-copy):

submit_single_write: Přenáší souvislou oblast ze zdroje do cíle

submit_paged_writes: Přenáší nesouvislé stránky určené indexy, kroky (strides) a posuny (offsets)

Obě operace přijímají volitelnou 32bitovou okamžitou hodnotu, která je po dokončení doručena přijímači atomicky.

Modul automaticky rozděluje velké přenosy mezi více síťových řadičů, aby maximalizoval propustnost.

Pro operace cílící na více partnerů (což je běžné při routování MoE) poskytuje TransferEngine optimalizovaná primitiva:

add_peer_group: Předem zaregistruje sadu adres partnerů pro rychlé vyhledávání

submit_scatter: Odešle různé datové řezy každému partnerovi ve skupině

submit_barrier: Odešle notifikaci obsahující pouze okamžitou hodnotu všem partnerům (přenos o nulové velikosti bajtů)

Tyto operace jsou optimalizovány pro nízkou latenci ukládáním interních datových struktur do mezipaměti a vytvářením šablon pro požadavky na práci RDMA, čímž se eliminují opakované režie nastavení.

Klíčovou výzvou v systémech LLM je koordinace přenosů RDMA s průběhem zpracování na GPU. V ideálním případě by měl přenos dat začít ihned, jakmile GPU vyprodukuje výstup, avšak dotazování GPU z CPU je neefektivní.

TransferEngine to řeší pomocí sledovačů UVM (UVM Watchers): sjednocených umístění virtuální paměti, která mohou být aktualizována jádry GPU (dokonce i v rámci grafů CUDA) a efektivně dotazována CPU pomocí knihovny GDRCopy. Při alokaci sledovače UVM je zaregistrováno zpětné volání, které se vyvolá vždy, když GPU aktualizuje hodnotu.

To umožňuje vzorce, jako je přenos KvCache vrstvu po vrstvě: po dokončení každé vrstvy transformátoru GPU inkrementuje sledovač, což spustí přenos KvCache dané vrstvy přes RDMA, zatímco následující vrstvy pokračují ve výpočtech.

TransferEngine poskytuje mechanismy notifikace o dokončení jak pro stranu iniciátora (volajícího funkce submit_write), tak pro stranu přijímače (cíle operací zápisu).

Notifikace na straně iniciátora: Při odeslání operace zápisu si volající může zvolit, zda chce být o dokončení informován prostřednictvím funkce zpětného volání, nebo atomického příznaku. Tato notifikace potvrzuje, že data byla úspěšně přenesena ze zdrojové vyrovnávací paměti.

Notifikace na straně přijímače: Klíčovou inovací je ImmCounter, komponenta, která sleduje čítače indexované 32bitovými okamžitými čísly. Když operace zápisu nese okamžitou hodnotu, čítač ImmCounter přijímače se při doručení dat atomicky inkrementuje. Funkce expect_imm_count umožňuje přijímači zaregistrovat zpětné volání, které se spustí, když konkrétní okamžitá hodnota dosáhne cílového počtu. Čítače lze transparentně synchronizovat s GPU pomocí GDRCopy, dotazovat se jich přímo nebo je sledovat prostřednictvím zpětných volání.

Tento mechanismus na straně přijímače poskytuje spolehlivou synchronizaci bez nutnosti řazení zpráv. Funguje stejně dobře na síťových řadičích ConnectX (které zaručují doručení v pořadí přes přenos RC) i na síťových řadičích EFA (které doručují zprávy mimo pořadí přes přenos SRD).

TransferEngine je implementován v jazyce Rust a pečlivě optimalizuje alokace, práci s vlákny a synchronizaci tak, aby minimalizoval latenci a dosáhl vysoké propustnosti. Modul spouští jedno pracovní vlákno na každou skupinu domén (DomainGroup), přičemž každé vlákno je připnuto k jádru CPU na uzlu NUMA, ke kterému jsou zařízení připojena, čímž se minimalizuje latence plánování i přístupu do paměti. Datové struktury specifické pro danou doménu jsou alokovány po připnutí, aby bylo zajištěno, že paměť je rezervována ve správném uzlu NUMA. Jedno pracovní vlákno obsluhuje až 4 domény (Domains) (z nichž každá spravuje jeden síťový řadič), zatímco jiné vyhrazené vlákno se dotazuje GPU za účelem aktualizace sledovačů UVM. Komunikace mezi vlákny probíhá prostřednictvím front bez zámků (lock-free queues).

API přeposílá požadavky do skupiny domén (DomainGroup) obsluhující příslušné zařízení, přičemž první z nich obsluhuje také hostitele. V těsné smyčce (tight loop) se pracovník domény dotazuje na novou práci a upřednostňuje odesílání nových požadavků. V závislosti na hardwaru a konfiguraci jsou požadavky rozděleny a zatížení je vyváženo mezi dostupnými doménami (Domains). První zápis složeného požadavku je okamžitě odeslán do odesílací fronty síťového řadiče. Jakmile jsou nové požadavky vyčerpány, pracovník přistoupí k provádění probíhajících operací a odesílá zápisy za účelem naplnění hardwarové pipeline. Následně se dotazuje front dokončení, aby zjistil dokončené přenosy a přírůstky okamžitých čítačů. Události jsou agregovány za účelem doručení notifikací pro jednotlivé přenosy a předání přenosu vyhrazenému vláknu pro zpětná volání sdílenému všemi skupinami.

Rozdělování uvnitř skupiny domén (DomainGroup) je flexibilní. Přenosy mohou cílit na konkrétní síťové řadiče podle indexu. Jeden zápis lze rozdělit. Stránkované přenosy, operace scatter a barrier, které se všechny překládají na více zápisů, se mohou rozdělit napříč všemi síťovými řadiči.

Domény (Domains) v rámci TransferEngine jsou specializované a optimalizované pro hardware, který mají pod kontrolou:

AWS EFA: Implementace pro EFA používá knihovnu libfabric a spravuje fabric doménu na jeden síťový řadič v rámci skupiny domén (DomainGroup). Klíčová implementační rozhodnutí a optimalizace zahrnují:

Zápisy o nulové velikosti: EFA se odchyluje od specifikace RDMA tím, že vyžaduje platné cílové deskriptory i pro zápisy o nulové velikosti. Proto TransferEngine vynucuje platné deskriptory pro všechny přenosy.

Vytváření šablon požadavků na práci: Předem vyplňuje a uchovává společná pole struktur knihovny libfabric před jejich odesláním, což snižuje režii na jednu operaci u hromadných přenosů a operací se skupinami partnerů.

NVIDIA ConnectX-7: Implementace pro ConnectX-7 používá knihovnu libibverbs s přenosem RC. Klíčová implementační rozhodnutí a optimalizace zahrnují:

Páry front UD pro navázání spojení (handshakes): Každá doména (Domain) používá pár front Unreliable Datagram (UD) k výměně navázání spojení RC.

Oddělené páry front RC: Vytváří 2 páry front RC na partnera – jeden pro obousměrné operace Send/Recv a druhý pro jednostranné operace Write a WriteImm. Toto oddělení je nutné, protože dokončení operací Recv i WriteImm spotřebovávají požadavky na práci v pořadí odeslání, což umožňuje sémantiku Recv na vysoké úrovni a zároveň podporuje WriteImm bez vzájemného rušení.

Vytváření šablon požadavků na práci: Předem vyplňuje společná pole struktur požadavků na práci pro hromadné přenosy a skupiny partnerů.

Zřetězení požadavků na práci: Propojuje až 4 požadavky na práci prostřednictvím ukazatele next struktury ibv_send_wr, čímž snižuje počet zazvoznění zvonku (doorbell rings) vůči síťovému řadiči.

Uvolněné řazení (Relaxed ordering): Povoluje IBV_ACCESS_RELAXED_ORDERING, což umožňuje transakce PCIe mimo pořadí mezi síťovým řadičem a pamětí GPU.

Tím, že se TransferEngine zaměřuje na průnik záruk napříč různými přenosy RDMA a zavádí synchronizaci nezávislou na pořadí prostřednictvím ImmCounter, poskytuje přenositelný základ pro komunikaci bod-bod v systémech LLM. Knihovna je open source a je k dispozici na GitHubu. Vyzýváme komunitu, aby ji prozkoumala a stavěla na ní při řešení nových komunikačních vzorců v rozsáhlých systémech strojového učení.

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

Motivace

TransferEngine

Návrh API

Registrace paměti

Odeslání/příjem (Send/Recv) pro obě strany

Jednostranný zápis (One-Sided Write)

Skupiny partnerů a hromadné operace

Sledovač UVM (UVM Watcher): Signalizace z GPU na CPU

Notifikace o dokončení

Implementace

Hardwarově specifické optimalizace

Závěr

Citace