Omogućavanje modela sa bilionima parametara na AWS EFA

U kompaniji Perplexity koristimo najbolje modele za naš proizvod, naše API interfejse i naše istraživačke timove. Veliki open-source Mixture-of-Experts modeli, kao što je Kimi-K2, postavljaju posebne izazove, s obzirom na to da najveći inferencijski čvorovi sa 8x NVIDIA H200 GPU-ova ne mogu efikasno da ih smeste, što uslovljava primenu u više čvorova (multi-node). Predstavljamo skup kernela za paralelizam stručnjaka (expert parallelism) koji postižu vrhunska kašnjenja na ConnectX-7 mreži, nadmašujući performanse sistema DeepEP. Isti kerneli su ujedno i prvi koji ostvaruju održiva kašnjenja na AWS Elastic Fabric Adapteru (EFA), omogućavajući primenu modela sa bilionima parametara.

Isprobajte naše kernele na GitHub-u i pročitajte kompletan istraživački rad na arXiv-u.

Uvod

Mixture-of-Experts (MoE) se nametnuo kao standardna arhitektura za skaliranje do stotina milijardi ili biliona parametara, istovremeno nudeći razumno kašnjenje u vreme inferencije. MoE modeli zamenjuju gusti sloj transformera skupom stručnjaka i malim linearnim usmeravajućim slojem koji određuje sa kojim stručnjakom se token množi. Grupisanjem tokena povezanih sa istim stručnjakom, problem postaje retak (sparse) i trivijalno paralelizabilan, pošto različiti stručnjaci sada mogu biti smešteni na različitim GPU-ovima raspoređenim na više čvorova.

Za razliku od drugih oblika paralelizma, kao što su paralelizam tenzora (Tensor Parallelism - TP) i paralelizam podataka (Data Parallelism - DP), koji se lako mogu primeniti pomoću biblioteka kolektivne komunikacije kao što su torch.distributed i NCCL, MoE usmeravanje uključuje retku (sparse) peer-to-peer komunikaciju. Kernel za raspodelu (dispatch) mora da podeli skup dolaznih tokena i prosledi ih rangovima koji hostuju stručnjake, dok njegov dualni problem, kernel za kombinovanje (combine), mora da primi obrađene ekvivalente i izračuna njihovu ponderisanu sredinu. Ovi prenosi se ne mapiraju na kolektive i zahtevaju prilagođene kernele radi postizanja niskog kašnjenja.

U kompaniji Perplexity razvijamo i prenosive međučvorne (inter-node) kernele i specijalizovane unutarčvorne (intra-node) kernele, a naš rad detaljnije prikazujemo u prethodnom članku o MoE komunikaciji. U drugom blog postu smo takođe pokazali da primena na više čvorova može doneti niže kašnjenje i veći protok za velike MoE modele.

U ovom članku se dublje bavimo implementacijom naših novih međučvornih kernela, čime se značajno unapređuju performanse našeg prethodnog rada zasnovanog na NVSHMEM-u. Pored prikazivanja vrhunskih rezultata na ConnectX-7, takođe pokazujemo da se u okviru našeg inferencijskog izvršnog okruženja ovi kerneli mogu koristiti za efikasno opsluživanje najvećih open-source modela.

Inferencija preko InfiniBand-a

Paralelizam tenzora i paralelizam stručnjaka mogu se efikasno skalirati unutar jednog čvora, pošto NVLink nudi visok protok (900 GB/s) uz kašnjenja u mikrosekundama. Iako su NVIDIA GB200 i GB300 sistemi opremljeni sa dovoljno memorije za čuvanje težina i KV keš memorija modela sa bilionima parametara, primene na AWS p5en instanci sa do 8 H200 GPU-ova ograničene su na 1120 GB HBM memorije, koja se mora podeliti između težina modela i KV keš memorije. Posledično, određeni modeli zahtevaju međučvorne primene preko InfiniBand-a.

