Efikasna i prenosiva komunikacija za Mixture-of-Experts modele

Predstavljamo visokoučinkovitu, prenosivu, open-source biblioteku za Mixture-of-Experts (MoE) komunikaciju koja postiže 10x brže performanse u poređenju sa standardnim All-to-All komunikacionim primitivima.

Naša implementacija sadrži nekoliko ključnih tehničkih inovacija koje obezbeđuju vrhunsku efikasnost MoE komunikacije:

Komunikacija inicirana od strane GPU-a (IBGDA): Podržava direktnu komunikaciju između GPU-a i NIC-a, značajno smanjujući latenciju zaobilaženjem učešća CPU-a

Preklapanje komunikacije i izračunavanja: Arhitektura sa podeljenim kernelom i odvojenim fazama slanja i prijema omogućava da se izračunavanje odvija istovremeno sa mrežnim prenosima

Najbrže performanse na jednom čvoru: 2,5x manja latencija u odnosu na prethodno najbržu implementaciju na konfiguracijama sa jednim čvorom

Efikasne i prenosive performanse na više čvorova: Naša implementacija postiže brzine do 10x brže od standardne all-to-all komunikacije. Iako je otprilike 2x sporija od visoko specijalizovanih implementacija, naš pristup nudi bolju prenosivost kroz različite NVSHMEM verzije i mrežna okruženja (NVLink, CX-7 i EFA)

Biblioteka je potpuno open-source i dostupna na https://github.com/ppl-ai/pplx-kernels.

U ovom članku istražujemo izazove paralelizma stručnjaka u velikim MoE modelima i opisujemo naš pristup efikasnom prosleđivanju i kombinovanju tokena u distribuiranim GPU okruženjima.

Mixture-of-Experts (MoE) modeli, kao što su DeepSeek R1 i Mixtral 8x7B, unapređuju guste modele ograničavanjem broja težina koje se aktiviraju za svaki token. Umesto postizanja velikog broja parametara povećanjem broja slojeva ili povećanjem veličine linearnih slojeva u Višeslojnom perceptronu (MLP) svakog sloja dekodera, MoE modeli zamenjuju tradicionalni MLP sa više stručnjaka i ruterom. Na primer, od 671B parametara modela DeepSeek R1, samo 37B se množi sa datim tokenom tokom inferencije. Ovo u velikoj meri smanjuje količinu izračunavanja potrebnu za dekodovanje tokena i poboljšava latenciju u poređenju sa gustim modelom.

MoE modeli predstavljaju dodatne izazove za inferenciju u poređenju sa gustim modelima. Dok su sami stručnjaci mali MLP slojevi, svaki sloj dekodera uključuje ruter koji odlučuje na koje stručnjake se token prosleđuje, pri čemu se svaki token šalje više stručnjaka. Ruter je obično mali linearni sloj koji proizvodi raspodelu verovatnoće. Obično se biraju stručnjaci sa top-K skorovima, a konačna aktivacija se izračunava kao ponderisana sredina, sabiranjem izlaza stručnjaka pomnoženih težinom izvedenom iz raspodele verovatnoće rutera.

Da bi se smanjila latencija u distribuiranim MoE sistemima, paralelizam se može iskoristiti na više uređaja, ali to uvodi komunikacione izazove. Modeli kao što su Mixtral 8x7B ili Llama-70B staju unutar 8 uređaja na jednom čvoru, koristeći brze NVLink međuveze sa niskom latencijom (do 900Gbps). Međutim, veći modeli zahtevaju deljenje (sharding) stručnjaka na više čvorova koristeći InfiniBand (sa vrhuncem od 400Gbps), što uvodi dodatne izazove u vezi sa latencijom.

Rutiranje tokena, ekspertno računanje i prikupljanje izlaza u sloju miksure eksperata.

U ovom članku istražujemo problem paralelizma stručnjaka, opisujući našu implementaciju GPU kernela koji efikasno prosleđuju tokene uređajima koji drže dodeljene stručnjake i sakupljaju ih nazad izračunavanjem ponderisane sume izlaza stručnjaka. Iako postoje agresivno optimizovane implementacije usmerene na specifičan mrežni hardver, mi predstavljamo visoko prenosivu implementaciju koja se oslanja na minimalan skup NVSHMEM primitiva, a ipak pruža performanse 10x brže od standardne all-to-all komunikacije. Naša implementacija postiže performanse na nivou savremenih rešenja na konfiguracijama sa jednim čvorom, uz održavanje izuzetne prenosivosti u različitim mrežnim okruženjima.

