Efficiënte en draagbare Mixture-of-Experts communicatie

Wij presenteren een hoogwaardige, draagbare, open-source bibliotheek voor Mixture-of-Experts (MoE)-communicatie die 10x snellere prestaties behaalt in vergelijking met standaard All-to-All communicatie-primitieven.

Onze implementatie bevat verschillende belangrijke technische innovaties die zorgen voor een superieure efficiëntie in MoE-communicatie:

Door GPU geïnitieerde communicatie (IBGDA): Ondersteunt directe GPU-naar-NIC-communicatie, wat de latentie aanzienlijk vermindert door CPU-betrokkenheid te omzeilen.

Overlap van communicatie en berekening: Architectuur met gespleten kernels met afzonderlijke verzend- en ontvangstfasen stelt berekeningen in staat gelijktijdig met netwerkoverdrachten plaats te vinden.

Snelste single-node prestaties: 2,5x lagere latentie dan de voorheen snelste implementatie op single-node configuraties.

Efficiënte en draagbare multi-node prestaties: Onze implementatie behaalt snelheden tot 10x sneller dan standaard all-to-all communicatie. Hoewel ongeveer 2x langzamer dan sterk gespecialiseerde implementaties, biedt onze aanpak betere draagbaarheid over NVSHMEM-versies en netwerkomgevingen (NVLink, CX-7 en EFA).

De bibliotheek is volledig open-source en beschikbaar op https://github.com/ppl-ai/pplx-kernels.

In dit artikel onderzoeken we de uitdagingen van expert-parallellisme in grootschalige MoE-modellen en beschrijven we onze aanpak voor efficiënte token-dispatch en -combinatie in gedistribueerde GPU-omgevingen.

Introductie

Mixture-of-Experts (MoE)-modellen, zoals DeepSeek R1 en Mixtral 8x7B, verbeteren de dichtheid van modellen door het aantal gewichten te beperken dat wordt geactiveerd voor elk token. In plaats van hoge parameteraantallen te bereiken door het aantal lagen te vergroten of de grootte van de lineaire lagen in de Multi-Layer Perceptron (MLP) van elke decoderlaag te vergroten, vervangen MoE-modellen de traditionele MLP door meerdere experts en een router. Bijvoorbeeld, van de 671B parameters van DeepSeek R1 worden er slechts 37B vermenigvuldigd met een gegeven token tijdens de inferentie. Dit vermindert de hoeveelheid berekeningen die nodig zijn om een token te decoderen aanzienlijk en verbetert de latentie in vergelijking met een dicht model.

MoE-modellen bieden enkele extra uitdagingen voor inferentie, vergeleken met dichte modellen. Hoewel de experts zelf kleine MLP-lagen zijn, bevat elke decoderlaag een router die bepaalt naar welke experts een token wordt gestuurd, waarbij elk token naar meerdere experts wordt gestuurd. De router is doorgaans een kleine lineaire laag die een waarschijnlijkheidsverdeling produceert. Meestal worden de experts met de top-K scores gekozen en wordt de uiteindelijke activering berekend als een gewogen gemiddelde, waarbij de expert-outputs worden opgeteld vermenigvuldigd met een gewicht afgeleid van de waarschijnlijkheidsverdeling van de router.

Om de latentie in gedistribueerde MoE-systemen te minimaliseren, kan parallellisme worden benut over meerdere apparaten, maar dit introduceert communicatie-uitdagingen. Modellen zoals Mixtral 8x7B of Llama-70B passen binnen 8 apparaten op een enkele node, waarbij ze profiteren van snelle en latentiearme NVLink-verbindingen (tot 900Gbps). Grotere modellen vereisen echter het sharden van experts over meerdere nodes met behulp van InfiniBand (met een piek op 400Gbps), wat extra latentie-uitdagingen introduceert.

Vergroten