Iako je InfiniBand i dalje izuzetno brz, on nudi protok od samo 400 Gbps (50 GB/s), što dodaje desetine do stotine mikrosekundi kašnjenja za MoE usmeravanje. Srećom, dodata kašnjenja se mogu sakriti upotrebom deljenih stručnjaka, mikro-grčenja (micro-batching) i preklapanja izračunavanja i komunikacije kako bi se isporučio održiv protok po konkurentnoj ceni. DeepSeek je pionir ovih tehnika kroz DeepSeek-R1, primenjen na klasterima H800 GPU-ova i NVIDIA ConnectX-7 InfiniBand adaptera.

Na AWS klasterima situacija je drugačija zbog upotrebe njihovih prilagođenih Elastic Fabric adaptera (EFA). Iako oni obezbeđuju vršni protok od 400 Gbps na kolektivima koji se tipično koriste u radnim opterećenjima za obuku, oni blago zaostaju za ConnectX-7 adapterima u pogledu veličina poruka razmenjenih tokom MoE raspodele i kombinovanja. Pored toga, EFA ne podržava GPUDirect Async, što zahteva pomoćnu nit (thread) CPU-a koja premošćava GPU i mrežni adapter (NIC) radi pokretanja prenosa. Pored dodate složenosti, dodatne PCIe transakcije dodaju režijske troškove vredne mikrosekundi svim prenosima. Okviri kao što je NVSHMEM omogućavaju autorima kernela da obezbede implementacije nezavisne od uređaja, međutim generički API interfejsi stvaraju režijske troškove koji primene čine neefikasnim u pogledu i troškova i kašnjenja.

pplx-kerneli zasnovani na NVSHMEM-u

Naši prethodni kerneli su apstrahovali osnovnu implementaciju prenosa kroz NVSHMEM. Oslanjali su se na nvshmemx_putmem_signal_nbi_warp za prenos tokena jedan po jedan, koristeći atomske brojače za sinhronizaciju parnjaka. Iako je ovo nudilo razumno  performanse na ConnectX-7 mrežnim adapterima sa IBGDA-om, implementacija na bazi niti za ConnectX-7 IBRC bila je znatno sporija, dok implementacija za EFA nije uspela da obezbedi kašnjenja manja od jedne milisekunde.

Korišćeni su kerneli kao testno okruženje za osnovni hardver, uz istovremeno utvrđivanje najpogodnijeg interfejsa za izlaganje našim modelima. Pretpostavili smo da česti prenosi i interakcije između kernela i mrežnih adaptera ili pomoćnih niti (threads) ometaju performanse. Pored toga, naša implementacija se nije oslanjala ni na kakve garancije hardvera osim pouzdanosti, što sugeriše da bismo mogli da ubrzamo proces smanjenjem režijskih troškova sinhronizacionih primitiva.

Prateći ova zapažanja, osmislili smo kernele zasnovane na EFA koji izlažu isti interfejs, optimizujući interakcije između kernela uređaja i host pomoćnog procesa ispod haube radi postizanja vršnih performansi.

MoE all-to-all preko EFA

Uvećaj

Novi kerneli prate hibridnu CPU-GPU arhitekturu, gde GPU kerneli komuniciraju sa modelom na uređaju, dok pomoćna nit (thread) koja se izvršava na host CPU-u upravlja interakcijama sa mrežnim adapterom. Novi kerneli ponovo koriste isti TransferEngine koji smo izgradili za prenose KV keša.

MoE usmeravanje se obavlja pomoću kernela za raspodelu (dispatch), koji prima tokene i stručnjake po tokenu od modela u cilju njihovog usmeravanja, šaljući ih preko mreže do rangova koji hostuju te stručnjake. Nakon što grupisani GEMM kerneli pomnože tokene sa težinama stručnjaka, kernel za kombinovanje prenosi tokene nazad na njihove izvorne rangove i izračunava njihovu ponderisanu sredinu pre njihovog vraćanja modelu.