Za efikasnu inferenciju, težine modela moraju se držati u memoriji uređaja, uz ostavljanje dovoljno prostora za aktivacije, KV keševe i druge bafere potrebne za prolaz napred (forward pass) kroz model. Najsposobniji modeli premašuju kapacitet čak i najsposobnijih GPU-ova, stoga se inferencija mora raspodeliti na više uređaja koji kolektivno mogu da skladište težine modela. Na osnovu šema deljenja težina, različite komunikacione šeme i šeme izračunavanja moraju se koristiti za rasterećenje izračunavanja i sinhronizaciju uređaja.

Paralelizam stručnjaka (EP), ilustrovan na prvoj slici, paralelizuje samo izračunavanje stručnjaka. Različiti stručnjaci su dodeljeni različitim uređajima, koji drže njihove težine. Nakon rutiranja, tokeni se šalju na odgovarajući uređaj, pri čemu se rezultati naknadno sakupljaju i akumuliraju. Složenost rutiranja zavisi od stepena paralelizma u drugim delovima modela: replikacija drugih slojeva mogla bi u potpunosti eliminisati potrebu za rutiranjem, pošto svaka rang jedinica može izabrati tokene iz lokalno replikovane tabele rutiranja. Međutim, ako samo jedna rang jedinica izvršava rutiranje, obavezno je emitovanje (broadcast) da bi se tokeni, indeksi i težine prosledili svojim odgovarajućim stručnjacima. Konačno, all-gather ili all-to-all emitovanje sinhronizuje izlazne tokene sa bilo kojom rang jedinicom koja nastavlja izvršavanje modela. Takva implementacija je relativno jednostavna pošto torch već izlaže potrebne primitive, mada bi se određene koristi mogle postići spajanjem akumulacije sa operacijom sakupljanja koja prikuplja tokene od stručnjaka.

Kombinovanje ekspertne paralelizacije sa tenzorskom paralelizacijom na dva GPU-a.

Paralelizam zasnovan samo na stručnjacima ne skalira se idealno, jer čvorovi u klasteru mogu biti neiskorišćeni dok model izvršava slojeve koji nisu vezani za stručnjake, kao što su pažnja (attention), normalizacija i uzorkovanje. Međutim, izračunavanje ovih slojeva, prvenstveno pažnje, takođe može imati koristi od tenzorskog paralelizma (TP). Većina modela se oslanja na višejezgarnu pažnju (multi-head attention), što znači da se glave pažnje i njihove odgovarajuće Q, K i V projekcije takođe mogu podeliti (shard) na uređaje, replikujući ili sakupljajući isečke između različitih slojeva. Ako je pažnja raspoređena na sve uređaje, all-gather može sinhronizovati aktivacije, omogućavajući replikaciju rutiranja, što zahteva primitive sinhronizacije slične slučaju paralelizma stručnjaka radi efikasne implementacije. Međutim, postoje ograničenja paralelizma na ovom nivou, jer smanjenje broja glava pažnje ispod određenog praga donosi sve manje koristi.

Kombinovanje ekspertne paralelizacije sa paralelizacijom podataka za odvojene zahteve na dva GPU-a.

