Effektiv og bærbar mixture-of-experts-kommunikasjon

En oversikt over bærbar Mixture-of-Experts (MoE)-kommunikasjon med fokus på å optimalisere GPU-parallelisme og redusere ventetid i storskala KI-modeller

ForfatterePerplexity Research

Vi presenterer et høyytelses, bærbar, åpen kildekode-bibliotek for Mixture-of-Experts (MoE)-kommunikasjon som oppnår 10x raskere ytelse sammenlignet med standard All-to-All-kommunikasjonsprimitiver.

Vår implementering har flere viktige tekniske innovasjoner som gir overlegen MoE-kommunikasjonseffektivitet:

GPU-initiert kommunikasjon (IBGDA): Støtter direkte GPU-til-NIC-kommunikasjon, noe som reduserer ventetiden betraktelig ved å omgå CPU-involvering

Overlapping av kommunikasjon og beregning: Delt kjernearkitektur med separate sende- og motta-trinn gjør at beregninger kan fortsette samtidig med nettverksoverføringer

Raskeste enkeltnøydeytelse: 2,5x lavere ventetid enn den tidligere raskeste implementeringen på enkeltnøydekonfigurasjoner

Effektiv og bærbar flernøydeytelse: Vår implementering oppnår hastigheter som er opptil 10x raskere enn standard all-to-all-kommunikasjon. Selv om den er omtrent 2x tregere enn svært spesialiserte implementeringer, tilbyr tilnærmingen vår bedre bærbarhet på tvers av NVSHMEM-versjoner og nettverksmiljøer (NVLink, CX-7 og EFA)

Biblioteket er fullstendig åpen kildekode og tilgjengelig på https://github.com/ppl-ai/pplx-kernels.

I denne artikkelen utforsker vi utfordringene med ekspertparallelisme i storskala MoE-modeller og beskriver vår tilnærming til effektiv tokendistribusjon og -kombinasjon på tvers av distribuerte GPU-miljøer.

Introduksjon

Mixture-of-Experts (MoE)-modeller, som DeepSeek R1 og Mixtral 8x7B, forbedrer tette modeller ved å begrense antall vekter som aktiveres for hvert token. I stedet for å oppnå høye parametertall ved å øke antall lag eller øke størrelsen på de lineære lagene i flerlagsperceptronen (MLP) i hvert dekoderlag, erstatter MoE-modeller den tradisjonelle MLP-en med flere eksperter og en ruter. For eksempel, ut av de 671B parametrene i DeepSeek R1, blir bare 37B multiplisert med et gitt token under inferens. Dette reduserer mengden beregninger som kreves for å dekode et token betraktelig, og forbedrer ventetiden sammenlignet med en tett modell.

MoE-modeller byr på noen ekstra utfordringer for inferens sammenlignet med tette modeller. Selv om ekspertene i seg selv er små MLP-lag, inkluderer hvert dekoderlag en ruter som avgjør hvilke eksperter et token sendes til, der hvert token sendes til flere eksperter. Ruteren er typisk et lite lineært lag som produserer en sannsynlighetsfordeling. Vanligvis velges ekspertene med de øverste K-skårene, og den endelige aktiveringen beregnes som et vektet gjennomsnitt ved å summere ekspertutgangene multiplisert med en vekt avledet fra ruterens sannsynlighetsfordeling.

For å minimere ventetiden i distribuerte MoE-systemer kan parallelisme utnyttes på tvers av flere enheter, men dette introduserer kommunikasjonsutfordringer. Modeller som Mixtral 8x7B eller Llama-70B får plass i 8 enheter på tvers av en enkelt node, og drar nytte av raske NVLink-sammankoblinger med lav ventetid (opptil 900 Gbps). Større modeller krever imidlertid deling (sharding) av eksperter på tvers av flere noder ved bruk av InfiniBand (som har en topp på 400 Gbps), noe som gir ytterligere utfordringer med ventetid.

Tokenruting, ekspertberegning og innsamling av resultater i et blandet ekspertlag.