I kerneli za raspodelu i kerneli za kombinovanje podeljeni su na polovinu za slanje (sender) i polovinu za primanje (receiver) radi lakšeg preklapanja i mikro-grčenja. Deo za slanje preuređuje tokene u baufere za slanje kako bi se mogli upisati u parnjake kroz jednu RDMA operaciju upisa. Nakon izvršenja dela za slanje, GPU može da obavlja druge poslove (kao što je množenje deljenog stručnjaka) dok su prenosi preko mreže na čekanju u pozadini. Kerneli za primanje blokiraju rad dok se svi prenosi ne prijave kao primljeni. U pozadini, radna nit na strani hosta (worker thread) upravlja prenosima: GPU joj šalje signal kada su bauberi za slanje popunjeni, nakon čega se pokreću RDMA upisi. Kada mrežni adapter signalizira završetak, on deblokira prijemnik, koji vraća tokene modelu.

Kerneli su usko povezani sa hostom, međusobno komunicirajući kroz unificiranu memoriju i GDRCopy. CUDA unificirana memorija omogućava da se memorija uređaja ili hosta preslika u adresne prostore oba, transparentno upravljajući komunikacijom putem PCIe-a.

Dok je unificirana memorija pogodna za masovne prenose, ispitivanje stanja (polling) sa niskim kašnjenjem obavlja se preko GDRCopy-ja. GDRCopy koristi GPUDirect RDMA za čitanje iz GPU memorije i upisivanje u nju uz kašnjenje u mikrosekundama.

Pored RDMA, unutar čvora kerneli takođe koriste NVLink za razmenu tokena. Ovo može rasteretiti značajan deo saobraćaja, do jedne četvrtine za EP32 i do jedne osmine za EP64, postižući ukupno niža kašnjenja.

Bauberi za slanje i primanje

Dizajn baubera za slanje za kernel raspodele je ključan za postizanje održivih kašnjenja koja uključuju pomoćni proces hosta. Ako postoji N rangova koji svaki raspoređuje T tokena, pri čemu je svaki usmeren na R stručnjaka od ukupno E, rang može primiti do N * T * max(R, E / N) tokena u najgorem mogućem slučaju, ukoliko su svi usmereni na njega. Štaviše, ova minimalna granica može se postići samo kada pošiljaoci rasporede svoje upise uzastopno (contiguously), pošto bi privatni bauberi za primanje potrošili previše memorije. Da bi pošiljaoci mogli da rasporede svoje upise uzastopno bez sukoba, potpune informacije o usmeravanju moraju biti prosleđene svakom rangu. Shodno tome, naši kerneli raspodele prvo razmenjuju broj tokena po stručnjaku, nakon čega svaki od njih može da utvrdi gde idu njihovi upisi i sa kog ofseta primaju tokene od svakog ranga.

Jednostavna razmena informacija o usmeravanju rezultirala bi nepotrebno visokim kašnjenjima, pošto bi svaki rang morao da čeka da upisi stignu pre slanja sopstvenih upisa. Umesto toga, rezervišemo malu količinu privatnog prostora po pošiljaocu na svakom rangu prijemnika i šaljemo neke od tokena zajedno sa upisima informacija o usmeravanju. Ovo osigurava puno iskorišćenje propusnog opsega tokom većeg dela procesa raspodele. Nakon što se rute razmene, preostali tokeni se prenose putem jednog upisa po udaljenom parnjaku. Razmena informacija o usmeravanju i baubera fiksne veličine, zajedno sa prenosima koji ih popunjavaju, prikazani su na slici ispod.

Uvećaj

Kerneli za kombinovanje su znatno jednostavniji, pošto su im informacije o usmeravanjima već dostupne iz faze raspodele. Oni mogu da kopiraju tokene u jednom upisu po parnjaku, nakon čega vrše čitanje iz baubera prijemnika kako bi izračunali ponderisanu sredinu tokena u izlaznim tenzorima.

