Modellen met biljoenen parameters mogelijk maken op AWS EFA

Bij Perplexity gebruiken we de beste modellen voor ons product, onze API's en onze onderzoeksteams. Grote open-source Mixture-of-Experts-modellen, zoals Kimi-K2, vormen specifieke uitdagingen omdat de grootste inferentie-nodes met 8x NVIDIA H200 GPU's deze niet efficiënt kunnen accommoderen, wat implementaties op meerdere nodes noodzakelijk maakt. We presenteren een reeks kernels voor expert-parallelisme die state-of-the-art latenties op ConnectX-7 bereiken, en daarmee de prestaties van DeepEP overtreffen. Dezelfde kernels zijn ook de eerste die levensvatbare latenties bereiken op AWS Elastic Fabric Adapter (EFA), waardoor de implementatie van modellen met biljoenen parameters mogelijk wordt.

Probeer onze kernels op GitHub en lees het volledige onderzoeksrapport op arXiv.

Mixture-of-Experts (MoE) is naar voren gekomen als de standaardarchitectuur voor het schalen naar honderden miljarden of biljoenen parameters, terwijl het ook een redelijke latentie biedt tijdens inferentie. MoE-modellen vervangen de dichte laag van een transformer door een set experts en een kleine lineaire routeringslaag die bepaalt met welke expert een token moet worden vermenigvuldigd. Door de tokens die aan dezelfde expert zijn gekoppeld te groeperen, wordt het probleem schaars en triviaal paralleliseerbaar, aangezien verschillende experts nu kunnen worden gehost op verschillende GPU's verspreid over meerdere nodes.

In tegenstelling tot andere vormen van parallelisme, zoals Tensor Parallelism (TP) en Data Parallelism (DP), die eenvoudig kunnen worden geïmplementeerd met behulp van collectieve communicatiebibliotheken zoals torch.distributed en NCCL, omvat MoE-routering schaarse peer-to-peer-communicatie. Een dispatch-kernel moet de set inkomende tokens splitsen en verzenden naar de ranks die de experts hosten, terwijl het duale probleem, de combine-kernel, de verwerkte tegenhangers moet ontvangen en hun gewogen gemiddelde moet berekenen. Deze overdrachten corresponderen niet met collectieven en vereisen aangepaste kernels om een lage latentie te bereiken.

Bij Perplexity ontwikkelen we zowel draagbare inter-node kernels als gespecialiseerde intra-node kernels, en we presenteren ons werk in grotere diepte in een eerder artikel over MoE-communicatie. We lieten in een andere blogpost ook zien dat multi-node-implementatie kan zorgen voor een lagere latentie en een hogere doorvoer voor grote MoE-modellen.

In dit artikel duiken we diep in de implementatie van onze nieuwe inter-node kernels, die de prestaties van ons eerdere werk op basis van NVSHMEM aanzienlijk verbeteren. Naast het tonen van state-of-the-art resultaten op ConnectX-7, tonen we ook aan dat de kernels binnen onze inferentie-runtime effectief kunnen worden gebruikt om de grootste open-source modellen te bedienen.

Tensor-parallelisme en expert-parallelisme kunnen effectief schalen binnen een enkele node, aangezien NVLink een hoge doorvoer (900 GB/s) biedt bij een latentie van microseconden. Hoewel NVIDIA GB200- en GB300-systemen zijn uitgerust met voldoende geheugen om de gewichten en KV-caches van modellen met biljoenen parameters te bevatten, zijn implementaties op AWS p5en-instanties met maximaal 8 H200 GPU's beperkt tot 1120 GB HBM, die moet worden verdeeld over zowel modelgewichten als KV-caches. Bijgevolg vereisen bepaalde modellen inter-node implementaties via InfiniBand.

