Aktivera modeller med biljoner parametrar på AWS EFA

På Perplexity använder vi de bästa modellerna för vår produkt, våra API:er och våra forskningsteam. Stora öppen källkodsmodeller av typen Mixture-of-Experts, som Kimi-K2, medför särskilda utmaningar eftersom de största inferensnoderna med 8x NVIDIA H200-GPU:er inte effektivt kan rymma dem, vilket gör distributioner över flera noder nödvändiga. Vi presenterar en uppsättning kärnor (kernels) för expertparallelism som uppnår banbrytande latenser på ConnectX-7 och överträffar prestandan hos DeepEP. Samma kärnor är dessutom de första som uppnår gångbara latenser på AWS Elastic Fabric Adapter (EFA), vilket möjliggör modelldistributioner på biljoner parametrar.

Testa våra kärnor på GitHub och läs hela forskningsartikeln på arXiv.

Introduktion

Mixture-of-Experts (MoE) har vuxit fram som standardarkitektur för att skala upp till hundratals miljarder eller biljoner parametrar samtidigt som det erbjuder rimlig latens vid inferenstid. MoE-modeller ersätter ett transformatorns (transformer) tätt lagerskikt (dense layer) med en uppsättning experter och ett litet linjärt routningslager som avgör vilken expert en token ska multipliceras med. Genom att gruppera de tokens som är associerade med samma expert blir problemet glest (sparse) och enkelt parallelliserbart, eftersom olika experter nu kan hostas på olika GPU:er utspridda över flera noder.

Till skillnad från andra former av parallelism, såsom Tensor Parallelism (TP) och Data Parallelism (DP) som enkelt kan implementeras med hjälp av kollektiva kommunikationsbibliotek som torch.distributed och NCCL, innefattar MoE-routning gles peer-to-peer-kommunikation. En dispatch-kärna måste dela upp uppsättningen inkommande tokens och skicka dem till de rader (ranks) som hostar experterna, medan dess dolda motsvarighet, combine-kärnan, måste ta emot sina bearbetade motsvarigheter och beräkna deras vägda genomsnitt. Dessa överföringar mappar inte till kollektiv och kräver anpassade kärnor för att uppnå låg latens.

På Perplexity har vi utvecklat både bärbara kärnor mellan noder (inter-node) och specialiserade kärnor inom noden (intra-node), och presenterar vårt arbete mer djupgående i en tidigare artikel om MoE-kommunikation. Vi visade också i ett annat blogginlägg att distribution över flera noder kan ge lägre latens och högre genomströmning för stora MoE-modeller.

I den här artikeln djupdyker vi i implementeringen av våra nya kärnor mellan noder, vilket avsevärt förbättrar prestandan jämfört med vårt tidigare NVSHMEM-baserade arbete. Förutom att visa upp banbrytande resultat på ConnectX-7 visar vi också att våra kärnor, inom vår inferenskörningstid, effektivt kan användas för att serva de största modellerna med öppen källkod.

Inferens över InfiniBand

Tensorparallelism och expertparallelism kan skalas effektivt inom en enskild nod, eftersom NVLink erbjuder hög genomströmning (900 GB/s) vid mikrosekunders latens. Även om NVIDIA GB200- och GB300-systemen är utrustade med tillräckligt med minne för att rymma vikterna (weights) och KV-cacheminnena för modeller med biljoner parametrar, är distributioner på AWS p5en-instanser med upp till 8 H200-GPU:er begränsade till 1120 GB HBM, vilket måste delas upp mellan både modellvikter och KV-cacher. Följaktligen kräver vissa modeller distributioner mellan noder över InfiniBand.

Även om InfiniBand fortfarande är oerhört snabbt erbjuder det bara 400 Gbps (50 GB/s) i genomströmning, vilket lägger till tiotals till hundratals mikrosekunder i latens för MoE-routning. Lyckligtvis kan de tillagda latenserna döljas genom användning av delade experter, mikro-batchning och överlappning mellan beräkning och kommunikation för att leverera gångbar genomströmning till en konkurrenskraftig kostnad. DeepSeek var pionjärer för dessa tekniker via DeepSeek-R1, distribuerat på kluster med H800-GPU:er och NVIDIA ConnectX-7 InfiniBand-adaktorer.