Za NVLink, bauberi za slanje i primanje unutar čvora mapirani su u parnjake kako bi oni mogli direktno da guraju ili povlače tokene sa susednih rangova.

Raspodela i kombinovanje

Slike ispod detaljno prikazuju vreme i interakcije između GPU-a, CPU-a i mrežnog adaptera, pokazujući jednu iteraciju raspodele i kombinovanja isprepletanu sa kernelima modela.

Rad počinje u slanju raspodele (dispatch send), pri čemu kernel agrigira informacije o usmeravanju i izračunava broj tokena po stručnjaku putem atomske akumulacije u deljenoj memoriji unutar jednog bloka. Kada su brojevi tokena dostupni, oni se kopiraju u memoriju hosta i pomoćni proces se odmah obaveštava da počne slanje informacija o usmeravanju svim ostalim partnerskim rangovima. Paralelno s tim, raspodela nastavlja da izračunava kumulativni zbir broja tokena, identifikujući ofsete potrebne za pakovanje tokena u kontigune baubere po pošiljaocu. Kada su ofseti dostupni, svi tokeni koji se šalju preko EFA pakuju se u baubere za slanje. Pomoćni proces dobija obaveštenje da postavi upise u baubere za slanje međučvornih parnjaka. Unutar čvora, kerneli takođe odmah guraju tokene u privatne baubere unutar istog čvora preko NVLink-a.

Uvećaj

Uvećaj

Kada rang primi informacije o usmeravanju od svih parnjaka, on ih centralizuje i izračunava različite ofsete neophodne za identifikaciju tokena za raspodelu i kombinovanje. Ovaj rad se obavlja na CPU-u dok su prenosi već na čekanju, tako da nije kritičan po pitanju kašnjenja. Naknadno, kada su svi ofseti poznati, upisi za ostatak tokena, koji premašuju veličinu privatnog baubera, se postavljaju i faza slanja raspodele je završena.

Na strani prijemnika, kernel prvo povlači iz istog čvora tokene koji premašuju kapacitet privatnih baubera preko NVLink-a. Zatim se blokira sve dok pomoćni proces hosta ne signalizira završetak svih prenosa. Kada se signal primi, tokeni se preuređuju iz baubera za prijem u izlazni tenzor prosleđen grupisanom GEMM-u. S obzirom na to da se bauberi ponovo koriste između raspodele i kombinovanja, i NVLink i RDMA barijere osiguravaju da su bauberi prijemnika pročitani pre nego što se dozvoli sledećem kernelu da ih prebriše.

Kerneli za kombinovanje prate otprilike sličan šablon: preuređuju tokene, signaliziraju pomoćnom procesu i čekaju završetak. Međutim, metapodaci za postavljanje njihovih upisa su odmah dostupni i svi udaljeni ofseti su poznati, stoga se svi tokeni preneti unutar čvora guraju preko NVLink-a na primaoce, dok se međučvorni tokeni prenose putem jednog upisa. Pre vraćanja rezultata, kernel akumulira tokene iz baubera prijemnika, upisujući ih u izlazne tenzore.

TransferEngine za MoE usmeravanje

Da bismo napravili prototip naše implementacije, inicijalno smo ponovo iskoristili TransferEngine izgrađen za razdvojeno popunjavanje (prefill) i dekodiranje u cilju izdavanja i praćenja RDMA upisa. Međutim, za razliku od razdvojenog dekodiranja gde svaki čvor postavlja veliki broj upisa na istog parnjaka, MoE raspodela postavlja male prenose na sve parnjake u jednoj operaciji. Out-of-the-box performanse nisu bile optimalne i naša implementacija je zaostajala za DeepEP-om na ConnectX-7 mreži za oko 20 mikrosekundi. Shodno tome, specijalizovali smo TransferEngine da podržava dve operacije: razbacivanje (scatter) i barijeru (barrier), kako bismo optimizovali interakcije sa velikim brojem parnjaka.