Hoewel InfiniBand nog steeds razendsnel is, biedt het slechts 400 Gbps (50 GB/s) aan doorvoer, wat tientallen tot honderden microseconden aan latentie toevoegt voor MoE-routering. Gelukkig kunnen de toegevoegde latenties worden verborgen door het gebruik van gedeelde experts, micro-batching en computation-communication overlapping om een levensvatbare doorvoer tegen acceptabele kosten te leveren. DeepSeek was pionier met deze technieken via DeepSeek-R1, geïmplementeerd op clusters van H800 GPU's en NVIDIA ConnectX-7 InfiniBand-adapters.

Op AWS-clusters is de situatie anders door het gebruik van hun aangepaste Elastic Fabric Adapters (EFA). Hoewel deze een piekdoorvoer van 400 Gbps leveren op de collectieven die doorgaans in trainingsworkloads worden gebruikt, blijven ze iets achter bij ConnectX-7 wat betreft de berichtgroottes die worden uitgewisseld tijdens MoE-dispatch en -combine. Bovendien ondersteunt EFA geen GPUDirect Async, waardoor een CPU-proxy-draad nodig is die de GPU met de NIC verbindt om overdrachten te initiëren. Naast de extra complexiteit voegen de extra PCIe-transacties microseconden aan overhead toe aan alle overdrachten. Frameworks zoals NVSHMEM stellen kernel-schrijvers in staat om apparaatonafhankelijke implementaties te bieden, maar de generieke API's brengen overheads met zich mee die implementaties inefficiënt maken in termen van zowel kosten als latentie.

Onze vorige kernels abstraheerden de onderliggende transportimplementatie weg via NVSHMEM. Ze vertrouwden op nvshmemx_putmem_signal_nbi_warp om tokens één voor één over te dragen, waarbij atomaire tellers werden gebruikt om peers te synchroniseren. Hoewel dit redelijke prestaties bood op ConnectX-7 NIC's met IBGDA, was de op proxy gebaseerde ConnectX-7 IBRC-implementatie aanzienlijk langzamer, terwijl de EFA-implementatie er niet in slaagde latenties onder een milliseconde te leveren.

We gebruikten de kernels als testomgeving voor de onderliggende hardware, terwijl we ook de meest geschikte interface bepaalden om aan onze modellen bloot te stellen. We veronderstelden dat de frequente overdrachten en interacties tussen de kernels en de NIC's of de proxy-draden de prestaties hinderden. Bovendien vertrouwde onze implementatie niet op garanties van de hardware behalve betrouwbaarheid, wat suggereert dat we de zaken konden versnellen door de overhead van synchronisatie-primitieven te verminderen.

Naar aanleiding van deze observaties hebben we onze op EFA gebaseerde kernels ontworpen die dezelfde interface blootstellen, waarbij we de interacties tussen de apparaat-kernel en host-proxy onder de motorkap optimaliseren voor topprestaties.

RDMA- en NVLink-hybride communicatie met een hostproxy, met initiërende en ontvangende ranks, proxyworkers en NIC's.

De nieuwe kernels volgen een hybride CPU-GPU-architectuur, waarbij GPU-kernels interageren met het model op het apparaat en een proxy-draad die op de host-CPU draait, de interacties met de NIC afhandelt. De nieuwe kernels hergebruiken dezelfde TransferEngine die we hebben gebouwd voor KV-cache-overdrachten.

MoE-routering wordt uitgevoerd door een dispatch-kernel, die de tokens en de per-token experts ontvangt om ze vanuit het model te routeren, en ze over het netwerk verzendt naar de ranks die de experts hosten. Nadat de gegroepeerde GEMM-kernels de tokens met de expertgewichten hebben vermenigvuldigd, draagt de combine-kernel de tokens over naar hun bron-ranks en berekent hun gewogen gemiddelde voordat ze naar het model worden teruggebracht.

