Habilitación de modelos de billóns de parámetros en AWS EFA

En Perplexity, utilizamos os mellores modelos para o noso produto, as nosas API e os nosos equipos de investigación. Os grandes modelos abertos de expertos mixtos (Mixture-of-Experts), como Kimi-K2, presentan desafíos particulares, xa que os nós de inferencia máis grandes con 8 GPUs NVIDIA H200 non poden acollelos de forma eficiente, o que fai necesarias implementacións en varios nós. Presentamos un conxunto de núcleos para o paralelismo de expertos que acadan latencias de vangarda en ConnectX-7, superando o rendemento de DeepEP. Os mesmos núcleos son tamén os primeiros en acadar latencias viables en AWS Elastic Fabric Adapter (EFA), o que permite implementacións de modelos de billóns de parámetros.

Proba os nosos núcleos en GitHub e le o artigo de investigación completo en arXiv.

Introdución

A arquitectura de expertos mixtos (MoE, Mixture-of-Experts) consolidouse como o estándar para escalar a centos de miles de millóns ou billóns de parámetros, á vez que ofrece unha latencia razoable no momento da inferencia. Os modelos MoE substituíron a capa densa dun transformador por un conxunto de expertos e unha pequena capa de encamiñamento lineal que decide con que expertos se debe multiplicar un token. Ao agrupar os tokens asociados ao mesmo experto, o problema vólvese disperso e paralelizable de forma trivial, xa que os diferentes expertos poden agora aloxarse en diferentes GPUs repartidas en varios nós.

A diferenza doutras formas de paralelismo, como o paralelismo de tensores (TP) e o paralelismo de datos (DP), que se poden implementar facilmente mediante bibliotecas de comunicación colectiva como torch.distributed e NCCL, o encamiñamento MoE implica comunicación punto a punto dispersa. Un núcleo de despacho debe dividir o conxunto de tokens entrantes e envialos aos rangos que aloxan os expertos, mentres que o seu problema dual, o núcleo de combinación, debe recibir as súas contrapartidas procesadas e calcular a súa media ponderada. estas transferencias non se mapean en colectivos e requirir núcleos personalizados para acadar unha baixa latencia.

En Perplexity, estivemos a desenvolver tanto núcleos portátiles entre nós como núcleos especializados intra-nó, presentando o noso traballo con maior profundidade nun artigo anterior sobre a comunicación MoE. Tamén mostramos noutra publicación do blog que a implementación en varios nós pode ofrecer unha menor latencia e un maior rendemento para grandes modelos MoE.

Neste artigo, profundamos na implementación dos nosos novos núcleos entre nós, que melloran notablemente o rendemento do noso traballo anterior baseado en NVSHMEM. Ademais de mostrar resultados de vangarda en ConnectX-7, tamén demostramos que, dentro do noso tempo de execución de inferencia, os núcleos se poden utilizar para servir eficazmente aos modelos de código aberto máis grandes.

Inferencia sobre InfiniBand

O paralelismo de tensores e o paralelismo de expertos poden escalar de forma efectiva dentro dun único nó, xa que NVLink ofrece un alto rendemento (900 GB/s) cunha latencia de microsegundos. Aínda que os sistemas NVIDIA GB200 e GB300 están equipados con suficiente memoria para conter os pesos e as caché KV de modelos de billóns de parámetros, as implementacións na instancia p5en de AWS con ata 8 GPUs H200 están limitadas a 1120 GB de HBM, que se deben dividir entre os pesos do modelo e as caché KV. En consecuencia, certos modelos requiren implementacións entre nós a través de InfiniBand.

Aínda que InfiniBand segue a ser extremadamente rápido, só ofrece un rendemento de 400 Gbps (50 GB/s), o que engade de decenas a centos de microsegundos de latencia para o encamiñamento MoE. Afortunadamente, as latencias engadidas pódense ocultar mediante o uso de expertos compartidos, microprazos e a superposición de cómputo e comunicación para ofrecer un rendemento viable a un custo competitivo. DeepSeek foi pioneira nestas técnicas a través de DeepSeek-R1, implementado en clústeres de GPUs H800 e adaptadores NVIDIA ConnectX-7 InfiniBand.