Da bi se na najbolji način iskoristili svi uređaji i podržao vrlo visok stepen paralelizma stručnjaka, do 128 ili 256 GPU-ova, potreban je paralelizam podataka (DP). U okviru ove šeme, uređaji na koje je model podeljen grupisani su radi istovremenog obuhvatanja zahteva, izračunavanja pažnje i održavanja KV keševa podeljenih unutar njihove lokalne grupe. Više instanci ovih paralelelnih grupa sarađuje na evaluaciji stručnjaka, pri čemu svaka hostuje različit podskup težina stručnjaka. Na osnovu broja glava pažnje, grupa obično može obuhvatiti veličinu celog čvora, pošto je komunikacija unutar čvora brža. Na primer, na slici iznad, jedan DP rang nezavisno servisira dva zahteva, A i B, obrađujući pažnju, normalizaciju i svako drugo vođenje evidencije za zahteve. Drugi DP rang obrađuje poseban zahtev C. Međutim, prvi čvor hostuje polovinu stručnjaka, dok drugi čvor drugu polovinu, stoga nakon rutiranja tokeni iz A i B mogu biti poslati na drugi čvor i obrnuto. Ovo dovodi do problema retke komunikacije: svaki uređaj može poslati različit broj tokena bilo kojoj drugoj ciljnoj rang jedinici. Postojeći primtivi iz torch-a, prvenstveno all_to_all, nisu naročito pogodni, jer mogu zahtevati neki oblik popunjavanja (padding) ili GPU-to-CPU sinhronizaciju i emitovanje metapodataka. Da bi se komunikacija sprovela efikasno, potrebni su prilagođeni kerneli za obradu informacija o rutiranju i iniciranje komunikacije sa uređaja bez intervencije CPU-a. Nakon rutiranja, kernel za prosleđivanje (dispatch) mora poslati tokene rang jedinicama na koje su rutirani, dok se na strani kombinovanja moraju prikupiti aktivacije koje pripadaju zahtevima u trenutnoj DP grupi. Pored toga, rad se mora izbalansirati unutar DP grupe koja može imati više uređaja, kako bi se smanjilo emitovanje i ispravno replikovali tokeni na svim rang jedinicama za naknadne slojeve.

NVSHMEM je OpenSHMEM implementacija specifična za NVIDIA-u, koja pruža prenosive objekte za međukomunikaciju uređaja koji apstrahuju složenost osnovnog hardvera. API može da izrazi čitanja i upise između uređaja, koji su mapirani na primitive pojedinačnih slojeva prenosa. Podržani su i NVLink i RDMA, što obezbeđuje stepen prenosivosti i omogućava NVSHMEM kernelima da rade unutar čvora ili preko klastera sa više čvorova. U našim kernelima koristimo upise u udaljenu memoriju i operacije slanja signala u cilju prenosa podataka i sinhronizacije uređaja koji učestvuju.

NVSHMEM operacije su izgrađene oko koncepta simetrične memorije: one rade na baferima koji su alocirani na svim uređajima koji učestvuju u međukomunikaciji uređaja. Lokalni uređaj može koristiti adresu izvedenu unutar sopstvenog lokalnog bafera zajedno sa indeksom ciljnog ranga da definiše odredište operacije. Slika ispod ilustruje ovaj koncept: oba GPU-a alociraju simetrične bafere iste veličine, zadržavajući src pokazivače na njih. Prvi rang želi da pošalje 3 cela broja drugom, postavljajući ih na početak bafera: nvshmem_int_put_nbi izvodi početnu adresu iz lokalnog bafera, navodeći ciljni uređaj. Drugi rang izvodi pomak iz sopstvenog bafera, šaljući jedan element prvom uređaju, sa pomeranjem za jedan. Dok odredišne adrese uvek moraju biti simetrični baferi alocirani korišćenjem nvshmem_alloc, izvorni baferi mogu biti proizvoljni regioni memorije uređaja, pod uslovom da su prethodno registrovani kod NVSHMEM-a.

Iako NVSHMEM pruža širok spektar primitiva, naši kerneli se oslanjaju na samo 3 funkcije, gradeći svu sinhronizaciju i ograđivanje (fencing) na njima.

Simetrična memorija i udaljeni upisi između dva GPU-a pomoću NVSHMEM-a.

nvshmemx_putmem_signal_nbi_warp: Prenosi blok podataka sa jednog uređaja na drugi, istovremeno postavljajući fleg na udaljenom uređaju. Operacija ili postavlja (NVSHMEM_SIGNAL_SET) ili uvećava (NVSHMEM_SIGNAL_ADD) 64-bitnu lokaciju. Fleg se ažurira nakon što se prebaci ceo blok memorije. Ako udaljeni uređaj primeti promenjenu vrednost flega, može bezbedno pristupiti baferu u sopstvenoj memoriji. Ova funkcija je korisna za spajanje prenosa podataka i sinhronizacije.

