Gør modeller med billioner af parametre mulige på AWS EFA

Hos Perplexity bruger vi de bedste modeller til vores produkt, vores API'er og vores forskningsteams. Store open-source Mixture-of-Experts-modeller som Kimi-K2 udgør særlige udfordringer, da de største inferensnoder med 8x NVIDIA H200 GPU'er ikke effektivt kan rumme dem, hvilket nødvendiggør udrulninger på tværs af flere noder. Vi præsenterer et sæt kerner til ekspertparallellisme, som opnår state-of-the-art latenser på ConnectX-7 og overgår ydeevnen af DeepEP. De samme kerner er desuden de første til at opnå levedygtige latenser på AWS Elastic Fabric Adapter (EFA), hvilket muliggør udrulning af modeller med billioner af parametre.

Prøv vores kerner på GitHub, og læs hele forskningsartiklen på arXiv.

Mixture-of-Experts (MoE) er opstået som standardarkitekturen til skalering til hundreder af milliarder eller billioner af parametre, samtidig med at den tilbyder rimelig latens ved inferenstid. MoE-modeller erstatter det tætte lag i en transformer med et sæt eksperter og et lille lineært routingslag, der afgør, hvilke eksperter et token skal multipliceres med. Ved at gruppere tokens, der er tilknyttet den samme ekspert, bliver problemet sparsomt og trivielt paralleliserbart, da forskellige eksperter nu kan hostes på forskellige GPU'er fordelt på flere noder.

I modsætning til andre former for parallellisme, såsom Tensor Parallelism (TP) og Data Parallelism (DP), som let kan implementeres ved hjælp af kollektive kommunikationsbiblioteker som torch.distributed og NCCL, involverer MoE-routing sparsom peer-to-peer-kommunikation. En dispatch-kerne skal opdele sættet af indgående tokens og sende dem videre til de rækker, der hoster eksperterne, mens dens duale problem, combine-kernen, skal modtage deres behandlede modstykker og beregne deres vægtede gennemsnit. Disse overførsler knytter sig ikke til kollektiver og kræver tilpassede kerner for at opnå lav latens.

Hos Perplexity har vi udviklet både bportable inter-node-kerner og specialiserede intra-node-kerner og præsenteret vores arbejde mere dybdegående i en tidligere artikel om MoE-kommunikation. Vi viste også i et andet blogindlæg, at udrulning på tværs af flere noder kan give lavere latens og højere gennemløb for store MoE-modeller.

I denne artikel dykker vi ned i implementeringen af vores nye inter-node-kerner, som i høj grad forbedrer ydeevnen i vores tidligere NVSHMEM-baserede arbejde. Ud over at fremvise state-of-the-art-resultater på ConnectX-7 viser vi også, at kernerne i vores inferensruntime effektivt kan bruges til at betjene de største open-source-modeller.

Tensorparallellisme og ekspertparallellisme kan skalere effektivt inden for en enkelt node, da NVLink tilbyder høj gennemløb (900 GB/s) ved mikrosekunds latens. Mens NVIDIA GB200- og GB300-systemer er udstyret med tilstrækkelig hukommelse til at rumme vægte og KV-cache i modeller med billioner af parametre, er udrulninger på AWS p5en-forekomster med op til 8 H200 GPU'er begrænset til 1120 GB HBM, som skal opdeles mellem både modelvægte og KV-cacher. Følgelig nødvendiggør visse modeller udrulning på tværs af noder over InfiniBand.

Selvom InfiniBand stadig er ekstremt hurtigt, tilbyder det kun 400 Gbps (50 GB/s) i gennemløb, hvilket tilføjer ti til hundredevis af mikrosekunder i latens for MoE-routing. Heldigvis kan de tilføjede latenser skjules ved hjælp af delte eksperter, mikrobatchning og overlapning af beregning og kommunikation for at levere levedygtigt gennemløb til en konkurrencedygtig pris. DeepSeek banebrød denne teknik via DeepSeek-R1, udrullet på klynger af H800 GPU'er og NVIDIA ConnectX-7 InfiniBand-adaptere.