Nos clústeres de AWS, a situación é diferente debido ao uso dos seus adaptadores personalizados Elastic Fabric Adapters (EFA). Aínda que estes ofrecen un rendemento máximo de 400 Gbps nos colectivos que se adoitan empregar nas cargas de traballo de adestramento, quedan lixeiramente por debaixo de ConnectX-7 nos tamaños de mensaxe intercambiados durante o despacho e a combinación de MoE. Ademais, EFA non é compatible con GPUDirect Async, o que require un fío de proxy da CPU que faga de ponte entre a GPU e a tarxeta de rede (NIC) para iniciar as transferencias. Ademais da complexidade engadida, as transaccións PCIe adicionais engaden unha sobrecarga equivalente a microsegundos a todas as transferencias. Os marcos de traballo como NVSHMEM permiten aos creadores de núcleos ofrecer implementacións independentes do dispositivo; con todo, as API xenéricas xeran sobrecargas que fan que as implementacións non sexan eficaces en termos de custo e latencia.

pplx-kernels baseados en NVSHMEM

Os nosos núcleos anteriores abstraían a implementación de transporte subxacente a través de NVSHMEM. Confiaban en nvshmemx_putmem_signal_nbi_warp para transferir tokens un a un, utilizando contadores atómicos para sincronizar os pares. Aínda que isto ofrecía un rendemento razoable nas NICs ConnectX-7 con IBGDA, a implementación IBRC de ConnectX-7 baseada en proxy era significativamente máis lenta, mentres que a implementación de EFA non lograva ofrecer latencias inferiores a un milisegundo.

Utilizamos os núcleos como banco de probas para o hardware subxacente, á vez que determinamos a interface máis axeitada para expoñer aos nosos modelos. Formulamos a hipótese de que as frecuentes transferencias e interaccións entre os núcleos e as NICs ou os fíos de proxy afectaban negativamente ao rendemento. Ademais, a nosa implementación non dependía de ningunha garantía por parte do hardware agás a fiabilidade, o que suxería que podíamos acelerar as cousas reducindo a sobrecarga das primitivas de sincronización.

Tras estas observacións, deseñamos os nosos núcleos baseados en EFA expoñendo a mesma interface, optimizando as interaccións entre o núcleo do dispositivo e o proxy do host baixo o capó para acadar o máximo rendemento.

MoE all-to-all sobre EFA

Ampliar

Os novos núcleos seguen unha arquitectura híbrida de CPU e GPU, onde os núcleos de GPU interactúan co modelo no dispositivo e un fío de proxy que se executa na CPU do host xestiona as interaccións coa NIC. Os novos núcleos reutilizan o mesmo TransferEngine que construímos para as transferencias de caché KV.

O encamiñamento MoE realízase mediante un núcleo de despacho, que recibe do modelo os tokens e os expertos por token para encamiñalos, envíalos a través da rede aos rangos que aloxan os expertos. Despois de que os núcleos GEMM agrupados multipliquen os tokens cos pesos dos expertos, o núcleo de combinación transfire os tokens aos seus rangos de orixe e calcula a súa media ponderada antes de devolvelos ao modelo.

Tanto o núcleo de despacho como o de combinación divídense nunha metade emisora e outra receptora para facilitar a superposición e o microprazamento. O emisor reorganiza os tokens en búferes de envío para que se poidan escribir nos pares mediante unha única operación de escritura RDMA. Despois de executar o emisor, a GPU pode realizar outras tarefas (como a multiplicación de expertos compartidos) mentres as transferencias a través da rede están pendentes en segundo plano. Os núcleos receptores bloquéanse ata que se reciben todas as transferencias. En segundo plano, un fío de traballo do lado do host xestiona as transferencias: a GPU avísao cando se encheren os búferes de envío, momento no cal despacha as escrituras RDMA. Unha vez que a NIC sinala a finalización, desboca o receptor, o cal devolve os tokens ao modelo.

Os núcleos están estreitamente acoplados co host, comunicándose entre si mediante memoria unificada e GDRCopy. A memoria unificada de CUDA permite que a memoria do dispositivo ou a memoria do host sexan mapeadas nos espazos de enderezos de ambos, xestionando de forma transparente a comunicación a través de PCIe.

Aínda que a memoria unificada é axeitada para transferencias masivas, a sondaxe de baixa latencia realízase a través de GDRCopy. GDRCopy utiliza GPUDirect RDMA para ler e escribir na memoria da GPU cunha latencia de microsegundos.