I denne artikkelen utforsker vi problemet med ekspertparallelisme og beskriver vår implementering av GPU-kjerner som effektivt distribuerer tokener til enheter som inneholder de tildelte ekspertene, og samler dem inn igjen ved å beregne den vekte summen av ekspertutganger. Selv om det finnes aggressivt optimaliserte implementeringer rettet mot spesiell nettverksmaskinvare, presenterer vi en svært bærbar implementering som baserer seg på et minimalt sett med NVSHMEM-primitiver, men likevel leverer en ytelse som er 10x raskere enn standard all-to-all-kommunikasjon. Vår implementering oppnår toppmoderne ytelse på enkeltnøydekonfigurasjoner samtidig som den opprettholder utmerket bærbarhet på tvers av ulike nettverksmiljøer.

Seksjonering og parallelisme

For effektiv inferens må en modells vekter holdes i enhetsminnet, samtidig som det er tilstrekkelig med rom for aktiveringer, KV-mellomlagre og andre buffere som kreves for fremoverpasseringen gjennom modellen. De mest kapable modellene overstiger kapasiteten til selv de mest kapable GPU-ene, og derfor må inferensen spres på tvers av flere enheter, som kollektivt kan lagre modellens vekter. Basert på vektskaleringsskjemaene må forskjellige kommunikasjons- og beregningsskjemaer brukes til å avlaste beregninger og synkronisere enheter.

Ekspertparallellisme over to GPU-er.

Ekspertparallelisme (EP), illustrert i den første figuren, parallelliserer bare ekspertberegningen. Ulike eksperter tilordnes forskjellige enheter som inneholder vektene deres. Etter rutjing sendes tokener til den tilsvarende enheten, og resultatene samles inn og akkumuleres etterpå. Kompleksiteten til rutjingen avhenger av graden av parallelisme i de andre delene av modellen: å replikere andre lag kan eliminere behovet for rutjing fullstendig, siden hver rangering kan velge tokenene fra en lokalt replikert rutjingstabell. Men hvis bare én rangering kjører rutjing, kreves det en kringkasting for å distribuere tokener, indekser og vekter til sine respektive eksperter. Til slutt synkroniserer en all-gather- eller en all-to-all-kringkasting utgangstokenene med uansett hvilken rangering som fortsetter kjøringen av en slik modell. En slik implementering er relativt enkel, siden torch allerede eksponerer de nødvendige primitivene, selv om man kan vinne noen fordeler ved å fusjonere akkumulering med gather-operasjonen som samler inn tokenene fra ekspertene.

Kombinere ekspertparallellisme med tensorparallellisme over to GPU-er.

Ekspert-enkel parallelisme skalerer ikke ideelt, da noder i en kluster kan være inaktive mens modellen kjører ikke-ekspertlag, som oppmerksomhet (attention), norm og sampling. Beregningen av disse lagene, primært oppmerksomhet, kan imidlertid også ha nytte av tensorparallelisme (TP). De fleste modeller baserer seg på flerhodet oppmerksomhet (multi-head attention), noe som betyr at oppmerksomhetshodene og deres tilsvarende Q-, K- og V-projeksjoner også kan deles på tvers av enheter, og replikere eller samle bitene mellom ulike lag. Hvis oppmerksomheten spres på tvers av alle enheter, kan en all-gather synkronisere aktiveringene, slik at rutjingen kan replikeres, noe som krever synkroniseringsprimitiver som ligner på tilfellet med ekspertparallelisme for en effektiv implementering. Det er imidlertid grenser for parallelisme på dette nivået, da reduksjon av antall oppmerksomhetshoder under en viss terskel vil gi avtagende avkastning.

Kombinere ekspertparallellisme med dataparallellisme for separate forespørsler over to GPU-er.