På AWS-kluster ser situationen annorlunda ut på grund av användningen av deras anpassade Elastic Fabric Adapters (EFA). Även om dessa levererar en toppgenomströmning på 400 Gbps för de kollektiv som typiskt används i träningsarbetsbelastningar, når de inte riktigt upp till ConnectX-7 när det gäller de meddelandestorlekar som utbyts under MoE-dispatch och -kombinering. Dessutom saknar EFA stöd för GPUDirect Async, vilket kräver en CPU-proxytråd som överbryggar GPU:n till nätverkskortet (NIC) för att initiera överföringar. Utöver den tillagda komplexiteten lägger de extra PCIe-transaktionerna till mikrosekunder av overhead för alla överföringar. Ramverk som NVSHMEM tillåter kärnskribenter att tillhandahålla enhetsberoende (device-independent) implementeringar, men de generiska API:erna medför en overhead som gör distributioner ineffektiva gällande både kostnad och latens.

NVSHMEM-baserade pplx-kernels

Våra tidigare kärnor abstraherade bort den underliggande transportimplementeringen genom NVSHMEM. De förlitade sig på nvshmemx_putmem_signal_nbi_warp för att överföra tokens en i taget, med hjälp av atomära räknare för att synkronisera parter (peers). Även om detta erbjöd rimlig prestanda på ConnectX-7-nätverkskort med IBGDA, var den proxybaserade ConnectX-7 IBRC-implementeringen betydligt långsammare, medan EFA-implementeringen misslyckades med att leverera latenser under en millisekund.

Vi använde kärnorna som en testbänk för den underliggande hårdvaran, samtidigt som vi fastställde det lämpligaste gränssnittet att exponera för våra modeller. Vi ställde hypotesen att de frekventa överföringarna och interaktionerna mellan kärnorna och nätverkskorten eller proxytrådarna hämmade prestandan. Dessutom förlitade sig vår implementering inte på några andra garantier från hårdvaran än tillförlitlighet, vilket antydde att vi kunde snabba upp saker genom att minska overheaden för synkroniseringsprimitiver.

Med utgångspunkt i dessa observationer utformade vi våra EFA-baserade kärnor som exponerar samma gränssnitt, och optimerade interaktionerna mellan enhetskärnan och värdproxyn (host proxy) under huven för maximal prestanda.

MoE all-to-all över EFA

Förstora

De nya kärnorna följer en hybrid CPU-GPU-arkitektur, där GPU-kärnor interagerar med modellen på enheten och en proxytråd som körs på värd-CPU:n hanterar interaktioner med nätverkskortet. De nya kärnorna återanvänder samma TransferEngine som vi byggde för överföringar av KV-cacher.

MoE-routning utförs av en dispatch-kärna som tar emot tokens och per-token-experter för att dirigera dem från modellen, och skickar dem över nätverket till de rader som hostar experterna. Efter att de grupperade GEMM-kärnorna multiplicerar tokens med expertvikterna, överför combine-kärnan tokens till deras källrader och beräknar deras vägda genomsnitt innan de returneras till modellen.

Både dispatch- och combine-kärnorna är uppdelade i en avsändar- och en mottagarhalva för att underlätta överlappning och mikro-batchning. Avsändaren sorterar tokens i sändbuffertar så att de kan skrivas till parter genom en enda RDMA-skrivoperation. Efter att ha kört avsändaren kan GPU:n utföra annat arbete (såsom delad expertmultiplikation) medan överföringarna över nätverket väntar i bakgrunden. Mottagarkärnorna blockeras tills alla överföringar har tagits emot. I bakgrunden hanterar en arbetartråd (worker thread) på värdsidan överföringarna: den signaleras av GPU:n när sändbuffertarna har fyllts, varefter den skickar iväg RDMA-skrivningarna. När nätverkskortet signalerar slutförande avblockeras mottagaren, som returnerar tokens till modellen.

Kärnorna är tätt kopplade till värden (host) och kommunicerar med varandra genom både enhetligt minne (unified memory) och GDRCopy. CUDA unified memory gör att enhetsminne eller värdminne kan mappas in i bådas adressrymder, vilket transparent hanterar kommunikation via PCIe.