Ademais de RDMA, dentro dun nó os núcleos tamén aproveitan NVLink para intercambiar tokens. Isto pode descargar unha fracción significativa do tráfico, ata un cuarto para EP32 e ata un oitavo para EP64, acadando en xeral latencias máis baixas.

Búferes de envío e recepción

O deseño dos búferes de envío para o núcleo de despacho é fundamental para acadar latencias viables que impliquen o proxy do host. Se hai N rangos e cada un despacha T tokens, cada un encamiñado a R expertos dun total de E, un rango pode recibir ata N * T * max(R, E / N) tokens no peor dos casos posibles, se todos son encamiñados cara a el. Ademais, este límite mínimo só se pode acadar cando os emisores dispoñen as súas escrituras de forma contiguas, xa que os búferes de recepción privados consumirían demasiada memoria. Para que os emisores poidan dispoñer as súas escrituras de maneira consecutiva sen conflitos, débese distribuír un coñecemento completo da información de encamiñamento a cada rango. En consecuencia, os nosos núcleos de despacho intercambian primeiro os recuentos de tokens por experto, tras o cal cada un pode determinar a onde van as súas escrituras e de que rango reciben os tokens e con que desprazamento.

O simple intercambio de información de encamiñamento daría lugar a latencias innecesariamente altas, xa que cada rango tería que agardar a que chegasen as escrituras antes de despachar calquera escritura propia. No seu lugar, reservamos unha pequena cantidade de espazo privado por emisor en cada rango receptor e despachamos algúns dos tokens xunto coas escrituras da información de encamiñamento. Isto garante a utilización plena do ancho de banda durante a maior parte do proceso de despacho. Unha vez intercambiadas as rutas, os tokens restantes transfírense mediante unha única escritura por cada par remoto. O intercambio de información de encamiñamento e os búferes de tamaño fixo, xunto coas transferencias que os encheren, ilústranse na seguinte figura.

Ampliar

Os núcleos de combinación son substancialmente máis sinxelos, xa que a información de encamiñamento xa está dispoñible para eles desde a fase de despacho. Poden copiar os tokens nunha única escritura por par, tras o cal len desde o búfer de recepción para calcular a media ponderada dos tokens nos tensores de saída.

Para NVLink, os búferes de envío e recepción dentro dun nó están mapeados nos pares para que estes poidan enviar ou extraer tokens directamente dos rangos veciños.

Despacho e combinación

As seguintes figuras detallan a temporización e as interaccións entre a GPU, a CPU e a NIC, mostrando unha iteración de despacho e combinación intercalada cos núcleos do modelo.

O traballo comeza no envío de despacho, co núcleo agregando a información de encamiñamento e calculando os recuentos de tokens por experto mediante acumulación atómica en memoria compartida dentro dun único bloque. Unha vez que os recuentos de tokens están dispoñibles, cópianse na memoria do host e notifícase inmediatamente ao proxy para que comece a despachar a información de encamiñamento a todos os demais rangos pares. En paralelo, o despacho continúa calculando a suma acumulada dos recuentos de tokens, identificando os desprazamentos necesarios para empaquetar os tokens en búferes de envío contiguos por par. Unha vez dispoñibles os desprazamentos, todos os tokens que se van enviar a través de EFA empaquetáronse nos búferes de envío. Notifícase ao proxy para que publique as escrituras nos búferes de envío dos pares entre nós. A nivel intra-nó, os núcleos tamén envían inmediatamente tokens a búferes privados dentro do mesmo nó a través de NVLink.

Ampliar

Ampliar

Unha vez que un rango recibe a información de encamiñamento de todos os pares, centralizámaa e calcula varios desprazamentos necesarios para identificar os tokens para o despacho e a combinación. Este traballo realízase na CPU mentres as transferencias xa están pendentes, polo que non é crítico para a latencia. Posteriormente, unha vez coñecidos todos os desprazamentos, publícanse as escrituras para o resto dos tokens, que superan o tamaño do búfer privado, e a fase de envío do despacho complétase.

No extremo receptor, o núcleo extrae primeiro a través de NVLink os tokens do mesmo nó que superan a capacidade dos búferes privados. A continuación, bloquéase ata que o proxy do host sinala a finalización de todas as transferencias. Unha vez recibio o sinal, os tokens transfírense desde os búferes de recepción ao tensor de saída pasado ao grupo GEMM. Dado que os búferes se reutilizan entre o despacho e a combinación, tanto as barreiras de NVLink como de RDMA garanten que os búferes de recepción se leu antes de permitir que o seguinte núcleo os sobrescriba.