I operacije razbacivanja i operacije barijere izlažu grupe parnjaka koje se mogu unapred registrovati radi unaprednog popunjavanja struktura podataka koje koristi osnovna implementacija prenosa (šabloni radnih zahteva - WR templating). Operacija razbacivanja kopira različite isečke iz istog izvornog baubera u različite ofsete na svim članovima grupe parnjaka. Barijera signalizira svim parnjacima preko neposredne poruke (immediate). Enkueovanje i rukovanje ovim operacijama u masovnim količinama štedi dragocene mikrosekunde na kritičnoj putanji. Na EFA mreži, gde se dva mrežna adaptera od 200 Gbps udružuju radi pružanja propusnog opsega od 400 Gbps, delimo grupe parnjaka na dva mrežna adaptera umesto deljenja bajtova svakog prenosa.

Iako smo inicijalno razvili TransferEngine isključivo na EFA koristeći libfabric, MoE kerneli rade na ConnectX-7 bez izmene i jedne linije koda nakon dodavanja podrške za ConnectX-7 u TransferEngine korišćenjem libibverbs. U poređenju sa EFA baziranim na pouzdanom datagramu sa mogućnošću skaliranja (SRD), ConnectX-7 zahteva dodatni rad na podešavanju veze i upravljanju parnjacima. Koristimo par redova čekanja nepouzdanog datagrama (UD) za pokretanje pouzdane veze (RC). Kreiramo dva para redova čekanja za svakog udaljenog parnjaka, jedan za dvosmerni SEND/RECV i jedan za jednosmerni WRITE_IMM, zato što i RECV i WRITE_IMM troše RECV operaciju redom, dok dve operacije zahtevaju različite veličine baubera. Pored toga, ublažavamo redosled PCIe transakcija između ConnectX-7 i GPU memorije. Slično kao kod EFA, primenjujemo optimizaciju šablona radnih zahteva (WR templating). Pored toga, prilikom slanja više malih poruka istom parnjaku, koristimo uvezivanje radnih zahteva (WR chaining) radi poboljšanja brzine poruka.

Pošto je rad na strani GPU-a identičan na oba adaptera, potisnuli smo osnovnu implementaciju prenosa do njenih krajnjih granica kako bismo nadmašili performanse DeepEP-a i istakli da je jaz između EFA i ConnectX-7 MoE usmeravanja određen mogućnostima osnovnog hardvera.

Evaluacija

Kašnjenje raspodele/kombinovanja

Procenili smo kašnjenja raspodele i kombinovanja na H200 sistemima opremljenim sa ConnectX-7 i EFA adapterima na standardnoj DeepSeek-V3 i Kimi-K2 konfiguraciji sa skrivenom dimenzijom od 7168, blok-skaliranom fp8 raspodelom i bf16 kombinovanjem. Uzimamo u obzir veličinu serije (batch size) od 128 za raspodelu i 4096 za kombinovanje. Na ConnectX-7 poređenje vršimo sa DeepEP-om i našim prethodnim pplx-kernelima baziranim na NVSHMEM-u, koristeći i GPUDirect Async (IBGDA) i pomoćni proces hosta (IBRC). Ostvarujemo kombinovana kašnjenja od 459 µs, 582 µs i 692 µs na EP16, EP32 i EP64, redom. Nadmašujemo druge kernele zasnovane na EFA koji se oslanjaju na generičke pomoćne procese, kao što je UCCL-EP, koji postižu 519 µs, 966 µs i 1159 µs.

Uvećaj