nvshmemx_signal_op radi na jednoj memorijskoj lokaciji, obično 64-bitnom flegu, atomski ga postavljajući ili uvećavajući. Koristan je za prenos metapodataka i sinhronizaciju uređaja.

nvshmem_uint64_wait_until se koristi na strani prijema signala, za anketiranje (polling) flega dok ga udaljeni uređaj ne ažurira.

Obezbeđivanje bilo kog oblika redosleda između operacija kroz nvshmem_quiet ili nvshmem_fence je skupo, stoga naša implementacija kernela izbegava upotrebu bilo kojih drugih NVSHMEM primitiva i svodi sve zavisnosti na minimum, izuzev pojedinačnih parova slanja i prijema. Prilikom slanja podataka, uvek koristimo neblokirajuću verziju funkcija, bez čekanja da podaci uopšte budu poslati van lokalnog ranga. Druga, implicitna sinhronizacija širom grupe kernela osigurava da ovi baferi ne budu destruktivno prepisani dok su NIC-ovima potrebni pristupi njima.

NVSHMEM koristi GPUDirect RDMA tako da NIC može direktno da čita i upisuje GPU memoriju radi razmene podataka bez uključivanja CPU-ova. Štaviše, na ConnectX NIC-ovima, NVSHMEM podržava komunikaciju iniciranu od strane GPU-a (takođe poznatu kao InfiniBand GPUDirect Async, ili IBGDA), što omogućava GPU-u da šalje RDMA operacije direktno na NIC i direktno anketira red čekanja završetaka (completion queue) NIC-a. Na platformama koje ne podržavaju komunikaciju iniciranu od strane GPU-a, NVSHMEM pokreće „proksi“ nit na glavnom CPU-u da inicira komunikaciju i anketira završetak u ime GPU-a. NVSHMEM program je prenosiv bez obzira na to da li je komunikacija inicirana od strane GPU-a podržana ili ne. Međutim, komunikacija inicirana od strane GPU-a značajno smanjuje latenciju jer u potpunosti zaobilazi zaobilaznicu do CPU-a.

Implementiramo MoE komunikaciju kroz par kernela za prosleđivanje (dispatch) i kombinovanje (combine). Kerneli za prosleđivanje su odgovorni za čitanje tokena i informacija o rutiranju na svakom rangu, prosleđujući ih odgovarajućim stručnjacima. Kerneli za kombinovanje sakupljaju aktivacije koje proizvode stručnjaci i šalju ih nazad svojim izvornim rang jedinicama, istovremeno izračunavajući ponderisanu sredinu od izabranih stručnjaka na osnovu težina koje izračunava ruter. Kerneli su dalje podeljeni na komponentu slanja i prijema, kako bi se omogućilo preklapanje prenosa podataka sa izračunavanjem. Kerneli za slanje su neblokirajući i nesinhronizujući: oni jednostavno prosleđuju sve upise udaljenim uređajima. Sa druge strane, kerneli za kombinovanje samo čitaju iz memorije, čekajući dok se svi potrebni podaci ne prenesu. Nakon prosleđivanja posla NIC-ovima, dok se podaci prenose asinhrono preko mreže, GPU-ovi mogu da obavljaju drugi koristan posao lokalno, kao što je primena deljenih stručnjaka ili izračunavanje pažnje.

Svaki par kernela ima sopstvene simetrične memorijske bafere, pri čemu svaki rang alocira privatno skladište za primanje podataka sa svih drugih rangova. Maksimalni broj tokena koji se može proslediti iz DP grupe je fiksiran na svim rangovima, što takođe postavlja gornju granicu koju svaki rang može primiti za svakog lokalnog stručnjaka iz svakog DP ranga. Ovo omogućava rangovima pošiljaocima da izvedu jedinstvenu adresu na odredišnom rangu u koju će upisivati, bez potrebe za bilo kakvom sinhronizacijom među njima. Nakon što su podaci primljeni, kerneli izmešaju tokene kako bi ih rasporedili u memoriji u optimalnom formatu koji zahtevaju dolazni kerneli. Iako baferi imaju značajnu dimenziju, oni se ponovo koriste kroz sve uzastopne slojeve modela.