Os núcleos de combinación seguen un patrón máis ou menos similar: reorganizan os tokens, envían un sinal ao proxy e agardan á finalización. Con todo, os metadatos para publicar as súas escrituras están dispoñibles de inmediato e todos os desprazamentos remotos son coñecidos; polo tanto, todos os tokens transferidos intra-nó envíanse a través de NVLink aos receptores e os tokens entre nós transfírense mediante unha única escritura. Antes de devolver os resultados, o núcleo acumula os tokens dos búferes de recepción, escribindo-os nos tensores de saída.

TransferEngine para o encamiñamento MoE

Para prototipar a nosa implementación, inicialmente reutilizamos o Transfer Engine construído para a pre-carga e a descodificación desaggregadas co fin de emitir e rastrexar escrituras RDMA. Con todo, a diferenza da descodificación desagregada, onde cada nó publica un gran número de escrituras no mesmo par, o despacho MoE publica pequenas transferencias a todos os pares nunha única operación. O rendemento predeterminado non era óptimo e a nosa implementación estaba atrasada con respecto a DeepEP en ConnectX-7 nuns 20 us. En consecuencia, especializamos o TransferEngine para admitir dúas operacións, dispersión (scatter) e barreira (barrier), co fin de optimizar as interaccións cun gran número de pares.

Tanto a operación de dispersión como a de barreira expoñen grupos de pares que se poden pre-rexistrar para pre-encher as estruturas de datos utilizadas pola implementación de transporte subxacente (plantillas de solicitude de traballo ou WR templating). A dispersión copia diferentes anacos do mesmo búfer de orixe en diferentes desprazamentos en todos os membros dun grupo de pares. A barreira sinala a todos os pares mediante un sinal inmediato. Encolar e xestionar estas operacións de forma masiva aforra valiosos microsegundos no camiño crítico. En EFA, onde dúas tarxetas de rede de 200 Gbps se agregan para proporcionar un ancho de banda de 400 Gbps, fragmentamos os pares do grupo entre as dúas tarxetas de rede en lugar de fragmentar os bytes de cada transferencia.

Aínda que inicialmente desenvolvemos o TransferEngine exclusivamente en EFA utilizando libfabric, os núcleos MoE funcionan en ConnectX-7 sen cambiar nin unha soa liña de código despois de engadir soporte para ConnectX-7 no TransferEngine utilizando libibverbs. En comparación con EFA baseado en Scalable Reliable Datagram (SRD), ConnectX-7 require traballo adicional para a configuración da conexión e a xestión de pares. Utilizamos un par de colas de datagramas non fiables (UD) para iniciar a Conexión Fiable (RC). Creamos dous pares de colas para cada par remoto, un para SEND/RECV bidireccional e outro para WRITE_IMM unidireccional, porque tanto RECV como WRITE_IMM consumen a operación RECV en orde, mentres que ambas as operacións requiren tamaños de búfer diferentes. Ademais, relaxamos a ordenación de transaccións PCIe entre ConnectX-7 e a memoria da GPU. De maneira similar a EFA, implementamos a optimización de plantillas de solicitudes de traballo. Ademais, ao enviar varias mensaxes pequenas ao mesmo par, utilizamos o encadeamento de solicitudes de traballo (WR chaining) para mellorar a taxa de mensaxes.

Dado que o traballo do lado da GPU é idéntico en ambos os adaptadores, levamos a implementación de transporte subxacente aos seus límites para superar o rendemento de DeepEP e resaltar que a brecha entre o encamiñamento MoE en EFA e en ConnectX-7 está determinada polas capacidades do hardware subxacente.

Avaliación

Latencia de despacho e combinación