Zowel de dispatch- als de combine-kernels zijn gesplitst in een zender- en ontvangerhelft om overlapping en micro-batching te faciliteren. De zender schudt de tokens in verzendbuffers zodat ze via een enkele RDMA-schrijfbewerking naar peers kunnen worden geschreven. Na uitvoering van de zender kan de GPU ander werk verrichten (zoals vermenigvuldiging van gedeelde experts) terwijl de overdrachten via het netwerk op de achtergrond in behandeling zijn. De ontvanger-kernels blokkeren totdat alle overdrachten zijn ontvangen. Op de achtergrond verwerkt een worker-draad aan de host-zijde de overdrachten: deze krijgt een signaal van de GPU wanneer de verzendbuffers zijn gevuld, waarna deze de RDMA-schrijfbewerkingen verzendt. Zodra de NIC de voltooiing signaleert, deblokkeert deze de ontvanger, die de tokens terugstuurt naar het model.

De kernels zijn nauw gekoppeld aan de host en communiceren met elkaar via zowel unified memory als GDRCopy. CUDA unified memory stelt apparaatgeheugen of hostgeheugen in staat om in beide adresruimten te worden toegewezen, waarbij de communicatie via PCIe transparant wordt afgehandeld.

Hoewel unified memory geschikt is voor bulkoverdrachten, gebeurt polling met lage latentie via GDRCopy. GDRCopy gebruikt GPUDirect RDMA om met een latentie van microseconden naar GPU-geheugen te lezen en te schrijven.

Naast RDMA benutten de kernels binnen een node ook NVLink om tokens uit te wisselen. Dit kan een aanzienlijk deel van het verkeer ontlasten, tot een kwart voor EP32 en tot een achtste voor EP64, wat resulteert in algeheel lagere latenties.

Het ontwerp van de verzendbuffers voor de dispatch-kernel is cruciaal voor het bereiken van levensvatbare latenties met de host-proxy. Als er N ranks zijn die elk T tokens verzenden, elk gerouteerd naar R experts van de E, kan een rank in het slechtst denkbare geval tot N * T * max(R, E / N) tokens ontvangen, als ze allemaal ernaartoe worden gerouteerd. Bovendien kan deze minimale grens alleen worden bereikt wanneer de zenders hun schrijfbewerkingen aaneengesloten inrichten, aangezien privé-ontvangstbuffers te veel geheugen zouden verbruiken. Om zenders in staat te stellen hun schrijfbewerkingen opeenvolgend in te richten zonder conflicten, moet volledige kennis van de routeringsinformatie naar elke rank worden gedistribueerd. Bijgevolg wisselen onze dispatch-kernels eerst per-expert tokentellingen uit, waarna ze elk kunnen bepalen waar hun schrijfbewerkingen heen gaan en van welke offset ze tokens van elke rank ontvangen.

Het simpelweg uitwisselen van routeringsinformatie zou resulteren in onnodig hoge latenties, aangezien elke rank zou moeten wachten tot de schrijfbewerkingen aankomen voordat er zelf schrijfbewerkingen worden verzonden. In plaats daarvan reserveren we een kleine hoeveelheid privé-ruimte per zender op elke ontvanger-rank en verzenden we enkele tokens naast de schrijfbewerkingen van de routeringsinformatie. Dit zorgt voor een volledig bandbreedtegebruik gedurende het grootste deel van het dispatch-proces. Nadat routes zijn uitgewisseld, worden de resterende tokens overgedragen via één schrijfbewerking per externe peer. De uitwisseling van routeringsinformatie en de buffers met vaste grootte, samen met de overdrachten die ze vullen, worden in de onderstaande figuur geïllustreerd.

Bufferindeling voor MoE-dispatch, met een weergave van hoe tokendata en routeringsmetadata worden gerangschikt en uitgewisseld tussen twee ranks.

De combine-kernels zijn aanzienlijk eenvoudiger, aangezien de routeringsinformatie voor hen al beschikbaar is vanaf de dispatch-fase. Ze kunnen de tokens in één schrijfbewerking per peer kopiëren, waarna ze uit de ontvangstbuffer lezen om het gewogen gemiddelde van de tokens naar de output-tensoren te berekenen.

Voor NVLink worden de verzend- en ontvangstbuffers binnen een node toegewezen aan de peers, zodat ze direct tokens van naburige ranks kunnen pushen of pullen.