På AWS-klynger er situationen anderledes på grund af brugen af deres tilpassede Elastic Fabric Adapters (EFA). Selvom disse leverer et topgennemløb på 400 Gbps på de kollektiver, der typisk bruges i træningsarbejdsbelastninger, falder de en smule til kort i forhold til ConnectX-7 på de meddelelsestørrelser, der udveksles under MoE-dispatch og -combine. Desuden understøtter EFA ikke GPUDirect Async, hvilket kræver en CPU-proxtråd, der bygger bro mellem GPU'en og NIC'et for at igangsætte overførsler. Ud over den øgede kompleksitet tilføjer de ekstra PCIe-transaktioner mikrosekunder i overhead til alle overførsler. Rammeværk som NVSHMEM giver kerneforfattere mulighed for at levere enhedsuafhængige implementeringer, men de generiske API'er medfører overhead, der gør udrulninger ineffektive med hensyn til både omkostninger og latens.

Vores tidligere kerner abstraherede den underliggende transportimplementering væk gennem NVSHMEM. De var afhængige af nvshmemx_putmem_signal_nbi_warp til at overføre tokens ét ad gangen ved hjælp af atomare tællere til at synkronisere parter. Selvom dette tilbød rimelig ydeevne på ConnectX-7 NIC'er med IBGDA, var den proxybaserede ConnectX-7 IBRC-implementering betydeligt langsommere, mens EFA-implementeringen ikke formåede at levere latenser under et millisekund.

Vi brugte kernerne som testmiljø for den underliggende hardware, samtidig med at vi fastsatte den bedst egnede grænseflade til vores modeller. Vi opstillede den tese, at de hyppige overførsler og interaktioner mellem kernerne og NIC'erne eller trådene forhindrede ydeevne. Desuden var vores implementering ikke afhængig af andre garantier fra hardwaren end pålidelighed, hvilket antydede, at vi kunne fremskynde tingene ved at reducere overheadet af synkroniseringsprimitiver.

Efter disse observationer udformede vi vores EFA-baserede kerner, der eksponerer den samme grænseflade og optimerer interaktionerne mellem enhedskernen og værts-proxyen under kølerhjelmen for maksimal ydeevne.

RDMA- og NVLink-hybridkommunikation med en hostproxy, der viser initierende og modtagende ranks, proxyarbejdere og NIC'er.

De nye kerner følger en hybrid CPU-GPU-arkitektur, hvor GPU-kerner interagerer med modellen på enheden, og en tråd, der kører på værts-CPU'en, håndterer interaktioner med NIC'et. De nye kerner genbruger den samme TransferEngine, som vi byggede til KV-cache-overførsler.

MoE-routing udføres af en dispatch-kerne, som modtager tokens og de per-token-eksperter, der skal routes til fra modellen, og sender dem over netværket til de rækker, der hoster eksperterne. Når de grupperede GEMM-kerner har multipliceret tokens med ekspertvægtene, overfører combine-kernen tokens til deres kilderækker og beregner deres vægtede gennemsnit, før de sendes tilbage til modellen.

Både dispatch- og combine-kernerne er opdelt i en afsender- og en modtagerdel for at lette overlapning og mikrobatchning. Afsenderen blander tokens ind i afsendelsesbuffere, så de kan skrives til parter via en enkelt RDMA-skrivehandling. Efter udførelse af afsenderen kan GPU'en udføre andet arbejde (såsom delt ekspertmultiplikation), mens overførslerne over netværket afventer i baggrunden. Modtagerkernerne blokerer, indtil alle overførsler er modtaget. I baggrunden håndterer en arbejderstråd på værtssiden overførslerne: den adviseres af GPU'en, når afsendelsesbufferne er blevet fyldt, hvorefter den udsender RDMA-skrivningerne. Når NIC'et signalerer fuldførelse, ophæver det blokeringen af modtageren, som returnerer tokens til modellen.

Kernerne er tæt koblet til værten og kommunikerer med hinanden gennem både unified memory og GDRCopy. CUDA unified memory gør det muligt at tilknytte enhedshukommelse eller værtshukommelse til begge adresserum, hvilket gennemsigtigt håndterer kommunikation via PCIe.

Selvom unified memory er velegnet til store overførsler, udføres lav-latens-pooling via GDRCopy. GDRCopy bruger GPUDirect RDMA til at læse fra og skrive til GPU-hukommelse ved mikrosekunds latens.

Ud over RDMA udnytter kernerne også NVLink til at udveksle tokens inden for en node. Dette kan aflaste en betydelig brøkdel af trafikken, op til en fjerdedel for EP32 og op til en ottendedel for EP64, hvilket opnår generelt lavere latenser.