Prilikom dekodiranja, uprkos tome što su naši kerneli nekoliko mikrosekundi sporiji od DeepEP-a, ostvarujemo ukupno niža kašnjenja jer su naši kerneli za kombinovanje znatno brži. Iako koristimo pomoćni proces hosta, zahvaljujemo se našim masovnim prenosima i prevazilazimo performanse svih implementacija bez pomoćnih procesa zasnovanih na IBGDA. Takođe pokazujemo da je implementacija pomoćnog procesa važna, pošto je naš visoko specijalizovani pomoćni proces reda veličine brži od generičke NVSHMEM implementacije. Uprkos činjenici da je EFA dvostruko sporiji od ConnectX-7 pri prenosu paketa od 256 KB, koji su blizu tipične veličine paketa razmenjene tokom MoE usmeravanja, jaz između ova dva adaptera je ukupno manje značajan.

Uvećaj

Iako se naši kerneli optimizovani za dekodiranje mogu skalirati na veličine serija popunjavanja (prefill) bez ikakvih dodatnih podešavanja, performanse zaostaju za DeepEP-om. To je prvenstveno zbog činjenice da mi akumuliramo tokene samo na strani prijemnika za kombinovanje, dok DeepEP akumulira u pošiljaocu kako bi se smanjila količina prenetog podataka. Međutim, ova optimizacija može imati implikacije na tačnost. Pored toga, pošto kerneli prenose podatke u jednom upisu, koristimo više memorije za baubere slanja i primanja.

Protok za Kimi-K2 / DeepSeek-V3

Da bismo dokazali održivost naših kernela, vršimo testiranje performansi end-to-end našeg internog inferencijskog endžina koji opslužuje DeepSeek-V3 (671B parametara) i Kimi-K2 (1T parametara) na više p5en instanci opremljenih sa Elastic Fabric adapterima (EFA) i H200 GPU-ovima. Pokrećemo model u režimu verifikacije nacrta (draft verification mode) i izveštavamo kašnjenje pod pretpostavkom stope prihvatanja od 80% koristeći jedan MTP sloj.

Na DeepSeek-V3 modelu vršimo poređenje sa referentnom vrednošću (baseline) unutar čvora EP=8 koja pokreće efikasne kernele raspodele i kombinovanja koji koriste isključivo NVLink. Međučvorne primene u celini ili dostižu ili premašuju performanse primene na jednom čvoru, naročito na serijama srednje veličine, što pokazuje da skaliranje preko čvorova donosi poboljšanja celokupnom protoku.

Na Kimi-K2 modelu, kerneli omogućavaju opsluživanje. Zbog svoje veličine, model ne može da stane unutar istog čvora koristeći paralelizam podataka (DP) na H200 čvoru, stoga je za opsluživanje modela potrebna međučvorna primena. Pošto model ima manje glava pažnje (attention heads), postižemo bolja kašnjenja nego kod DeepSeek-V3 modela i možemo da opslužujemo model na EFA mreži uz održivo kašnjenje.

Uvećaj

Uvećaj

Takođe smo razmotrili mikro-grčenje, mada je ono poboljšalo kašnjenja samo pri veoma velikim veličinama serija ili kada se koristi više od jednog tokena nacrta. Da bi se imala korist od deljenja problema grupisane GEMM na dva dela, veličine serija tokena moraju biti dovoljno guste. Veći deo vremena se i dalje troši na MoE raspodelu i kombinovanje, što ukazuje na to da će svako dalje poboljšanje kernela posljedično poboljšati kašnjenja i olakšati mikro-grčenje.

Budući rad

Blisko sarađujemo sa inženjerima iz AWS-a na poboljšanju performansi na EFA mreži. Sa naše strane, pratićemo ažuriranja biblioteke libfabric koja će smanjiti režijske troškove ravni podataka. Pored toga, eksperimentisaćemo sa efa-direct tehnologijom kako bismo dodatno smanjili režijske troškove u korisničkom prostoru i direktno komunicirali sa hardverom.

Reference

Tehnički izveštaj za DeepSeek R1

Tehnički izveštaj za Kimi K2

DeepEP

UCCL-EP

Poboljšanje mrežnih performansi HPC sistema korišćenjem sistema NVIDIA Magnum IO NVSHMEM i GPUDirect Async

Citiranje

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}