De onderstaande figuren geven details over de timing en de interacties tussen de GPU, CPU en NIC, en tonen één iteratie van dispatch en combine afgewisseld met de kernels van het model.

Het werk begint bij dispatch-send, waarbij de kernel de routeringsinformatie aggregeert en per-expert tokentellingen berekent via atomaire accumulatie in gedeeld geheugen binnen één blok. Zodra de tokentellingen beschikbaar zijn, worden ze gekopieerd naar het hostgeheugen en wordt de proxy onmiddellijk op de hoogte gesteld om te beginnen met het verzenden van de routeringsinformatie naar alle andere peer-ranks. Parallel hieraan blijft dispatch de cumulatieve som van de tokentellingen berekenen, waarbij de offsets worden geïdentificeerd die nodig zijn om tokens in aaneengesloten verzendbuffers per peer in te pakken. Zodra de offsets beschikbaar zijn, worden alle tokens die via EFA moeten worden verzonden in de verzendbuffers verpakt. De proxy krijgt bericht om de schrijfbewerkingen in de verzendbuffers van inter-node peers te plaatsen. Intra-node pushen de kernels ook onmiddellijk tokens in privé-buffers binnen dezelfde node via NVLink.

Timing en coördinatie voor MoE-dispatch, met gefaseerde GPU-, CPU- en NIC-bewerkingen en hun synchronisatie.

Timing en coördinatie voor MoE-combine, met GPU-, CPU- en NIC-bewerkingen en hun synchronisatie.

Zodra een rank routeringsinformatie van alle peers ontvangt, centraliseert deze deze en berekent de verschillende offsets die nodig zijn om tokens voor dispatch en combine te identificeren. Dit werk wordt op de CPU gedaan terwijl er al overdrachten in behandeling zijn, dus het is niet latentiekritisch. Wanneer alle offsets bekend zijn, worden schrijfbewerkingen voor de rest van de tokens, die de grootte van de privé-buffer overschrijden, geplaatst en is de verzendfase van dispatch voltooid.

Aan de ontvangende kant haalt de kernel eerst tokens uit dezelfde node die de capaciteit van de privé-buffers via NVLink overschrijden. Vervolgens blokkeert het totdat de host-proxy de voltooiing van alle overdrachten signaleert. Zodra het signaal is ontvangen, worden tokens vanuit de ontvangstbuffers verplaatst naar de output-tensor die aan de groep GEMM wordt doorgegeven. Omdat buffers worden hergebruikt tussen dispatch en combine, zorgen zowel NVLink- als RDMA-barrières ervoor dat de ontvangstbuffers zijn gelezen voordat de volgende kernel ze kan overschrijven.

De combine-kernels volgen een ongeveer vergelijkbaar patroon: tokens verplaatsen, de proxy signaleren en wachten op voltooiing. De metadata voor het plaatsen van hun schrijfbewerkingen is echter direct beschikbaar en alle externe offsets zijn bekend, dus alle tokens die intra-node worden overgedragen, worden via NVLink naar de ontvangers gepusht en de inter-node tokens worden overgedragen via een enkele schrijfbewerking. Voordat de resultaten worden geretourneerd, accumuleert de kernel de tokens uit de ontvangstbuffers en schrijft ze in de output-tensoren.

Om onze implementatie te prototypen, hebben we aanvankelijk de TransferEngine hergebruikt die was gebouwd voor disaggregate prefill en decode om RDMA-schrijfbewerkingen uit te voeren en te volgen. In tegenstelling tot disaggregated decoding waarbij elke node een groot aantal schrijfbewerkingen naar dezelfde peer plaatst, plaatst MoE-dispatch echter kleine overdrachten naar alle peers in één bewerking. Out-of-the-box waren de prestaties niet optimaal en liep onze implementatie op ConnectX-7 ongeveer 20us achter op DeepEP. Bijgevolg hebben we de TransferEngine gespecialiseerd om twee bewerkingen te ondersteunen: scatter en barrier, om de interacties met een groot aantal peers te optimaliseren.