In dit artikel onderzoeken we het probleem van expert-parallellisme en beschrijven we onze implementatie van GPU-kernels die efficiënt tokens naar apparaten sturen die de toegewezen experts bevatten en ze weer verzamelen door de gewogen som van expert-outputs te berekenen. Hoewel er agressief geoptimaliseerde implementaties bestaan die zich richten op specifieke netwerkhardware, presenteren wij een zeer draagbare implementatie die vertrouwt op een minimale set NVSHMEM-primitieven, maar die prestaties levert die 10x sneller zijn dan standaard all-to-all communicatie. Onze implementatie behaalt state-of-the-art prestaties op single-node configuraties met behoud van uitstekende draagbaarheid over verschillende netwerkomgevingen.

Sharding en parallellisme

Voor efficiënte inferentie moeten de gewichten van een model in het apparaatgeheugen worden gehouden, terwijl er ook voldoende ruimte overblijft voor activeringen, KV-caches en andere buffers die nodig zijn voor de voorwaartse pass door het model. De meest capabele modellen overschrijden de capaciteit van zelfs de meest krachtige GPU's, dus moet inferentie worden verspreid over meerdere apparaten die gezamenlijk de gewichten van het model kunnen opslaan. Op basis van de sharding-schema's van gewichten moeten verschillende communicatie- en berekeningsschema's worden gebruikt om berekeningen uit te besteden en apparaten te synchroniseren.

Vergroten

Expert-parallellisme (EP), geïllustreerd in de eerste figuur, parallelliseert alleen de expertberekening. Verschillende experts worden toegewezen aan verschillende apparaten, die hun gewichten vasthouden. Na het routeren worden tokens naar het corresponderende apparaat gestuurd, waarbij de resultaten daarna worden verzameld en geaccumuleerd. De complexiteit van het routeren hangt af van de graad van parallellisme in de andere delen van het model: het repliceren van andere lagen zou de noodzaak voor routeren volledig kunnen elimineren, aangezien elke rang de tokens kan selecteren uit een lokaal gerepliceerde routeringstabel. Als echter slechts een van de rangen routering uitvoert, is een broadcast vereist om tokens, indices en gewichten naar hun respectievelijke experts te sturen. Ten slotte synchroniseert een all-gather of een all-to-all broadcast de output-tokens met de rang die de uitvoering van het model voortzet. Een dergelijke implementatie is relatief eenvoudig aangezien torch de vereiste primitieven al blootstelt, hoewel er enige voordelen zouden kunnen worden behaald door accumulatie te versmelten met de gather-operatie die de tokens van de experts verzamelt.

Vergroten

Expert-only parallellisme schaalt niet ideaal, aangezien nodes in een cluster inactief kunnen zijn terwijl het model niet-expert lagen uitvoert, zoals aandacht, norm en sampling. De berekening van deze lagen, voornamelijk aandacht, kan echter ook profiteren van Tensor-parallellisme (TP). De meeste modellen vertrouwen op multi-head aandacht, wat betekent dat de aandachts-heads en hun corresponderende Q-, K- en V-projecties ook over apparaten kunnen worden geshard, waarbij de plakjes tussen verschillende lagen worden gerepliceerd of verzameld. Als aandacht over alle apparaten is verspreid, kan een all-gather de activeringen synchroniseren, waardoor routering kan worden gerepliceerd, wat voor een efficiënte implementatie synchronisatie-primitieven vereist die vergelijkbaar zijn met het geval van expert-parallellisme. Er zijn echter limieten aan parallellisme op dit niveau, aangezien het verminderen van het aantal aandachts-heads onder een bepaalde drempel waarde verminderende opbrengsten zal opleveren.

Vergroten