For å utnytte alle enheter best mulig og støtte en svært høy grad av ekspertparallelisme på opptil 128 eller 256 GPU-er, kreves dataparallelisme (DP). Under dette oppsettet grupperes enhetene som en modell er delt på, for å håndtere forespørsler samtidig, beregne oppmerksomhet og opprettholde KV-mellomlagre delt på tvers av sin lokale gruppe. Flere forekomster av disse parallelle gruppene samarbeider om ekspertvurdering, der hver enkelt er vert for et forskjellige delsett av ekspertvektene. Avhengig av antall oppmerksomhetshoder kan en gruppe typisk strekke seg opp til størrelsen på en hel node, da intranode-kommunikasjon er raskere. For eksempel, i figuren ovenfor, betjener én DP-rangering uavhengig to forespørsler, A og B, og håndterer oppmerksomhet, norm og annen bokføring for forespørslene. Den andre DP-rangeringen behandler en distinkt forespørsel C. Den første noden er imidlertid vert for halvparten av ekspertene, mens den andre noden er vert for den andre halvparten, og derfor kan tokener fra A og B bli sendt til den andre noden etter rutjing, og omvendt. Dette fører til et glissent kommunikasjonsproblem: hver enhet kan sende et forskjellig antall tokener til en hvilken som helst annen destinasjonsrangering. Eksisterende primitiver fra torch, primært all_to_all, egner seg ikke spesielt godt, da de kan kreve en form for fylling (padding) eller GPU-til-CPU-synkronisering og metadata-kringkasting. For å implementere kommunikasjon effektivt kreves egendefinerte kjerner for å behandle rutjingsinformasjonen og initiere kommunikasjon fra enhetene uten CPU-inngripen. Etter rutjing må en distribusjonskjerne sende tokener til rangeringene de ble rutjet til, mens aktiveringene som tilhører forespørslene i den gjeldende DP-gruppen, må samles inn på kombinasjonssiden. I tillegg må arbeidet balanseres innenfor en DP-gruppe som kan ha flere enheter, for å minimere kringkasting og korrekt replikere tokenene på alle rangeringer for de etterfølgende lagene.

NVSHMEM

NVSHMEM er en NVIDIA-spesifikk OpenSHMEM-implementering som tilbyr bærbar kommunikasjonsfasiliteter mellom enheter som abstraherer bort kompleksiteten til den underliggende maskinvaren. API-et kan uttrykke enhet-til-enhet-lesinger og -skrivinger, som kartlegges til primitivene for individuelle transportlag. Både NVLink og RDMA støttes, noe som gir en viss grad av bærbarhet og gjør det mulig for NVSHMEM-kjerner å operere innenfor en node eller på tvers av et kluster med flere noder. I kjernene våre bruker vi eksterne minneskrivinger og signaliseringsoperasjoner for å overføre data og synkronisere de deltakende enhetene.

NVSHMEM-operasjoner er bygget rundt konseptet symmetrisk minne: de opererer på buffere som er allokert på alle enhetene som deltar i kommunikasjon mellom enheter. En lokal enhet kan bruke en adresse som er avledet i sin egen lokale buffer, sammen med indeksen til målrangeringsenheten, for å spesifisere destinasjonen for en operasjon. Figuren nedenfor illustrerer dette konseptet: begge GPU-ene allokerer symmetriske buffere av samme størrelse, og beholder src-pekere til dem. Den første rangeringen ønsker å sende 3 heltall til den andre, og plasserer dem i starten av bufferen: nvshmem_int_put_nbi utleder startadressen fra den lokale bufferen og spesifiserer målenheten. Den andre rangeringen utleder en offset fra sin egen buffer og sender ett element til den første enheten, med en offset på én. Mens destinasjonsadresser alltid må være symmetriske buffere allokert ved hjelp av nvshmem_alloc, kan kildebuffere være vilkårlige regioner i enhetsminnet, forutsatt at de er forhåndsregistrert med NVSHMEM.

Selv om NVSHMEM tilbyr et bredt spekter av primitiver, stoler kjernene våre på bare 3 funksjoner, og bygger all synkronisering og gjerding (fencing) på dem.

Symmetrisk minne og eksterne skrivinger mellom to GPU-er ved bruk av NVSHMEM.