Avaliamos as latencias de despacho e combinación en sistemas H200 equipados con adaptadores ConnectX-7 e EFA na configuración estándar de DeepSeek-V3 e Kimi-K2 cunha dimensión oculta de 7168, despacho fp8 escalado por bloques e combinación bf16. Consideramos un tamaño de lote de 128 para o despacho e 4096 para a combinación. En ConnectX-7, comparamos con DeepEP e cos nosos núcleos pplx anteriores baseados en NVSHMEM, utilizando tanto GPUDirect Async (IBGDA) como un proxy de host (IBRC). Acadamos latencias combinadas de 459 us, 582 us e 692 us en EP16, EP32 e EP64, respectivamente. Superamos a outros núcleos baseados en EFA que dependen de proxys xenéricos, como UCCL-EP, que acadan 519 us, 966 us e 1159 us.

Ampliar

Na descodificación, a pesar de que os nosos núcleos son uns poucos microsegundos máis lentos que DeepEP, acadamos en xeral latencias máis baixas xa que os nosos núcleos de combinación son substancialmente máis rápidos. Aínda que utilizamos un proxy de host, debido ás nosas transferencias masivas superamos o rendemento de todas as implementacións sen proxy baseadas en IBGDA. Tamén mostramos que a implementación do proxy importa, xa que o noso proxy altamente especializado é unha orde de magnitude máis rápido que a implementación xenérica de NVSHMEM. A pesar de que EFA é dúas veces máis lento que ConnectX-7 cando transfire paquetes de 256 KB, que están preto do tamaño típico de paquete intercambiado polo encamiñamento MoE, a brecha entre os dous adaptadores é globalmente menos substancial.

Ampliar

Aínda que os nosos núcleos optimizados para a descodificación poden escalar aos tamaños de lote de pre-carga sen ningún axuste adicional, o rendemento está atrasado con respecto a DeepEP. Isto débese principalmente a que acumulamos tokens só no lado receptor da combinación, mentres que DeepEP acumula no emisor para reducir a cantidade de datos transferidos. Con todo, esta optimización pode ter implicacións na precisión. Ademais, dado que os núcleos transfieren datos nunha única escritura, utilizamos máis memoria para os búferes de envío e recepción.

Rendemento de Kimi-K2 / DeepSeek-V3

Para demostrar a viabilidade dos nosos núcleos, realizamos unha avaliación comparativa do rendemento de extremo a extremo do noso motor de inferencia propio que serve a DeepSeek-V3 (671B de parámetros) e Kimi-K2 (1T de parámetros) en múltiples instancias p5en equipadas con Elastic Fabric Adapters (EFA) e GPUs H200. Executamos o modelo en modo de verificación de borrador e informamos da latencia asumiindo unha taxa de aceptación do 80% utilizando unha única capa MTP.

En DeepSeek-V3, comparamos cunha liña de base intra-nó con EP=8 que executa núcleos de despacho e combinación eficientes baseados exclusivamente en NVLink. As implementacións entre nós igualan ou superan en xeral o rendemento da implementación nun único nó, particularmente en lotes de tamaño mediano, o que demostra que escalar a través de nós achega melloras ao rendemento global.

En Kimi-K2, os núcleos fan posible a serviz. Debido ao seu tamaño, o modelo non cabe dentro do mesmo nó utilizando o paralelismo de datos (DP) nun nó H200, polo que se require unha implementación entre nós para servir o modelo. Como o modelo ten menos cabezas de atención, acadamos mellores latencias que con DeepSeek-V3 e podemos servir o modelo en EFA cunha latencia viable.

Ampliar

Ampliar

Tamén consideramos o microprazamento, aínda que só mellorou as latencias en tamaños de lote moi elevados ou cando se utilizan máis dun token de borrador. Para beneficiarse da división dos problemas GEMM agrupados en dous, os tamaños de lote de tokens deben ser o suficientemente densos. A maior parte do tempo gástase aínda no despacho e a combinación MoE, o que indica que calquera mellora adicional nos núcleos mellorará consecuentemente as latencias e facilitará o microprazamento.

Traballo futuro

Estamos a traballar en estreita colaboración cos enxeñeiros de AWS para mellorar o rendemento en EFA. Pola nosa banda, seguiremos as actualizacións de libfabric, o que reducirá a sobrecarga do plano de datos. Ademais, experimentaremos con efa-direct para reducir aínda máis as sobrecargas no espazo de usuario e interactuar directamente co hardware.

Referencias

Informe técnico de DeepSeek R1

Informe técnico de Kimi K2

DeepEP

UCCL-EP

Mellora do rendemento de rede dos sistemas HPC utilizando NVIDIA Magnum IO NVSHMEM e GPUDirect Async

Cita

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