Om alle apparaten optimaal te benutten en een zeer hoge graad van expert-parallellisme te ondersteunen, van maximaal 128 of 256 GPU's, is Data-parallellisme (DP) vereist. Onder dit schema worden de apparaten waarover een model is verdeeld gegroepeerd om verzoeken gelijktijdig te verwerken, waarbij aandacht wordt berekend en KV-caches worden onderhouden die over hun lokale groep zijn geshard. Meerdere instanties van deze parallelle groepen werken samen bij expert-evaluatie, waarbij elk een andere subset van de expert-gewichten host. Op basis van het aantal aandachts-heads kan een groep doorgaans de grootte van een volledige node omspannen, aangezien intra-node communicatie sneller is. Bijvoorbeeld, in de figuur hierboven bedient één DP-rang onafhankelijk twee verzoeken, A en B, die aandacht, norm en eventuele andere boekhouding voor de verzoeken afhandelen. De andere DP-rang verwerkt een apart verzoek C. De eerste node host echter de helft van de experts, terwijl de andere node de andere helft host, dus na het routeren kunnen tokens van A en B naar de tweede node worden gestuurd en vice-versa. Dit leidt tot een schaars communicatieprobleem: elk apparaat kan een ander aantal tokens naar een willekeurige andere doelrang sturen. Bestaande primitieven van torch, voornamelijk all_to_all, zijn niet bijzonder geschikt, omdat ze een vorm van padding of GPU-naar-CPU-synchronisatie en metadata-broadcast kunnen vereisen. Om communicatie effectief te implementeren, zijn aangepaste kernels vereist om de routeringsinformatie te verwerken en communicatie vanaf de apparaten te initiëren zonder CPU-interventie. Na het routeren moet een dispatch-kernel tokens naar de rangen sturen waarnaar ze zijn gerouteerd, terwijl aan de combine-kant de activeringen die bij de verzoeken in de huidige DP-groep horen, moeten worden verzameld. Bovendien moet het werk binnen een DP-groep die meerdere apparaten kan hebben, worden gebalanceerd om broadcast te minimaliseren en de tokens correct op alle rangen te repliceren voor de daaropvolgende lagen.

NVSHMEM

NVSHMEM is een NVIDIA-specifieke OpenSHMEM-implementatie die draagbare communicatiefaciliteiten tussen apparaten biedt die de complexiteit van de onderliggende hardware abstraheren. De API kan apparaat-naar-apparaat lees- en schrijfacties uitdrukken, die worden toegewezen aan de primitieven van individuele transportlagen. Zowel NVLink als RDMA worden ondersteund, wat een graad van draagbaarheid verleent en NVSHMEM-kernels in staat stelt om binnen een node of over een cluster van meerdere nodes te werken. In onze kernels gebruiken we schrijf- en signaleringsoperaties op afstand om gegevens over te dragen en de deelnemende apparaten te synchroniseren.

NVSHMEM-operaties zijn opgebouwd rond het concept van symmetrisch geheugen: ze werken op buffers die zijn toegewezen op alle apparaten die deelnemen aan communicatie tussen apparaten. Een lokaal apparaat kan een adres gebruiken dat is afgeleid binnen zijn eigen lokale buffer, samen met de index van de doelrang, om de bestemming van een operatie te specificeren. De figuur hieronder illustreert dit concept: beide GPU's wijzen symmetrische buffers van dezelfde grootte toe, waarbij src-pointers naar hen behouden blijven. De eerste rang wil 3 gehele getallen naar de tweede sturen en ze aan het begin van de buffer plaatsen: nvshmem_int_put_nbi leidt het startadres af uit de lokale buffer en specificeert het doelapparaat. De tweede rang leidt een offset af uit zijn eigen buffer en stuurt één element naar het eerste apparaat, met een offset van één. Hoewel doeladressen altijd symmetrische buffers moeten zijn die zijn toegewezen met nvshmem_alloc, kunnen bronbuffers willekeurige gebieden van het apparaatgeheugen zijn, mits ze vooraf zijn geregistreerd bij NVSHMEM.

Hoewel NVSHMEM een breed scala aan primitieven biedt, vertrouwen onze kernels op slechts 3 functies, waarop alle synchronisatie en fencing zijn gebouwd.

Vergroten

nvshmemx_putmem_signal_nbi_warp: Verplaatst een blok gegevens van het ene apparaat naar het andere, terwijl ook een vlag op het externe apparaat wordt ingesteld. De operatie stelt een 64-bit locatie in (NVSHMEM_SIGNAL_SET) of verhoogt deze (NVSHMEM_SIGNAL_ADD). De vlag wordt bijgewerkt nadat het volledige geheugenblok is overgedragen. Als het externe apparaat een wijziging in de vlag heeft waargenomen, kan het veilig toegang krijgen tot de buffer in zijn eigen geheugen. Deze functie is nuttig voor het samenvoegen van gegevensoverdracht en synchronisatie.