Designet af afsendelsesbufferne til dispatch-kernen er afgørende for at opnå levedygtige latenser, der involverer værts-proxyen. Hvis der er N rækker, der hver dispatchet T tokens, som hver især routes til R eksperter ud af E, kan en række modtage op til N * T * max(R, E / N) tokens i værst tænkelige fald, hvis de alle routes til den. Desuden kan denne minimumsgrænse kun opnås, når afsenderne lægger deres skrivninger sammenhængende, da private modtagerbuffere ville opbruge for meget hukommelse. For at afsendere kan lægge deres skrivninger ud i rækkefølge uden konflikter, skal fuldstændig kendskab til routinginformationen distribueres til hver række. Følgelig udveksler vores dispatch-kerner først per-ekspert-tokenantal, hvorefter de hver især kan bestemme, hvor deres skrivninger går hen, og fra hvilken forskydning de modtager tokens fra hver række.

Simpel udveksling af routinginformation ville resultere i unødvendig høj latens, da hver række ville være nødt til at vente på, at skrivningerne ankom, før den udsendte sine egne skrivninger. I stedet reserverer vi en lille mængde privat Space pr. afsender på hver modtagerrække og dispatchet nogle af tokens sammen med skrivningerne af routinginformationen. Dette sikrer fuld båndbreddeudnyttelse under størstedelen af dispatch-processen. Når ruter er udvekslet, overføres de resterende tokens via en enkelt skrivning pr. fjernpart. Udvekslingen af routinginformation og de faststørrelsesbuffere, sammen med de overførsler, der udfylder dem, er illustreret i figuren nedenfor.

Bufferlayout for MoE-dispatch, der viser, hvordan token-data og routingmetadata er arrangeret og udvekslet mellem to ranks.

Combine-kernerne er væsentligt enklere, da routinginformationen allerede er tilgængelig for dem fra dispatch-stadiet. De kan kopiere tokens over i en enkelt skrivning pr. part, hvorefter de læser fra modtagerbufferen for at beregne det vægtede gennemsnit af tokens ind i output-tensorerne.

For NVLink er afsendelses- og modtagerbufferne inden for en node knyttet til parterne, så de direkte kan skubbe eller trække tokens fra naborækker.

Figurerne nedenfor beskriver tidsplanen og interaktionerne mellem GPU, CPU og NIC og viser én iteration af dispatch og combine flettet sammen med modellens kerner.

Arbejdet starter i dispatch-send, hvor kernen aggregerer routinginformationen og beregner per-ekspert-tokenantal via atomar akkumulering i delt hukommelse inden for en enkelt blok. Når tokenantallet er tilgængeligt, kopieres det til værtshukommelsen, og proxyen underrettes med det samme om at begynde at dispatchet routinginformationen til alle andre part-rækker. Parallelt hermed fortsætter dispatch med at beregne den kumulative sum af tokenantallet og identificerer de forskydninger, der kræves for at pakke tokens i sammenhængende afsendelsesbuffere pr. række. Når forskydningerne er tilgængelige, pakkes alle tokens, der skal sendes via EFA, i afsendelsesbufferne. Proxyen underrettes om at poste skrivningerne i inter-node-parternes afsendelsesbuffere. Intra-node skubber kernerne også øjeblikkeligt tokens ind i private buffere inden for samme node via NVLink.

Tidsplan og koordinering for MoE-dispatch, der viser trinvis GPU-, CPU- og NIC-operationer samt deres synkronisering.

Tidsplan og koordinering for MoE-combine, der viser GPU-, CPU- og NIC-operationer samt deres synkronisering.

Når en række modtager routinginformation fra alle parter, centraliserer den den og beregner forskellige forskydninger, der er nødvendige for at identificere tokens til dispatch og combine. Dette arbejde udføres på CPU'en, mens overførsler allerede afventer, så det er ikke latenskritisk. Når alle forskydninger er kendt, postes efterfølgende skrivninger for resten af de tokens, der overstiger den private bufferstørrelse, og afsendelsesfasen af dispatch er fuldført.

I modtagerenden trækker kernen først tokens fra den samme node, som overstiger de private buffers kapacitet via NVLink. Derefter blokerer den, indtil værts-proxyen signalerer, at alle overførsler er fuldført. Når signalet er modtaget, flyttes tokens fra modtagerbufferne til den output-tensor, der sendes videre til gruppe-GEMM. Da buffere genbruges mellem dispatch og combine, sikrer både NVLink- og RDMA-barrierer, at modtagerbufferne er blevet læst, før den næste kerne tillades at overskrive dem.