nvshmemx_putmem_signal_nbi_warp: Overfører en datablokk fra en enhet til en annen, samtidig som den setter et flagg på den eksterne enheten. Operasjonen enten setter (NVSHMEM_SIGNAL_SET) eller øker (NVSHMEM_SIGNAL_ADD) en 64-bits plassering. Flagget oppdateres etter at hele minneblokken er overført. Hvis den eksterne enheten oppdaget en endring i flagget, kan den trygt få tilgang til bufferen i sitt eget minne. Denne funksjonen er nyttig for å samle dataoverføring og synkronisering.

nvshmemx_signal_op opererer på en enkelt minneplassering, typisk et 64-bits flagg, og setter eller øker det atomisk. Det er nyttig for å sende over metadata og synkronisere enheter.

nvshmem_uint64_wait_until brukes på mottakersiden av et signal for å avspørre et flagg til den eksterne enheten oppdaterer det.

Å sikre enhver form for rekkefølge mellom operasjoner gjennom nvshmem_quiet eller nvshmem_fence er kostbart, og derfor unngår implementeringen vår av kjernene bruken av andre NVSHMEM-primitiver og minimerer eventuelle avhengigheter utover individuelle send-mottak-par. Når vi sender data, bruker vi alltid den ikke-blokkerende versjonen av funksjoner, uten å vente på at dataene engang skal bli sendt ut fra den lokale rangeringen. Annen implisitt synkronisering på tvers av kjernegruppen sikrer at disse bufferne ikke blir destruktivt overskrevet mens NIC-ene trenger tilgang til dem.

NVSHMEM bruker GPUDirect RDMA slik at NIC-en direkte kan lese og skrive GPU-minne for datautveksling uten å involvere CPU-er. Videre, på ConnectX NIC-er, støtter NVSHMEM GPU-initiert kommunikasjon (også kjent som InfiniBand GPUDirect Async, eller IBGDA), som lar GPU-en poste RDMA-operasjoner direkte til NIC-en og avspørre fullføringskøen til NIC-en direkte. På plattformer som ikke støtter GPU-initiert kommunikasjon, spawner NVSHMEM en «proxy»-tråd på vert-CPU-en for å initiere kommunikasjon og avspørre fullføring på vegne av GPU-en. NVSHMEM-programmet er bærbar uavhengig av om GPU-initiert kommunikasjon støttes eller ikke. Imidlertid kuttet GPU-initiert kommunikasjon ventetiden betraktelig fordi den fullstendig omgår omveien til CPU-en.

Bærbare kjerner

Vi implementerer MoE-kommunikasjon gjennom et par med distribusjons- og kombinasjonskjerner (dispatch og combine). Distribusjonskjernene har ansvar for å lese tokener og rutjingsinformasjon på hver rangering, og distribuere dem til de aktuelle ekspertene. Kombinasjonskjernene samler inn aktiveringene produsert av ekspertene og sender dem tilbake til kilderangeringene sine, samtidig som de beregner det vekter gjennomsnittet fra de utvalgte ekspertene basert på vektene beregnet av ruteren. Kjernene er videre delt inn i en sende- og mottakskomponent for å tillate at dataoverføringer overlappes med beregninger. Sendekjernene er alle ikke-blokkerende og ikke-synkroniserende: de bare distribuerer alle skrivingene til de eksterne enhetene. På den andre siden leser kombinasjonskjernene bare fra minnet, og venter til alle nødvendige data er overført. Etter å ha distribuert arbeid til NIC-ene, mens dataene overføres asynkront over nettet, kan GPU-ene gjøre annet nyttig arbeid lokalt, som å bruke delte eksperter eller beregne oppmerksomhet.

Hvert kjernebar har sine egne symmetriske minnebuffere, der hver rangering allokerer privat lagring for å motta data fra alle andre rangeringer. Det maksimale antallet tokener som kan distribueres fra en DP-gruppe, er fast på tvers av rangeringene, noe som også setter den øvre grensen hver rangering kan motta for hver lokale ekspert fra hver DP-rangering. Dette gjør at kilderangeringer kan utlede en unik adresse på destinasjonsrangeringen å skrive til, uten å kreve synkronisering mellom dem. Etter at data er mottatt, stokker kjernene rundt på tokenene for å legge dem ut i minnet i det optimale formatet som kreves av de mottakende kjernene. Selv om bufferne har en betydelig størrelse, gjenbrukes de på tvers av alle sekvensielle lag i en modell.