nvshmemx_signal_op werkt op een enkele geheugenlocatie, meestal een 64-bit vlag, en stelt deze atomair in of verhoogt deze. Het is nuttig bij het verzenden van metadata en het synchroniseren van apparaten.

nvshmem_uint64_wait_until wordt gebruikt aan de ontvangende kant van een signaal om een vlag te pollen totdat de externe deze bijwerkt.

Het waarborgen van enige vorm van ordening tussen operaties via nvshmem_quiet of nvshmem_fence is kostbaar, dus onze implementatie van de kernels vermijdt het gebruik van andere NVSHMEM-primitieven en minimaliseert eventuele afhankelijkheden anders dan individuele verzend-ontvangst-paren. Bij het verzenden van gegevens gebruiken we altijd de niet-blokkerende versie van functies, zonder te wachten tot de gegevens zelfs maar vanuit de lokale rang zijn verzonden. Andere, impliciete synchronisatie over de groep kernels zorgt ervoor dat deze buffers niet destructief worden overschreven terwijl de NIC's toegang tot hen nodig hebben.

NVSHMEM gebruikt GPUDirect RDMA zodat de NIC direct GPU-geheugen kan lezen en schrijven voor gegevensuitwisseling zonder CPU's te betrekken. Bovendien ondersteunt NVSHMEM op ConnectX NIC's door GPU geïnitieerde communicatie (ook bekend als Infiniband GPUDirect Async, of IBGDA), waardoor GPU RDMA-operaties direct naar de NIC kan sturen en de voltooiingswachtrij van de NIC direct kan pollen. Op platforms die geen door GPU geïnitieerde communicatie ondersteunen, start NVSHMEM een “proxy”-draad op de host-CPU om communicatie te initiëren en voltooiing namens de GPU te pollen. Het NVSHMEM-programma is draagbaar, ongeacht of door GPU geïnitieerde communicatie wordt ondersteund of niet. Door GPU geïnitieerde communicatie verkort echter de latentie aanzienlijk omdat het de omweg naar de CPU volledig omzeilt.

Draagbare kernels

We implementeren MoE-communicatie via een paar dispatch- en combine-kernels. De dispatch-kernels zijn verantwoordelijk voor het lezen van tokens en routeringsinformatie op elke rang, en het verzenden ervan naar de juiste experts. De combine-kernels verzamelen de activeringen die door de experts zijn geproduceerd en sturen ze terug naar hun bronrangen, terwijl ze ook het gewogen gemiddelde van de geselecteerde experts berekenen op basis van de gewichten die door de router zijn berekend. De kernels zijn verder opgesplitst in een verzend- en ontvangstcomponent om gegevensoverdrachten te laten overlappen met berekeningen. De verzend-kernels zijn allemaal niet-blokkerend en niet-synchroniserend: ze sturen eenvoudigweg alle schrijfacties naar de remotes. Aan de andere kant lezen de combine-kernels alleen uit het geheugen en wachten ze tot alle vereiste gegevens zijn overgedragen. Na het verzenden van werk naar de NIC's, terwijl de gegevens asynchroon over de kabel worden overgedragen, kunnen de GPU's ander nuttig werk lokaal doen, zoals het toepassen van gedeelde experts of het berekenen van aandacht.

Elk paar kernels heeft zijn eigen symmetrische geheugenbuffers, waarbij elke rang privé-opslag toewijst om gegevens van alle andere rangen te ontvangen. Het maximale aantal tokens dat vanuit een DP-groep kan worden verzonden, is vastgelegd over de rangen, wat ook de bovengrens bepaalt die elke rang kan ontvangen voor elke lokale expert van elke DP-rang. Hierdoor kunnen zender-rangen een uniek adres op de doelrang afleiden om naar te schrijven, zonder dat er enige synchronisatie tussen hen vereist is. Nadat gegevens zijn ontvangen, husselen de kernels de tokens om ze in het geheugen te leggen in het optimale formaat dat vereist is door de ontvangende kernels. Hoewel de buffers een aanzienlijke dimensie hebben, worden ze opnieuw gebruikt over alle sequentiële lagen van een model.