Odvojeni kernel-i za slanje i prijem za operacije raspodele i kombinovanja.

Jedini oblik globalne sinhronizacije, kao što je ilustrovano na slici iznad, implementiran je u kernelu za kombinovanje-prijem. Jednom kada se podaci proslede iz slanja, rang se vraća i nastavlja izvršavanje kada primi sve tokene koji su mu potrebni od drugih rangova i ulazi u kornene za slanje-kombinovanje čim se oni obrade. Barijera u kombinovanje-prijem osigurava da nijedan rang ne može da ide ispred i počne prosleđivanje-slanje dok bilo koji drugi rang i dalje čeka na prijem podataka, što potencijalno izaziva destruktivno preklapanje. Sinhronizacija se vrši u kernelu za kombinovanje radi jednostavnosti, a kao optimizacija, ista barijera bi mogla da se primeni između kernela za prijem-prosleđivanje i slanje-kombinovanje, dodatno smanjujući latenciju dozvoljavajući NIC-ovima više vremena da završe prenose dok GPU pokreće druge kernele.

Oba kernela su podeljena na sve dostupne SM-ove uređaja: dok kerneli za slanje-prosleđivanje i prijem-kombinovanje moraju da se paralelizuju preko maksimalnog broja tokena po rangu (max_m), kerneli za prijem-prosleđivanje i slanje-kombinovanie moraju da obrađuju taj broj tokena iz svakog ranga i lokalnog stručnjaka, do maksimuma od max_m * (num_experts // EP) * (EP // TP).

Kernel za prosleđivanje je odgovoran za prosleđivanje tokena i faktora skaliranja bloka kolegama koji hostuju stručnjake na koje su tokeni rutirani. Strana pošiljaoca se oslanja na specijalizaciju varpova (warp specialization) da paralelizuje dva zadatka: agregiranje informacija o rutiranju radi utvrđivanja koliko se tokena šalje sa svakog ranga i pakovanje tokena u poruke koje se šalju udaljenom uređaju. Strana primaoca prvo čeka da se prime sve vrednosti broja tokena, zatim čita i raspakuje korisne terete (payloads) iz deljene memorije u tenzore koji će biti prosleđeni nizvodnim Grouped GEMM kernelima. Ona takođe skladišti informacije u baferima koji se dele sa kernelima za kombinovanje, da naznači gde svaki token treba da bude poslat nazad. Ove informacije su potrebne pošto kernel za prijem meša tokene u susednim baferima u nedeterminističkom redosledu.

Raspodeljivanje tokena i signaliziranje njihovog dolaska između izvornog i odredišnog ranga.

U delu pošiljaoca, jedan varp od 32 niti na svakom bloku je odgovoran za čitanje informacija o rutiranju i brojanje broja tokena za koje se očekuje da svaki ciljni stručnjak primi od ovog ranga, šaljući zbir uvećan za jedan koristeći nvshmem_signal_op. Brojač se uvećava za jedan, pošto prelaz sa nule na nenultu vrednost na udaljenoj strani označava prijem brojača sa izvornog ranga. Paralelno, preostali varpovi sarađuju da kopiraju tokene u simetričnu memoriju preko svih blokova paralelno, pakujući aktivacije, faktore skaliranja i njihov indeks na lokalnom rangu u jedan neprekinuti komad. Indeks je potreban pošiljaocu kombinovanja da odredi adresu na koju će token biti upisan. Sledeće, nakon osiguravanja da su svi podaci kopirani kroz barijeru preko grupa varpova, varpovi ponovo rade nezavisno, pri čemu svaki šalje isti bafer drugom stručnjaku paralelno. Tokeni se šalju korišćenjem nvshmemx_putmem_signal_nbi_warp, što takođe atomski povećava broj poslatih tokena sa lokalnog ranga na udaljenom uređaju. Unutar DP grupe, pošto svaki rang poseduje repliku tokena koji se šalje, prosleđivanje je izbalansirano ravnomerno, pri čemu svaki uređaj šalje podskup tokena.

Na strani prijema, svi kerneli se prvo sinhronizuju sa rangovima pošiljaocima čekajući da se ukupni brojevi tokena prenesu. Nakon toga, svi čekaju da se atomski uvećani brojevi poslatih tokena stabilizuju na ukupne brojeve, što ukazuje da su i svi korisni tereti sa izvornih rangova takođe poslati, zahvaljujući semantici putmem poziva. Kerneli anketiraju brojače koristeći nvshmem_uint64_wait_until, paralelizujući operaciju na sve blokove i niti. Naknadno, međublokovska barijera osigurava da nijedan blok ne čita bafere osim ako svi podaci nisu ispravno primljeni. Raspoređeni po blokovima i sinhronizovani preko atomski uvećanog brojača, tokeni se kopiraju iz simetrične memorije u regularne neprekinute bafere koji postaju ulazi u nizvodne kernele koji implementiraju stručnjake. Izvorni rang, indeks stručnjaka i indeks tokena se skladište odvojeno, tačno locirajući mesto gde kernel za kombinovanje mora da pošalje aktivacije. Iako se tokeni unutar DP grupe šalju sa različitih uređaja, oni su svi grupisani zajedno da bi bili prosleđeni odgovarajućem stručnjaku.

Kerneli za kombinovanje su ponovo podeljeni na polovine slanja i prijema: pošiljaoci kopiraju nekvantizovane 16-bitne aktivacije na udaljene uređaje, pri čemu kerneli prijema čekaju da se podaci prenesu i lokalno izračunavaju ponderisanu sredinu svih doprinosa stručnjaka. Pored toga, oni takođe deluju kao barijera, da sinhronizuju sekvencu prosleđivanje-kombinovanje: svaki rang postavlja fleg na svakom kolegi pri ulasku u kernele za slanje, pri čemu kernelima prijema nije dozvoljeno da se vrate osim ako ne primete da je fleg postavljen. Latencija sinhronizacije je minimalna, jer se preklapa sa stvarnom komunikacijom i izračunavanjem.

Kombinovanje izlaza eksperata u tokene na svakom ranku sa paralelizmom podataka.

Na strani pošiljaoca, kerneli paralelno prolaze kroz listu tokena dodeljenih svim lokalnim stručnjacima, upisujući ih u bafer na odredišnom rangu. Ciljni rang, indeks stručnjaka i indeks tokena čitaju se iz bafera po tokenu koje su popunili kerneli rasipanja (scatter) po prijemu tokena. Svaki pošiljalac ima svoj privatni memorijski region po stručnjaku u koji upisuje, kao što je naznačeno na slici iznad, čime se izbegava potreba za sinhronizacijom. Slično kao kod prosleđivanja, kombinovanje atomski uvećava brojače po tokenu na odredišnom rangu da naznači prijem podataka: kada brojač odgovara broju stručnjaka kojima je token prosleđen, sadržaju tokena se može pristupiti.

U kernelu za prijem, lista tokena se obrađuje paralelno kroz više blokova, čekajući da njihov sadržaj stigne anketiranjem flega koji je postavila operacija slanja signala. Po dolasku, korisni tereti se čitaju iz privatnih bafera, pri čemu tabela rutiranja ukazuje iz kog bafera treba čitati i koju težinu dodeliti svakom stručnjaku. Rezultati se zatim upisuju u spolja alocirane tenzore, pri čemu kernel završava izvršavanje kada svi uređaji prođu barijeru. Preko DP grupe, svi rangovi primaju kopiju svake aktivacije stručnjaka da izračunaju sopstvene replike.

Evaluiramo naše kernele na klasteru od 128 GPU-ova raspoređenih na 16 čvorova povezanih preko InfiniBand-a i CX-7 NIC-ova. Isprobavamo i GPUDirect Async (IBGDA) i Reliable Connection (RC) sa CPU proksijem. Poređenja radi, upoređujemo ih sa visoko optimizovanom DeepSeek implementacijom, kao i sa gustim primitivima koje obezbeđuje PyTorch (preko NCCL) ili NVSHMEM.

Latencija komunikacije za 128 tokena po GPU-u u višestrukone-konfiguraciji EP128.

Iako je kod guste NVSHMEM all-to-all operacije performansa IBRC-a i IBGDA-a slična (6378 µs naspram 6180 µs), IBGDA je znatno brži sa retkim kernelima. Sabiranjem prosleđivanja i kombinovanja, IBGDA koristi 902 µs dok IBRC troši 3223 µs – poboljšanje latencije od 3,6x. Dok je all-to-all ograničen protokom (bandwidth bound), retki kerneli emituju znatno manje podataka, te su ograničeni latencijom. Pokretanjem mrežnih prenosa direktno sa GPU-a, bez potrebe za CPU proksijem koji bi koordinisao GPU i NIC, krajnja latencija (end-to-end latency) se značajno smanjuje.

Iako su naši prenosivi kerneli otprilike 2x sporiji od visoko optimizovanih DeepSeek kernela, oni poboljšavaju latenciju za 10x u poređenju sa gustim kernelima (902 µs naspram 9944 µs).

Pored toga, podela na komponente pošiljaoca i primaoca takođe omogućava da se deo latencije sakrije, za razliku od primitiva biblioteke.

Latencija komunikacije za 128 tokena po GPU-u u konfiguraciji EP8 na jednom čvoru.

Na jednom čvoru (EP8), NVSHMEM koristi NVLINK za prenos, pružajući manju latenciju i veći protok od umrežavanja između čvorova. Naši prenosivi kerneli su oko 2,5x brži od DeepEP-a na jednom čvoru (186 µs naspram 481 µs).

Kerneli opisani ovde nadmašuju ugrađene primitive ML okvira i nude pristojne performanse bez preterane specijalizacije za određene prenose između uređaja, kao što su NVLink ili InfiniBand. Pored već pomenutih mogućnosti, dalji dobici u performansama mogu se ostvariti zamenom komunikacionih primitiva specijalizovanijim verzijama. Na primer, preko NVLink-a, upotreba simetrične memorije mogla bi se zameniti baferima koji se dele između uređaja, čime se eliminiše izvesno kopiranje i omogućavaju finije šeme sinhronizacije, preko pojedinačnih tokena umesto serija tokena. Preko InfiniBand-a, implementacija bi mogla direktno da pristupi osnovnim parovima redova čekanja, eliminišući potrebu za prenosom i očekivanog i poslatog broja tokena i ispuštanjem tokena iz redova u izlazne bafere kako pristižu. Međutim, takve implementacije dolaze po cenu prenosivosti, dok naše implementacije omogućavaju evaluaciju klastera sa specifičnim konfiguracijama pre ulaganja razvojnog truda u optimizaciju za određeni hardver.

Predstavili smo visokoučinkovitu, prenosivu biblioteku za MoE komunikaciju koja postiže 10x brže performanse u poređenju sa standardnom all-to-all komunikacijom uz očuvanje kompatibilnosti kroz različite hardverske konfiguracije. Na primenama na jednom čvoru sa NVLink-om, naše rešenje pokazuje 2,5x manju latenciju od prethodnih implementacija.

Naš pristup balansira performanse sa prenosivošću kroz ključne inovacije, uključujući podršku za komunikaciju iniciranu od strane GPU-a, arhitekturu sa podeljenim kernelom koja omogućava preklapanje izračunavanja i komunikacije, kao i efikasno prosleđivanje tokena korišćenjem minimalnih NVSHMEM primitiva. Iako je otprilike 2x sporija od visoko specijalizovanih implementacija na postavkama sa više čvorova, naša biblioteka nudi superiornu fleksibilnost u različitim mrežnim okruženjima (NVLink, CX-7 i EFA).

Kako MoE modeli nastavljaju da se skaliraju, efikasne komunikacione strategije poput naše postaće sve važnije za praktičnu primenu. Naša potpuno open-source implementacija dostupna je na https://github.com/ppl-ai/pplx-kernels.

DeepEP

Tehnički izveštaj za DeepSeek-V3

NVIDIA NVSHMEM

BibTeX

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}

Uvod

Deljenje (Sharding) i paralelizam

NVSHMEM

Prenosivi kerneli

Prosleđivanje (Dispatch)

Kombinovanje (Combine)

Poređenja (Benchmarks)

Dalja unapređenja

Zaključak

Reference

Citiranje