Combine-kernerne følger et omtrent lignende mønster, hvor tokens flyttes rundt, proxyen adviseres, og der ventes på fuldførelse. Metadataene til at poste deres skrivninger er dog let tilgængelige, og alle fjerntliggende forskydninger er kendte, så alle tokens, der overføres intra-node, skubbes via NVLink til modtagerne, og inter-node-tokens overføres via en enkelt skrivning. Før resultaterne returneres, akkumulerer kernen tokens fra modtagerbufferne og skriver dem i output-tensorerne.

For at prototype vores implementering genbrugte vi oprindeligt den TransferEngine, der blev bygget til opdelt prefill og dekodning, for at udstede og spore RDMA-skrivninger. I modsætning til opdelt dekodning, hvor hver node poster et stort antal skrivninger til den samme part, poster MoE-dispatch imidlertid små overførsler til alle parter i én handling. Ydeevnen var ikke optimal fra start, og vores implementering lå ca. 20 us efter DeepEP på ConnectX-7. Følgelig specialiserede vi TransferEngine til at understøtte to handlinger, scatter og barrier, for at optimere interaktioner med et stort antal parter.

Både scatter- og barrier-handlinger eksponerer partgrupper, som kan forhåndsregistreres for at forhåndsbefolke de datastrukturer, der bruges af den underliggende transportimplementering (Work Request (WR)-skabeloner). Scatter kopierer forskellige udsnit fra den samme kildebuffer til forskellige forskydninger på alle medlemmer i en partgruppe. Barrier signalerer til alle parter via en immediate. Kø-sættelse og håndtering af disse handlinger i bulk sparer værdifulde mikrosekunder på den kritiske sti. På EFA, hvor to 200 Gbps NIC'er aggregerer for at levere 400 Gbps båndbredde, sharder vi gruppe-parter på tværs af de to NIC'er i stedet for at sharde bytes for hver overførsel.

Selvom vi oprindeligt udviklede TransferEngine udelukkende på EFA ved hjælp af libfabric, fungerer MoE-kernerne på ConnectX-7 uden at ændre en eneste linje kode efter tilføjelse af ConnectX-7-understøttelse i TransferEngine ved hjælp af libibverbs. Sammenlignet med SRD-baseret EFA (Scalable Reliable Datagram) kræver ConnectX-7 ekstra arbejde til opsætning af forbindelser og parthåndtering. Vi bruger et UD-køpar (Unreliable Datagram) til at boottrapme RC (Reliable Connection). Vi opreter to køpar for hver fjernpart, én til tovejs SEND/RECV og én til envejs WRITE_IMM, fordi både RECV og WRITE_IMM forbruger RECV-handlinger i rækkefølge, mens de to handlinger kræver forskellige bufferstørrelser. Desuden lempes rækkefølgen af PCIe-transaktioner mellem ConnectX-7 og GPU-hukommelse. Ligesom med EFA implementerer vi WR-skabelonoptimeringen. Derudover bruger vi WR-kædning til at forbedre meddelelsens hastighed, når vi sender flere små meddelelser til den samme part.

Da GPU-siden er identisk på tværs af begge adaptere, skubbede vi den underliggende transportimplementering til dens grænser for at overgå ydeevnen af DeepEP og fremhæve, at kløften mellem EFA- og ConnectX-7 MoE-routing bestemmes af den underliggende hardwares kapaciteter.

Vi evaluerede dispatch- og combine-latenser på H200-systemer udstyret med ConnectX-7- og EFA-adaptere på standard DeepSeek-V3- og Kimi-K2-konfigurationen med en skjult dimension på 7168, blokskaleret fp8-dispatch og bf16-combine. Vi tager udgangspunkt i en batchstørrelse på 128 for dispatch og 4096 for combine. På ConnectX-7 sammenligner vi med DeepEP og vores tidligere NVSHMEM-baserede pplx-kerner ved hjælp af både GPUDirect Async (IBGDA) og en værts-proxy (IBRC). Vi opnår kombinerede latenser på henholdsvis 459 us, 582 us og 692 us på EP16, EP32 og EP64. Vi overgår andre EFA-baserede kerner, der er afhængige af generiske proxyer, såsom UCCL-EP, som opnår 519 us, 966 us og 1159 us.

Dispatch- og combine-latens for EP64 MoE-decode, der sammenligner de viste implementeringer og kommunikationsstier.