Vergroten

De enige vorm van globale synchronisatie, zoals geïllustreerd in de figuur hierboven, wordt geïmplementeerd in de combine-receive kernel. Zodra gegevens vanuit verzenden zijn verzonden, keert een rang terug en zet de uitvoering voort wanneer deze alle tokens ontvangt die het nodig heeft van de andere rangen en gaat de combine-send kernels in zodra ze zijn verwerkt. De barrière in combine-receive zorgt ervoor dat geen enkele rang vooruit kan lopen en dispatch-send kan starten terwijl een andere rang nog wacht op het ontvangen van gegevens, wat potentieel destructieve overlapping veroorzaakt. Synchronisatie wordt gedaan in de combine-kernel voor de eenvoud; als optimalisatie zou dezelfde barrière kunnen worden geïmplementeerd tussen de dispatch-receive en de combine-send kernels, waardoor de latentie verder wordt verminderd door de NIC's meer tijd te geven om de overdrachten te regelen terwijl de GPU andere kernels uitvoert.

Beide kernels zijn verdeeld over alle beschikbare SM's van de apparaten: terwijl de dispatch-send en combine-receive kernels moeten parallelliseren over een per-rang maximaal aantal tokens (max_m), moeten de dispatch-receive en combine-send kernels zoveel tokens van elke rang en lokale expert verwerken, voor een maximum van max_m * (num_experts // EP) * (EP // TP).

Dispatch

De dispatch-kernel is verantwoordelijk voor het verzenden van tokens en blokschalingsfactoren naar peers die de experts hosten waarnaar de tokens werden gerouteerd. De zenderkant vertrouwt op warp-specialisatie om twee taken te parallelliseren: het aggregeren van routeringsinformatie om te bepalen hoeveel tokens van elke rang worden verzonden en het inpakken van de tokens in berichten die naar de remote moeten worden verzonden. De ontvangende kant wacht eerst tot alle aantallen tokens zijn ontvangen, daarna leest en pakt het de payloads uit vanuit het gedeelde geheugen in de tensoren die worden doorgegeven aan de downstream Grouped GEMM-kernels. Het slaat ook informatie op in buffers die worden gedeeld met de combine-kernels, om aan te geven waar elk token moet worden teruggestuurd. Deze informatie is vereist omdat de receive-kernel tokens in aaneengesloten buffers in een niet-deterministische volgorde husselt.

Vergroten

In het zendergedeelte is een enkele warp van 32 draden op elk blok verantwoordelijk voor het lezen van de routeringsinformatie en het tellen van het aantal tokens dat elke doel-expert naar verwachting van deze rang zal ontvangen, waarbij het aantal plus één wordt verzonden met nvshmem_signal_op. Het aantal wordt met één verhoogd, aangezien de overgang van nul naar niet-nul aan de remote kant de ontvangst van de aantallen van een bronrang betekent. Parallel hieraan werken de overgebleven warps samen om tokens naar symmetrisch geheugen te kopiëren over alle blokken heen, waarbij activeringen, schalingsfactoren en hun index op de lokale rang worden ingepakt in een aaneengesloten brok. De index is vereist door de combine-zender om het adres te bepalen waar het token naar zal worden geschreven. Vervolgens, na ervoor te hebben gezorgd dat alle gegevens zijn gekopieerd via een barrière over de warp-groepen, werken de warps opnieuw onafhankelijk, waarbij elk dezelfde buffer parallel naar een andere expert stuurt. De tokens worden verzonden met nvshmemx_putmem_signal_nbi_warp, wat ook atomair het aantal verzonden tokens van de lokale rang op het externe apparaat verhoogt. Binnen een DP-groep, aangezien elke rang eigenaar is van een replica van het te verzenden token, wordt de dispatch gelijkmatig gebalanceerd, waarbij elk apparaat een subset van de tokens verstuurt.