Medan unified memory är lämpligt för bulköverföringar utförs polling med låg latens genom GDRCopy. GDRCopy använder GPUDirect RDMA för att läsa från och skriva till GPU-minne vid mikrosekunders latens.

Utöver RDMA utnyttjar kärnorna även NVLink för att utbyta tokens inom en nod. Detta kan avlasta en betydande del av trafikvolymen, upp till en fjärdedel för EP32 och upp till en åttondel för EP64, vilket ger generellt lägre latenser.

Sänd- och mottagningsbuffertar

Utformningen av sändbuffertarna för dispatch-kärnan är avgörande för att uppnå gångbara latenser som involverar värdproxyn. Om det finns N rader som var och en skickar T tokens, där varje token dirigeras till R experter av E, kan en rad i värsta möjliga fall ta emot upp till N * T * max(R, E / N) tokens om alla dirigeras till den. Dessutom kan denna minimigräns endast uppnås när avsändarna lägger ut sina skrivningar contiguously (sammanhängande), eftersom privata mottagningsbuffertar skulle förbruka för mycket minne. För att avsändare ska kunna lägga ut sina skrivningar efter varandra utan konflikter måste fullständig kännedom om routningsinformationen distribueras till varje rad. Följaktligen utbyter våra dispatch-kärnor först per-expert-tokencounter, varefter var och en kan avgöra vart deras skrivningar tar vägen och från vilken förskjutning (offset) de tar emot tokens från varje rad.

Ett enkelt utbyte av routningsinformation skulle resultera i onödigt höga latenser, eftersom varje rad skulle behöva vänta på att skrivningarna ska anlända innan den skickar några egna skrivningar. Istället reserverar vi en liten mängd privat utrymme per avsändare på varje mottagarrad och skickar några av tokens tillsammans med skrivningarna för routningsinformationen. Detta säkerställer fullständig bandbreddsutnyttjelse under större delen av dispatch-processen. När rutter har utbytts överförs resterande tokens via en enda skrivning per fjärrpart. Utbytet av routningsinformation och buffertarna med fast storlek, tillsammans med de överföringar som fyller dem, illustreras i figuren nedan.

Förstora

Combine-kärnorna är väsentligt enklare, eftersom routningsinformation redan är tillgänglig för dem från dispatch-stadiet. De kan kopiera över tokens i en enda skrivning per part, varefter de läser från mottagningsbufferten för att beräkna det vägda genomsnittet av tokens i utdatatensorerna.

För NVLink mappas sänd- och mottagningsbuffertarna inom en nod in i parterna så att de direkt kan skicka (push) eller hämta (pull) tokens från grannrader.

Dispatch och Combine

Figurerna nedan beskriver tidsinställningen och interaktionerna mellan GPU, CPU och nätverkskort, och visar en iteration av dispatch och combine som varvas med modellens kärnor.

Arbetet börjar i dispatch send, där kärnan aggregerar routningsinformationen och beräknar per-expert-tokencounter via atomär ackumulering i det delade minnet inom ett enda block. När tokencountsen är tillgängliga kopieras de till värdminnet och proxyn meddelas omedelbart om att börja skicka routningsinformationen till alla andra partrader. Parallellt fortsätter dispatch att beräkna den kumulativa summan av tokencountsen, vilket identifierar de förskjutningar som krävs för att packa tokens i sammanhängande per-part-sändbuffertar. När förskjutningarna är tillgängliga packas alla tokens som ska skickas via EFA in i sändbuffertarna. Proxyn meddelas om att posta skrivningarna i sändbuffertarna för noder mellan noder (inter-node). Internt i noden (intra-node) skickar kärnorna också omedelbart tokens till privata buffertar inom samma nod via NVLink.

Förstora

Förstora

När en rad tar emot routningsinformation från alla parter centraliserar den informationen och beräknar olika förskjutningar som är nödvändiga för att identifiera tokens för dispatch och combine. Detta arbete utförs på CPU:n medan överföringar redan väntar, och det är därmed inte latenskritiskt. När alla förskjutningar väl är kända postas skrivningar för resten av de tokens som överstiger den privata buffertstorleken, och sändfasen för dispatch är färdig.