Zowel scatter- als barrier-bewerkingen stellen peer-groepen bloot die vooraf kunnen worden geregistreerd om de datastructuren te vullen die door de onderliggende transportimplementatie (Work Request (WR) templating) worden gebruikt. Scatter kopieert verschillende slices van dezelfde bronbuffer naar verschillende offsets op alle leden van een peer-groep. Barrier signaleert alle peers via een immediate. Het enqueuen en verwerken van deze bewerkingen in bulk bespaart waardevolle microseconden op het kritieke pad. Op EFA, waar twee 200 Gbps NIC's aggregeren om 400 Gbps bandbreedte te leveren, shardden we groep-peers over de twee NIC's in plaats van de bytes van elke overdracht te sharden.

Hoewel we de TransferEngine aanvankelijk uitsluitend op EFA ontwikkelden met behulp van libfabric, werken de MoE-kernels op ConnectX-7 zonder één regel code te wijzigen na toevoeging van ConnectX-7-ondersteuning in de TransferEngine met behulp van libibverbs. Vergeleken met SRD-gebaseerde EFA vereist ConnectX-7 extra werk voor verbindingsinstelling en peer-beheer. We gebruiken een Unreliable Datagram (UD) queue pair om Reliable Connection (RC) te bootstrappen. We maken twee queue pairs voor elke externe peer, één voor tweezijdige SEND/RECV en één voor eenzijdige WRITE_IMM, omdat zowel RECV als WRITE_IMM de RECV-bewerking in volgorde verbruiken terwijl de twee bewerkingen verschillende buffergroottes vereisen. Bovendien versoepelen we de volgorde van PCIe-transacties tussen ConnectX-7 en GPU-geheugen. Net als bij EFA implementeren we de WR-templating-optimalisatie. Bovendien gebruiken we bij het verzenden van meerdere kleine berichten naar dezelfde peer WR-chaining om de berichtfrequentie te verbeteren.

Aangezien het werk aan de GPU-zijde identiek is bij beide adapters, hebben we de onderliggende transportimplementatie tot het uiterste gedreven om de prestaties van DeepEP te overtreffen, om te benadrukken dat het gat tussen EFA en ConnectX-7 MoE-routering wordt bepaald door de mogelijkheden van de onderliggende hardware.

We hebben de dispatch- en combine-latentie geëvalueerd op H200-systemen uitgerust met ConnectX-7- en EFA-adapters op de standaard DeepSeek-V3- en Kimi-K2-configuratie met een verborgen dimensie van 7168, block-scaled fp8-dispatch en bf16-combine. We overwegen een batchgrootte van 128 voor dispatch en 4096 voor combine. Op ConnectX-7 vergelijken we met DeepEP en onze eerdere NVSHMEM-gebaseerde pplx-kernels, waarbij we zowel GPUDirect Async (IBGDA) als een host-proxy (IBRC) gebruiken. We bereiken gecombineerde latenties van 459us, 582us en 692us op EP16, EP32 en EP64, respectievelijk. We presteren beter dan andere op EFA gebaseerde kernels die afhankelijk zijn van generieke proxy's, zoals UCCL-EP, die 519us, 966us en 1159us bereiken.

Dispatch- en combinatielatentie voor EP64 MoE-decode, met een vergelijking van de getoonde implementaties en communicatiepaden.

Bij decode bereiken we, ondanks dat onze kernels enkele microseconden trager zijn dan DeepEP, algeheel lagere latenties omdat onze combine-kernels aanzienlijk sneller zijn. Hoewel we een host-proxy gebruiken, overtreffen we vanwege onze bulkoverdrachten de prestaties van alle IBGDA-gebaseerde proxyloze implementaties. We laten ook zien dat proxy-implementatie ertoe doet, aangezien onze gespecialiseerde proxy een orde van grootte sneller is dan de generieke NVSHMEM-implementatie. Ondanks het feit dat EFA twee keer zo traag is als ConnectX-7 bij het overdragen van 256 KB-pakketten, die dicht bij de typische pakketgrootte liggen die door MoE-routering wordt uitgewisseld, is het gat tussen de twee adapters over het algemeen minder substantieel.