Aan de ontvangende kant synchroniseren alle kernels eerst met de zender-rangen door te wachten tot de totale aantallen tokens zijn verzonden. Daarna wachten ze allemaal tot de atomair verhoogde verzonden token-aantallen zich stabiliseren op de totale aantallen, wat aangeeft dat alle payloads van de bronrangen ook zijn verzonden, dankzij de semantiek van de putmem-aanroep. De kernels pollen op de aantallen met nvshmem_uint64_wait_until, waarbij de operatie over alle blokken en draden wordt geparallelliseerd. Vervolgens zorgt een cross-block barrière ervoor dat geen enkel blok de buffers leest tenzij alle gegevens correct zijn ontvangen. Verspreid over blokken en gesynchroniseerd via een atomair verhoogde teller, worden de tokens gekopieerd vanuit symmetrisch geheugen naar reguliere aaneengesloten buffers die de inputs worden voor de downstream kernels die de experts implementeren. De bronrang, expert-index en token-index worden afzonderlijk opgeslagen, wat precies de locatie aangeeft waar de combine-kernel de activeringen naartoe moet sturen. Hoewel tokens binnen een DP-groep vanaf verschillende apparaten worden verzonden, worden ze allemaal gegroepeerd om te worden doorgegeven aan de corresponderende expert.

Combine

De combine-kernels zijn opnieuw opgesplitst in verzend- en ontvangsthelften: de zenders kopiëren de niet-gekwantiseerde 16-bit activeringen naar de externe apparaten, waarbij de receive-kernels wachten tot de gegevens zijn verzonden en het gewogen gemiddelde van alle expert-bijdragen lokaal berekenen. Bovendien fungeren ze ook als een barrière om de dispatch-combine-volgorde te synchroniseren: elke rang stelt een vlag in op elke peer bij het betreden van de verzend-kernels, waarbij de receive-kernels niet mogen terugkeren tenzij ze waarnemen dat de vlag is ingesteld. De latentie van synchronisatie is minimaal, aangezien deze overlapt met de werkelijke communicatie en berekening.

Vergroten

Aan de zenderkant doorlopen de kernels de lijst met tokens die aan alle lokale experts zijn toegewezen, parallel, en schrijven ze naar een buffer op de doelrang. De doelrang, expert-index en token-index worden gelezen uit de buffers per token die door de scatter-kernels zijn gevuld bij de ontvangst van de tokens. Elke zender heeft zijn eigen privé-geheugenregio per expert om naar te schrijven, zoals aangegeven in de figuur hierboven, waardoor synchronisatie niet nodig is. Net als bij dispatch verhoogt combine atomair tellers per token op de doelrang om de ontvangst van de gegevens aan te geven: wanneer de teller overeenkomt met het aantal experts waarvoor een token was verzonden, kunnen de token-inhouden worden benaderd.

In de receive-kernel wordt de lijst met tokens parallel over meerdere blokken doorlopen, waarbij wordt gewacht tot de inhouden arriveren door de vlag te pollen die is ingesteld door de signaleringsoperatie. Bij aankomst worden de payloads gelezen uit de privé-buffers, waarbij de routeringstabel aangeeft uit welke buffer moet worden gelezen en welk gewicht aan elke expert moet worden toegewezen. De resultaten worden vervolgens geschreven naar extern toegewezen tensoren, waarbij de kernel de uitvoering voltooit zodra alle apparaten de barrière zijn gepasseerd. Over een DP-groep heen ontvangen alle rangen een kopie van elke expert-activering om hun eigen replica's te berekenen.

Benchmarks

We evalueren onze kernels op een cluster van 128 GPU's verspreid over 16 nodes die zijn verbonden via InfiniBand en CX-7 NIC's. We proberen zowel GPUDirect Async (IBGDA) als Reliable Connection (RC) met een CPU-proxy. We vergelijken ze met de hooggeoptimaliseerde DeepSeek-implementatie, evenals de dichte primitieven geleverd door PyTorch (via NCCL) of NVSHMEM.

Vergroten