På mottagarsidan hämtar kärnan först tokens från samma nod som överskrider de privata buffertarnas kapacitet via NVLink. Därefter blockeras den tills värdproxyn signalerar att alla överföringar har slutförts. När signalen har tagits emot flyttas tokens om (shuffled) från mottagningsbuffertarna till den utdatatensor som skickats till grupp-GEMM. Eftersom buffertar återanvänds mellan dispatch och combine säkerställer både NVLink- och RDMA-barriärer att mottagningsbuffertarna har lästs innan nästa kärna tillåts skriva över dem.

Combine-kärnorna följer ett någorlunda liknande mönster där tokens flyttas om, proxyn signaleras och man väntar på slutförande. Metadata för att posta deras skrivningar är dock lätt tillgänglig och alla fjärrförskjutningar är kända, vilket gör att alla tokens som överförs inom noden skickas via NVLink till mottagarna och tokens mellan noder överförs via en enda skrivning. Innan resultaten returneras ackumulerar kärnan tokens från mottagningsbuffertarna och skriver dem till utdatatensorerna.

TransferEngine för MoE-routning

För att prototypa vår implementering återanvände vi initialt den Transfer Engine som byggts för disaggregerad prefill och decode i syfte att utföra och spåra RDMA-skrivningar. Till skillnad från disaggregerad avkodning där varje nod postar ett stort antal skrivningar till samma part, postar MoE dispatch dock små överföringar till alla parter i en enda operation. Prestandan var inte optimal direkt ur lådan, och vår implementering låg efter DeepEP på ConnectX-7 med ungefär 20 us. Följaktligen specialiserade vi TransferEngine för att ha stöd för två operationer: scatter och barrier, för att optimera interaktioner med ett stort antal parter.

Båda operationerna, scatter och barrier, exponerar partgrupper som kan förregistreras för att i förväg befolka de datastrukturer som används av den underliggande transportimplementeringen (Work Request (WR)-templating). Scatter kopierar olika delar från samma källbuffert till olika förskjutningar på alla medlemmar i en partgrupp. Barrier signalerar alla parter via en omedelbar (immediate) signal. Att köa och hantera dessa operationer i bulk sparar värdefulla mikrosekunder på den kritiska sökvägen. På EFA, där två 200 Gbps-nätverkskort aggregeras för att tillhandahålla 400 Gbps bandbredd, shardar vi partgrupper över de två nätverkskorten istället för att sharda bytena i varje överföring.

Även om vi initialt utvecklade TransferEngine enbart på EFA med hjälp av libfabric, fungerar MoE-kärnorna på ConnectX-7 utan att ändra en enda rad kod efter att ha lagt till stöd för ConnectX-7 i TransferEngine med hjälp av libibverbs. Jämfört med Scalable Reliable Datagram (SRD)-baserad EFA kräver ConnectX-7 extra arbete för anslutningsetablering och parthantering. Vi använder ett Unreliable Datagram (UD)-köpar för att starta upp Reliable Connection (RC). Vi skapar två köpar för varje fjärrpart, ett för dubbelsidig SEND/RECV och ett för ensidig WRITE_IMM, eftersom både RECV och WRITE_IMM förbrukar RECV-operationer i ordning medan de två operationerna kräver olika buffertstorlekar. Dessutom mjukar vi upp ordningsföljden för PCIe-transaktioner mellan ConnectX-7 och GPU-minne. I likhet med EFA implementerar vi WR-templating-optimeringen. När vi skickar flera små meddelanden till samma part använder vi dessutom WR-kedjning för att förbättra meddelandefrekvensen.

Eftersom arbetet på GPU-sidan är identiskt över båda adaptrarna pressade vi den underliggande transportimplementeringen till dess gränser för att överträffa prestandan hos DeepEP, för att belysa att gapet mellan EFA och ConnectX-7 MoE-routning bestäms av den underliggande hårdvarans kapacitet.

Utvärdering

Latens för Dispatch/Combine