Ved dekodning opnår vi generelt lavere latenser, da vores combine-kerner er væsentligt hurtigere, på trods af at vores kerner er et par mikrosekunder langsommere end DeepEP. Selvom vi bruger en værts-proxy, overgår vi på grund af vores bulkoversendelser ydeevnen af alle IBGDA-baserede implementeringer uden proxy. Vi viser også, at proxy-implementeringen betyder noget, da vores stærkt specialiserede proxy er en størrelsesorden hurtigere end den generiske NVSHMEM-implementering. På trods af at EFA er dobbelt så langsom som ConnectX-7 ved overførsel af 256 KB-pakker, som ligger tæt på den typiske pakkestørrelse udvekslet af MoE-routing, er forskellen mellem de to adaptere generelt mindre væsentlig.

Dispatch- og combine-latens for EP64 MoE-prefill, der sammenligner pplx- og DeepEP-kommunikationsstier.

Selvom vores dekodningsoptimerede kerner kan skalere til prefill-batchstørrelser uden yderligere tilpasninger, halter ydeevnen bagefter DeepEP. Dette skyldes primært, at vi kun akkumulerer tokens på modtagersiden af combine, hvorimod DeepEP akkumulerer i afsenderen for at reducere mængden af overført data. Denne optimering kan dog have konsekvenser for nøjagtigheden. Da kernerne overfører data i én skrivning, udnytter vi desuden mere hukommelse til afsendelses- og modtagerbuffere.

For at bevise vores kerners levedygtighed benchmark-tester vi end-to-end-ydeevnen for vores interne inferensmotor, der betjener DeepSeek-V3 (671B parametre) og Kimi-K2 (1T parametre) på flere p5en-forekomster udstyret med Elastic Fabric Adapters (EFA) og H200 GPU'er. Vi kører modellen i udkast-verificeringstilstand (draft verification mode) og rapporterer latens under forudsætning af en acceptrate på 80 % ved hjælp af et enkelt MTP-lag.

På DeepSeek-V3 sammenligner vi med en EP=8-baselinje inden for samme node, der kører effektive dispatch- og combine-kerner udelukkende med NVLink. Udrulninger på tværs af noder matcher eller overgår generelt ydeevnen af udrulningen på en enkelt node, især ved mellemstore batches, hvilket viser, at skalering på tværs af noder giver forbedringer i det samlede gennemløb.

På Kimi-K2 muliggør kernerne betjening. På grund af sin størrelse kan modellen ikke være der inden for den samme node ved hjælp af Data Parallelism (DP) på en H200-node, hvorfor en udrulning på tværs af noder er påkrævet for at betjene modellen. Da modellen har færre opmærksomhedshoveder (attention heads), opnår vi bedre latenser end DeepSeek-V3, og vi kan betjene modellen på EFA ved en levedygtig latens.

DeepSeek-V3 per-GPU-throughput versus outputhastighed for EP32-, EP16- og EP8-konfigurationer.

Kimi-K2 per-GPU-throughput versus outputhastighed for EP32- og EP16-konfigurationer.

Vi har også overvejet mikrobatchning, om end det kun forbedrede latenser ved meget store batchstørrelser eller ved brug af mere end ét udkast-token. For at drage fordel af at opdele grupperede GEMM-problemer i to skal token-batchstørrelserne være tilstrækkeligt tætte. Hovedparten af tiden bruges stadig på MoE-dispatch og -combine, hvilket indikerer, at yderligere forbedringer af kernerne vil forbedre latenserne og lette mikrobatchning.

Vi arbejder tæt sammen med AWS-ingeniører om at forbedre ydeevnen på EFA. På vores side vil vi følge opdateringer til libfabric, som vil reducere dataplane-overheadet. Derudover vil vi eksperimentere med efa-direct for yderligere at reducere userspace-overhead og interagere direkte med hardwaren.

DeepSeek R1 Technical Report

Kimi K2 Technical Report

DeepEP

UCCL-EP

Improving Network Performance of HPC Systems Using NVIDIA Magnum IO NVSHMEM and 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}
}

Introduktion

Inferens over InfiniBand

NVSHMEM-baserede pplx-kerner

MoE all-to-all over EFA

Afsendelses- og modtagerbuffere

Dispatch og Combine

TransferEngine til MoE-routing

Evaluering

Dispatch/Combine-latens

Kimi-K2 / DeepSeek-V3 Gennemløb

Fremtidigt arbejde

Referencer

Citering