Hoewel de prestaties van IBRC en IBGDA bij de dichte NVSHMEM all-to-all operatie vergelijkbaar zijn (6378 µs vs 6180 µs), is IBGDA aanzienlijk sneller met de schaarse kernels. Bij het optellen van Dispatch en Combine gebruikt IBGDA 902 µs terwijl IBRC 3223 µs duurt - een verbetering van 3,6x in latentie. Hoewel all-to-all bandbreedte-gebonden is, verzenden de schaarse kernels ordes van grootte minder gegevens, waardoor ze latentie-gebonden zijn. Door netwerkoverdrachten direct vanaf een GPU te activeren, zonder dat een CPU-proxy nodig is om de GPU en de NIC te coördineren, wordt de end-to-end latentie aanzienlijk verminderd.

Hoewel onze draagbare kernels ongeveer 2x langzamer zijn dan de hooggeoptimaliseerde DeepSeek-kernels, verbeteren ze de latentie met 10x in vergelijking met de dichte kernels (902 µs vs 9944 µs).

Bovendien staat de opsplitsing in zender-ontvanger-componenten toe dat een deel van de latentie wordt verborgen, in tegenstelling tot de bibliotheek-primitieven.

Vergroten

Op single-node (EP8) gebruikt NVSHMEM NVLINK voor transport, wat lagere latentie en hogere doorvoer levert dan netwerken tussen nodes. Onze draagbare kernels zijn ongeveer 2,5x sneller dan DeepEP op single-node (186µs vs 481 µs).

Verdere verbeteringen

De kernels die hier worden beschreven presteren beter dan de ingebouwde primitieven van ML-frameworks en bieden behoorlijke prestaties zonder over-specialisatie voor specifieke transporten tussen apparaten, zoals NVLink of InfiniBand. Naast de reeds genoemde mogelijkheden zijn verdere prestatiewinsten haalbaar door de communicatie-primitieven te vervangen door meer gespecialiseerde versies. Bijvoorbeeld, over NVLink zou het gebruik van symmetrisch geheugen kunnen worden vervangen door buffers die over de apparaten worden gedeeld, waardoor kopiëren wordt geëlimineerd en fijnmazigere synchronisatieschema's mogelijk worden, over individuele tokens in plaats van token-batches. Over InfiniBand zou een implementatie direct toegang kunnen krijgen tot de onderliggende wachtrijparen, waardoor de noodzaak om zowel de verwachte als de verzonden token-aantallen over te dragen en tokens uit de wachtrijen naar output-buffers te flushen zodra ze arriveren, wordt geëlimineerd. Dergelijke implementaties gaan echter ten koste van de draagbaarheid, terwijl onze implementaties clusters met specifieke configuraties toelaten om te worden geëvalueerd voordat ontwikkelingsinspanningen worden vastgelegd voor het optimaliseren voor specifieke hardware.

Conclusie

We hebben een hoogwaardige, draagbare bibliotheek voor MoE-communicatie gepresenteerd die 10x snellere prestaties behaalt in vergelijking met standaard all-to-all communicatie, terwijl de compatibiliteit over diverse hardwareconfiguraties behouden blijft. Op single-node implementaties met NVLink vertoont onze oplossing 2,5x lagere latentie dan eerdere implementaties.

Onze aanpak balanceert prestaties met draagbaarheid door belangrijke innovaties, waaronder ondersteuning voor door GPU geïnitieerde communicatie, een architectuur met gespleten kernels die overlap van berekening en communicatie mogelijk maakt, en efficiënte token-dispatch met minimale NVSHMEM-primitieven. Hoewel ongeveer 2x langzamer dan sterk gespecialiseerde implementaties op multi-node opstellingen, biedt onze bibliotheek superieure flexibiliteit over verschillende netwerkomgevingen (NVLink, CX-7 en EFA).

Naarmate MoE-modellen verder schalen, zullen efficiënte communicatiestrategieën zoals de onze steeds belangrijker worden voor praktische inzet. Onze volledig open-source implementatie is beschikbaar op https://github.com/ppl-ai/pplx-kernels.

Referenties

DeepEP

DeepSeek-V3 technisch rapport

NVIDIA NVSHMEM

Citaat

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