Separate sende- og mottakskjerner for dispatch- og kombinasjonsoperasjoner.

Den eneste formen for global synkronisering, som illustrert i figuren ovenfor, er implementert i kombinasjon-mottak-kjernen. Når data er distribuert fra sending, returnerer en rangering og fortsetter kjøringen når den mottar alle tokenene den trenger fra de andre rangeringene, og går inn i kombinasjon-sende-kjernene så snart de er behandlet. Sperren (barrier) i kombinasjon-mottak sikrer at ingen rangering kan løpe foran og starte distribusjon-sending mens en annen rangering fremdeles venter på å motta data, noe som potensielt kan forårsake destruktiv overlapping. Synkronisering gjøres i kombinasjonskjernen for enkelhets skyld; som en optimalisering kan den samme sperren implementeres mellom distribusjon-mottak- og kombinasjon-sende-kjernene, noe som reduserer ventetiden ytterligere ved å gi NIC-ene mer tid til å avvikle overføringene mens GPU-en kjører andre kjerner.

Begge kjernene er delt på tvers av alle tilgjengelige SM-er på enhetene: mens distribusjon-sende- og kombinasjon-mottak-kjernene må parallellisere over et per-rangering-maksimum tokenantall (max_m), må distribusjon-mottak- og kombinasjon-sende-kjernene håndtere så mange tokener fra hver rangering og lokal ekspert, for et maksimum på max_m * (num_experts // EP) * (EP // TP).

Distribusjon (Dispatch)

Distribusjonskjernen har ansvaret for å distribuere tokener og blokkskaleringsfaktorer til likemenn som er vert for ekspertene tokenene ble rutjet til. Senderdelen stoler på warp-spesialisering for å parallellisere to oppgaver: aggregere rutjingsinformasjon for å fastslå hvor mange tokener som sendes fra hver rangering, og pakke tokenene inn i meldinger som skal distribueres til den eksterne enheten. Mottakersiden venter først på at alle tokenantallene skal bli mottatt, og deretter leser og pakker den ut nyttelastene fra det delte minnet til tensorene som skal gis videre til de nedstrøms grupperte GEMM-kjernene. Den lagrer også informasjon i buffere som deles med kombinasjonskjernene, for å indikere hvor hvert token skal sendes tilbake. Denne informasjonen kreves siden mottakskjernen stokker tokener rundt i sammenhengende buffere i en ikke-deterministisk rekkefølge.

Distribuere tokens og signalisere ankomst mellom kilde- og destinasjonsranger.

I senderdelen er en enkelt warp med 32 tråder på hver blokk ansvarlig for å lese rutjingsinformasjonen og telle antall tokener hver destinasjonsekspert forventes å motta fra denne rangeringen, og sende antallet pluss én ved hjelp av nvshmem_signal_op. Antallet økes med én, ettersom overgangen fra null til ikke-null på den eksterne enden betyr mottak av antallene fra en kilderangering. Parallelt samarbeider de resterende warpene om å kopiere tokener inn i symmetrisk minne på tvers av alle blokker parallelt, og pakker aktiveringer, skaleringsfaktorer og indeksen deres på den lokale rangeringen inn i en sammenhengende bit. Indeksen kreves av kombinasjonssenderen for å fastslå adressen der tokenet vil bli skrevet. Deretter, etter å ha sikret at alle data er kopiert gjennom en sperre på tvers av warp-gruppene, opererer warpene nok en gang uavhengig, der hver sender den samme bufferen til en annen ekspert parallelt. Tokenene sendes ved hjelp av nvshmemx_putmem_signal_nbi_warp, som også atomisk øker antallet sendte tokener fra den lokale rangeringen på den eksterne enheten. Innenfor en DP-gruppe, siden hver rangering eier en replika av tokenet som skal sendes, balanseres distribusjonen jevnt, der hver enhet sender ut et delsett av tokenene.

På mottakssiden synkroniserer alle kjernene seg først med senderangeringene ved å vente på at det totale tokenantallet skal sendes over. Deretter venter de alle på at de atomisk økte sendte tokenantallene skal avleire seg til de totale antallene, noe som indikerer at all nyttelast fra kilderangeringene også er sendt over, takket være semantikken til putmem-kallene. Kjernene avspørrer antallene ved hjelp av nvshmem_uint64_wait_until, og parallelliserer operasjonen på tvers av alle blokker og tråder. Deretter sørger en tverrblokksperre for at ingen blokker leser bufferne med mindre alle data er korrekt mottatt. Spredt over blokker og synkronisert via en atomisk økt teller, kopieres tokenene fra det symmetriske minnet til vanlige sammenhengende buffere som blir inngangene til de nedstrøms kjernene som implementerer ekspertene. Kilderangeringen, ekspertindeksen og tokenindeksen lagres separat, og presiserer nøyaktig plasseringen der kombinasjonskjernen må sende aktiveringene. Selv om tokener fra en DP-gruppe sendes fra forskjellige enheter, grupperes de alle sammen for å bli sendt videre til den tilsvarende eksperten.

Kombinasjon (Combine)

Kombinasjonskjernene er igjen delt inn i sende- og mottakshalvdelene: senderne kopierer de ukvantiserte 16-bits aktiveringene over til de eksterne enhetene, mens mottakskjernene venter på at dataene skal sendes over og beregner det vektede gjennomsnittet av alle ekspertbidrag lokalt. I tillegg fungerer de også som en sperre for å synkronisere distribusjon-kombinasjon-sekvensen: hver rangering setter et flagg på hver likemann ved inngang til sendekjernene, og mottakskjernene får ikke lov til å returnere med mindre de observerer at flagget er satt. Ventetiden for synkronisering er minimal, da den overlapper med den faktiske kommunikasjonen og beregningen.

Kombinere ekspertutdata inn i tokenene på hver dataparallell rang.

På senderiden traverserer kjernene listen over tokener som er tilordnet alle lokale eksperter parallelt, og skriver dem til en buffer på destinasjonsrangeringen. Målrangeringsenheten, ekspertindeksen og tokenindeksen leses fra per-token-bufferne fylt ut av spredningskjernene (scatter) ved mottak av tokenene. Hver sender har sin egen private minneregion per ekspert å skrive til, som vist i figuren ovenfor, noe som unngår behovet for synkronisering. I likhet med distribusjon øker kombinasjon per-token-tellere atomisk på destinasjonsrangeringen for å indikere mottak av dataene: når telleren samsvarer med antall eksperter et token ble distribusjon til, kan tokeninnholdet nås.

I mottakskjernen traverseres listen over tokener parallelt på tvers av flere blokker, og venter på at innholdet deres skal ankomme ved å avspørre flagget satt av signaliseringsoperasjonen. Ved ankomst leses nyttelastene fra de private bufferne, der rutjingstabellen indikerer hvilken buffer som skal leses fra, og hvilken vekt som skal tilordnes hver ekspert. Resultatene skrives deretter til eksternt allokerte tensorer, der kjernen fullfører kjøringen når alle enheter har passert sperren. På tvers av en DP-gruppe mottar alle rangeringer en kopi av hver ekspertaktivering for å beregne sine egne replikaer.

Ytelsestester (Benchmarks)

Vi evaluerer kjernene våre på et kluster med 128 GPU-er fordelt på 16 noder tilkoblet via InfiniBand- og CX-7-NIC-er. Vi prøver både GPUDirect Async (IBGDA) og Reliable Connection (RC) med en CPU-proxy. Vi sammenligner dem med den svært optimaliserte DeepSeek-implementeringen, samt de tette primitivene levert av PyTorch (gjennom NCCL) eller NVSHMEM.

Kommunikasjonslatens for 128 tokens per GPU i en fler-node EP128-konfigurasjon.

Selv om ytelsen til IBRC og IBGDA er lik på den tette NVSHMEM all-to-all-operasjonen (6378 µs vs 6180 µs), er IBGDA betydelig raskere med de glisne (sparse) kjernene. Hvis man legger sammen distribusjon og kombinasjon, bruker IBGDA 902 µs, mens IBRC tar 3223 µs – en 3,6x forbedring i ventetid. Mens all-to-all er båndbreddebundet, kringkaster de glisne kjernene størrelsesordener mindre data og er bundet av ventetid. Ved å utløse nettverksoverføringer direkte fra en GPU, uten å kreve en CPU-proxy for å koordinere GPU-en og NIC-en, reduseres ende-til-ende-ventetiden betraktelig.

Selv om de bærbar kjernene våre er omtrent 2x tregere enn de svært optimaliserte DeepSeek-kjernene, forbedrer de ventetiden med 10x sammenlignet med de tette kjernene (902 µs vs 9944 µs).

I tillegg gjør oppdelingen i sende-mottak-komponenter det også mulig å skjule noe av ventetiden, i motsetning til bibliotekprimitivene.

Kommunikasjonslatens for 128 tokens per GPU i en en-node EP8-konfigurasjon.

På enkeltnode (EP8) bruker NVSHMEM NVLINK til transport, noe som leverer lavere ventetid og høyere gjennomstrømning enn nettverk mellom noder. Våre bærbar kjerner er omtrent 2,5x raskere enn DeepEP på enkeltnode (186 µs vs 481 µs).

Ytterligere forbedringer

Kjernene som er beskrevet her, utkonkurrerer de innebygde primitivene til maskinlæringsrammeverk og tilbyr anstendig ytelse uten overspesialisering for spesifikke enhet-til-enhet-transporter, som NVLink eller InfiniBand. I tillegg til de allerede nevnte mulighetene kan ytterligere ytelsesgevinster oppnås ved å erstatte kommunikasjonsprimitivene med mer spesialiserte versjoner. For eksempel, på tvers av NVLink, kan bruken av symmetrisk minne erstattes med buffere som deles på tvers av enhetene, noe som eliminerer noe kopiering og muliggjør finere synkroniseringsskjemaer på tvers av individuelle tokener i stedet for tokenbunter. På tvers av InfiniBand kan en implementering få direkte tilgang til de underliggende køparene, noe som eliminerer behovet for å overføre både de forventede og sendte tokenantallene og tømme tokener fra køene inn i utdatabuffere når de ankommer. Slike implementeringer koster imidlertid på bekostning av bærbarhet, mens implementeringen vår lar klynger med spesifikke konfigurasjoner evalueres før man forplikter seg til utviklingsinnsats for å optimalisere for spesiell maskinvare.

Konklusjon

Vi har presentert et høyytelses, bærbart bibliotek for MoE-kommunikasjon som oppnår 10x raskere ytelse sammenlignet med standard all-to-all-kommunikasjon, samtidig som det opprettholder kompatibilitet på tvers av ulike maskinvarekonfigurasjoner. På enkeltnoders utplasseringer med NVLink viser løsningen vår 2,5x lavere ventetid enn tidligere implementeringer.

Tilnærmingen vår balanserer ytelse med bærbarhet gjennom nøkkelinnovasjoner, inkludert støtte for GPU-initiert kommunikasjon, en delt kjernearkitektur som muliggjør overlapping av beregning og kommunikasjon, og effektiv tokendistribusjon ved hjelp av minimale NVSHMEM-primitiver. Selv om biblioteket vårt er omtrent 2x tregere enn svært spesialiserte implementeringer på oppsett med flere noder, tilbyr det overlegen fleksibilitet på tvers av ulike nettverksmiljøer (NVLink, CX-7 og EFA).

Ettersom MoE-modeller fortsetter å skalere, vil effektive kommunikasjonsstrategier som vår bli stadig viktigere for praktisk utplassering. Vår fullt åpne kildekode-implementering er tilgjengelig på https://github.com/ppl-ai/pplx-kernels.

Referanser

DeepEP

Teknisk rapport om DeepSeek-V3

NVIDIA NVSHMEM

Sitat

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://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}