Vi utvärderade latenser för dispatch och combine på H200-system utrustade med ConnectX-7- och EFA-adaptrar på standardkonfigurationen för DeepSeek-V3 och Kimi-K2 med en dold dimension på 7168, blockskalad fp8-dispatch och bf16-kombinering. Vi betraktar en batchstorlek på 128 för dispatch och 4096 för combine. På ConnectX-7 jämför vi med DeepEP och våra tidigare NVSHMEM-baserade pplx-kernels, med användning av både GPUDirect Async (IBGDA) och en värdproxy (IBRC). Vi uppnår kombinerade latenser på 459 us, 582 us respektive 692 us på EP16, EP32 och EP64. Vi presterar bättre än andra EFA-baserade kärnor som förlitar sig på generiska proxyservrar, såsom UCCL-EP, vilka uppnår 519 us, 966 us och 1159 us.

Förstora

Vid avkodning (decode), trots att våra kärnor är några mikrosekunder långsammare än DeepEP, uppnår vi totalt sett lägre latenser eftersom våra combine-kärnor är avsevärt snabbare. Även om vi använder en värdproxy överträffar vi, på grund av våra bulköverföringar, prestandan hos alla IBGDA-baserade implementeringar utan proxy. Vi visar också att proxyimplementeringen spelar roll, eftersom vår högt specialiserade proxy är en storleksordning snabbare än den generiska NVSHMEM-implementeringen. Trots det faktum att EFA är dubbelt så långsam som ConnectX-7 vid överföring av 256 KB-paket, vilka ligger nära den typiska paketstorlek som utbyts av MoE-routning, är gapet mellan de två adaptrarna mindre påtagligt totalt sett.

Förstora

Även om våra avkodningsoptimerade kärnor kan skalas till prefill-batchstorlekar utan några extra justeringar, släpar prestandan efter DeepEP. Detta beror främst på att vi endast ackumulerar tokens på mottagarsidan av combine, medan DeepEP ackumulerar i avsändaren för att minska mängden överförd data. Denna optimering kan dock ha konsekvenser för noggrannheten. Eftersom kärnorna överför data i en enda skrivning använder vi dessutom mer minne för sänd- och mottagningsbuffertar.

Genomströmning för Kimi-K2 / DeepSeek-V3

För att bevisa våra kärnors bärkraft kör vi ett prestandatest (benchmark) på end-to-end-prestandan för vår interna inferensmotor som servar DeepSeek-V3 (671B parametrar) och Kimi-K2 (1T parametrar) på flera p5en-instanser utrustade med Elastic Fabric Adapters (EFA) och H200-GPU:er. Vi kör modellen i utkastverifieringsläge (draft verification mode) och redovisar latens under antagandet av en acceptansgrad på 80 % med ett enda MTP-lager.

På DeepSeek-V3 jämför vi med en EP=8-baslinje inom noden som kör effektiva dispatch- och combine-kärnor som enbart använder NVLink. Distributioner mellan noder matchar eller överträffar totalt sett prestandan för distributionen på en enskild nod, i synnerhet på medelstora batcher, vilket visar att skalning över noder faktiskt ger förbättringar av den totala genomströmningen.

På Kimi-K2 möjliggör kärnorna servering. På grund av dess storlek får modellen inte plats inom samma nod med Data Parallelism (DP) på en H200-nod, varför en distribution mellan noder krävs för att serva modellen. Eftersom modellen har färre uppmärksamhetshuvuden (attention heads) uppnår vi bättre latenser än DeepSeek-V3 och vi kan serva modellen på EFA vid en gångbar latens.

Förstora

Förstora

Vi har även övervägt mikro-batchning, även om det bara förbättrade latenserna vid mycket stora batchstorlekar eller när fler än en utkast-token (draft token) används. För att dra nytta av att dela upp grupperade GEMM-problem i två måste token-batchstorlekarna vara tillräckligt glesa (dense). Huvud delen av tiden läggs fortfarande på MoE dispatch och combine, vilket indikerar att alla ytterligare förbättringar av kärnorna följaktligen kommer att förbättra latenserna och underlätta mikro-batchning.

Framtida arbete

Vi har ett nära samarbete med AWS-ingenjörer för att förbättra prestandan på EFA. På vår sida kommer vi att följa uppdateringar av libfabric som kommer att minska dataplanets overhead. Dessutom kommer vi att experimentera med efa-direct för att ytterligare minska användarrymdens (userspace) overhead och gränssnitta direkt mot hårdvaran.

Referenser

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

Citering

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