Dispatch- en combinatielatentie voor EP64 MoE-prefill, met een vergelijking van de pplx- en DeepEP-communicatiepaden.

Hoewel onze decode-geoptimaliseerde kernels kunnen schalen naar prefill-batchgroottes zonder enige extra aanpassingen, lopen de prestaties achter op DeepEP. Dit komt voornamelijk door het feit dat we tokens alleen aan de ontvangerzijde van combine accumuleren, terwijl DeepEP accumuleert in de zender om de hoeveelheid overgedragen gegevens te verminderen. Deze optimalisatie kan echter gevolgen hebben voor de nauwkeurigheid. Bovendien, aangezien de kernels gegevens in één schrijfbewerking overdragen, gebruiken we meer geheugen voor verzend- en ontvangstbuffers.

Om de levensvatbaarheid van onze kernels te bewijzen, benchmarken we de end-to-end prestaties van onze eigen inferentie-engine die DeepSeek-V3 (671B parameters) en Kimi-K2 (1T parameters) bedient op meerdere p5en-instanties uitgerust met Elastic Fabric Adapters (EFA) en H200 GPU's. We voeren het model uit in draft verification-modus en rapporteren de latentie uitgaande van een acceptatiepercentage van 80% met behulp van één MTP-laag.

Op DeepSeek-V3 vergelijken we met een EP=8 intranode-baseline die efficiënte NVLink-only dispatch- en combine-kernels uitvoert. Inter-node implementaties evenaren of overtreffen in het algemeen de prestaties van de single-node implementatie, met name bij middelgrote batches, wat aantoont dat schalen over nodes heen verbeteringen in de algehele doorvoer oplevert.

Op Kimi-K2 maken de kernels serveren mogelijk. Vanwege de omvang kan het model niet in dezelfde node passen met Data Parallelism (DP) op een H200-node, dus is een inter-node implementatie vereist om het model te bedienen. Omdat het model minder attention heads heeft, bereiken we betere latenties dan DeepSeek-V3 en kunnen we het model op EFA bedienen met een levensvatbare latentie.

DeepSeek-V3-doorvoer per GPU versus uitvoersnelheid voor EP32-, EP16- en EP8-configuraties.

Kimi-K2-doorvoer per GPU versus uitvoersnelheid voor EP32- en EP16-configuraties.

We hebben ook micro-batching overwogen, hoewel dit de latenties alleen verbeterde bij zeer hoge batchgroottes of bij gebruik van meer dan één draft-token. Om te profiteren van het splitsen van gegroepeerde GEMM-problemen in twee, moeten token-batchgroottes voldoende dicht zijn. Het grootste deel van de tijd wordt nog steeds besteed aan MoE-dispatch en -combine, wat aangeeft dat verdere verbeteringen aan de kernels de latenties dienovereenkomstig zullen verbeteren en micro-batching zullen vergemakkelijken.

We werken nauw samen met AWS-engineurs om de prestaties op EFA te verbeteren. Van onze kant zullen we updates van libfabric volgen, wat de overhead van het datavlak zal verminderen. Daarnaast zullen we experimenteren met efa-direct om de overhead van de gebruikersruimte verder te verminderen en direct met de hardware te interageren.

Technisch rapport DeepSeek R1

Technisch rapport Kimi K2

DeepEP

UCCL-EP

Verbetering van netwerkprestaties van HPC-systemen met NVIDIA Magnum IO NVSHMEM en GPUDirect Async

BibTeX

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

Introductie

Inferentie over InfiniBand

Op NVSHMEM gebaseerde pplx-kernels

MoE all-to-all over EFA

Verzend- en ontvangstbuffers

Dispatch en Combine

TransferEngine voor MoE-routering

Evaluatie

Dispatch/Combine-latentie

Kimi-K2 / DeepSeek-V3-doorvoer

Toekomstig werk

Referenties

Citaat