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

En Perplexity utilizamos los mejores modelos para nuestro producto, nuestras API y nuestros equipos de investigación. Los modelos grandes de código abierto de tipo mezcla de expertos (MoE), como Kimi-K2, plantean retos particulares, ya que los nodos de inferencia más grandes con 8 GPU NVIDIA H200 no pueden alojarlos de manera eficiente, lo que requiere implementaciones multinodo. Presentamos un conjunto de kernels para paralelismo de expertos que logran latencias de vanguardia en ConnectX-7, superando el rendimiento de DeepEP. Estos mismos kernels son también los primeros en alcanzar latencias viables en AWS Elastic Fabric Adapter (EFA), lo que permite implementaciones de modelos de un billón de parámetros.

Pruebe nuestros kernels en GitHub y lea el artículo de investigación completo en arXiv.

Introducción

La mezcla de expertos (MoE) ha surgido como la arquitectura estándar para escalar a cientos de miles de millones o billones de parámetros, ofreciendo además una latencia razonable en el momento de la inferencia. Los modelos MoE reemplazan la capa densa de un transformer por un conjunto de expertos y una pequeña capa de enrutamiento lineal que decide con qué expertos debe multiplicar un token. Al agrupar los tokens asociados con el mismo experto, el problema se vuelve disperso y trivialmente paralelizable, ya que diferentes expertos pueden ahora alojarse en diferentes GPU repartidas en múltiples nodos.

A diferencia de otras formas de paralelismo, como el paralelismo de tensores (TP) y el paralelismo de datos (DP), que pueden implementarse fácilmente mediante bibliotecas de comunicación colectiva como torch.distributed y NCCL, el enrutamiento MoE implica una comunicación dispersa entre pares. Un kernel de despacho debe dividir el conjunto de tokens entrantes y enviarlos a los rangos que alojan a los expertos, mientras que su problema dual, el kernel de combinación, debe recibir sus contrapartes procesadas y calcular su media ponderada. Estas transferencias no se asignan a colectivos y requieren kernels personalizados para lograr una baja latencia.

En Perplexity hemos estado desarrollando tanto kernels portátiles entre nodos como kernels especializados dentro de un nodo, presentando nuestro trabajo con mayor profundidad en un artículo anterior sobre comunicación MoE. También mostramos en otra entrada de blog que la implementación multinodo puede aportar una menor latencia y un mayor rendimiento para modelos MoE grandes.

En este artículo, profundizamos en la implementación de nuestros nuevos kernels entre nodos que mejora enormemente el rendimiento de nuestro trabajo anterior basado en NVSHMEM. Además de mostrar resultados de vanguardia en ConnectX-7, demostramos también que dentro de nuestro entorno de ejecución de inferencia, los kernels pueden utilizarse para servir eficazmente a los modelos de código abierto más grandes.

Inferencia sobre InfiniBand

El paralelismo de tensores y el paralelismo de expertos pueden escalar eficazmente dentro de un solo nodo, ya que NVLink ofrece un alto rendimiento (900 GB/s) con una latencia de microsegundos. Aunque los sistemas NVIDIA GB200 y GB300 están equipados con suficiente memoria para contener los pesos y las cachés KV de modelos de un billón de parámetros, las implementaciones en instancias p5en de AWS con hasta 8 GPU H200 están limitadas a 1120 GB de HBM, que deben repartirse entre los pesos del modelo y las cachés KV. En consecuencia, ciertos modelos requieren implementaciones entre nodos sobre InfiniBand.

Aunque InfiniBand sigue siendo increíblemente rápido, solo ofrece 400 Gbps (50 GB/s) de rendimiento, lo que añade decenas a cientos de microsegundos de latencia para el enrutamiento MoE. Afortunadamente, las latencias añadidas pueden ocultarse mediante el uso de expertos compartidos, micro-batching y la superposición de cálculo y comunicación para ofrecer un rendimiento viable a un coste competitivo. DeepSeek fue pionero en estas técnicas a través de DeepSeek-R1, desplegado en clústeres de GPU H800 y adaptadores InfiniBand NVIDIA ConnectX-7.

En los clústeres de AWS, la situación es diferente debido al uso de sus adaptadores personalizados Elastic Fabric Adapters (EFA). Aunque estos ofrecen un rendimiento máximo de 400 Gbps en los colectivos utilizados habitualmente en las cargas de trabajo de entrenamiento, se quedan ligeramente cortos respecto a ConnectX-7 en los tamaños de mensaje intercambiados durante el despacho y la combinación de MoE. Además, EFA no admite GPUDirect Async, lo que requiere un hilo proxy de CPU que sirva de puente entre la GPU y la NIC para iniciar las transferencias. Además de la complejidad añadida, las transacciones PCIe adicionales añaden microsegundos de sobrecarga a todas las transferencias. Marcos de trabajo como NVSHMEM permiten a los escritores de kernels proporcionar implementaciones independientes del dispositivo; sin embargo, las API genéricas incurren en gastos generales que hacen que las implementaciones no sean eficaces en términos de coste y latencia.

pplx-kernels basados en NVSHMEM

Nuestros kernels anteriores abstraían la implementación de transporte subyacente a través de NVSHMEM. Se basaban en nvshmemx_putmem_signal_nbi_warp para transferir tokens uno a uno, utilizando contadores atómicos para sincronizar a los pares. Aunque esto ofrecía un rendimiento razonable en las NIC ConnectX-7 con IBGDA, la implementación IBRC de ConnectX-7 basada en proxy era significativamente más lenta, mientras que la implementación EFA no lograba ofrecer latencias inferiores a un milisegundo.

Utilizamos los kernels como banco de pruebas para el hardware subyacente, al tiempo que determinamos la interfaz más adecuada para exponer a nuestros modelos. Planteamos la hipótesis de que las frecuentes transferencias e interacciones entre los kernels y las NIC o los hilos proxy dificultaban el rendimiento. Además, nuestra implementación no se basaba en ninguna garantía del hardware más allá de la fiabilidad, lo que sugería que podíamos acelerar las cosas reduciendo la sobrecarga de las primitivas de sincronización.

Tras estas observaciones, ideamos nuestros kernels basados en EFA exponiendo la misma interfaz, optimizando las interacciones entre el kernel del dispositivo y el proxy del host bajo el capó para obtener el máximo rendimiento.

MoE all-to-all sobre EFA

Ampliar

Los nuevos kernels siguen una arquitectura híbrida CPU-GPU, en la que los kernels de GPU interactúan con el modelo en el dispositivo y un hilo proxy que se ejecuta en la CPU del host gestiona las interacciones con la NIC. Los nuevos kernels reutilizan el mismo TransferEngine que creamos para transferencias de caché KV.

El enrutamiento MoE es realizado por un kernel de despacho, que recibe los tokens y los expertos por token a los que enrutarlos desde el modelo, enviándolos a través de la red a los rangos que alojan a los expertos. Después de que los kernels GEMM agrupados multipliquen los tokens con los pesos de los expertos, el kernel de combinación transfiere los tokens a sus rangos de origen y calcula su media ponderada antes de devolverlos al modelo.

Tanto el kernel de despacho como el de combinación se dividen en una mitad emisora y otra receptora para facilitar la superposición y el micro-batching. El emisor mezcla los tokens en los búferes de envío para que puedan ser escritos a los pares a través de una única operación de escritura RDMA. Después de ejecutar el emisor, la GPU puede realizar otras tareas (como la multiplicación de expertos compartidos) mientras las transferencias a través de la red están pendientes en segundo plano. Los kernels receptores se bloquean hasta que se hayan recibido todas las transferencias. En segundo plano, un hilo de trabajador en el lado del host gestiona las transferencias: la GPU le avisa cuando se han rellenado los búferes de envío, tras lo cual despacha las escrituras RDMA. Una vez que la NIC señala la finalización, desbloquea al receptor, que devuelve los tokens al modelo.

Los kernels están estrechamente acoplados con el host, comunicándose entre sí a través de memoria unificada y GDRCopy. La memoria unificada de CUDA permite asignar la memoria del dispositivo o la memoria del host a los espacios de direcciones de ambos, gestionando de forma transparente la comunicación a través de PCIe.

Aunque la memoria unificada es adecuada para transferencias masivas, el sondeo de baja latencia se realiza a través de GDRCopy. GDRCopy utiliza GPUDirect RDMA para leer y escribir en la memoria de la GPU con una latencia de microsegundos.

Además del RDMA, dentro de un nodo los kernels también aprovechan NVLink para intercambiar tokens. Esto puede descargar una fracción significativa del tráfico, hasta una cuarta parte para EP32 y hasta una octava para EP64, logrando latencias generales más bajas.

Búferes de envío y recepción

El diseño de los búferes de envío para el kernel de despacho es crucial para lograr latencias viables que impliquen al proxy del host. Si hay N rangos, cada uno despachando T tokens, cada uno enrutado a R expertos de E, un rango puede recibir hasta N * T * max(R, E / N) tokens en el peor de los casos, si todos se enrutan a él. Además, este límite mínimo solo puede alcanzarse cuando los emisores disponen sus escrituras de forma contigua, ya que los búferes de recepción privados consumirían demasiada memoria. Para que los emisores puedan disponer sus escrituras consecutivamente sin conflictos, debe distribuirse a cada rango un conocimiento completo de la información de enrutamiento. En consecuencia, nuestros kernels de despacho intercambian primero los recuentos de tokens por experto, tras lo cual cada uno puede determinar adónde van sus escrituras y desde qué desplazamiento reciben tokens de cada rango.

El simple intercambio de información de enrutamiento daría lugar a latencias innecesariamente altas, ya que cada rango tendría que esperar a que llegaran las escrituras antes de despachar ninguna escritura propia. En su lugar, reservamos una pequeña cantidad de espacio privado por emisor en cada rango receptor y despachamos algunos de los tokens junto con las escrituras de la información de enrutamiento. Esto garantiza la utilización total del ancho de banda durante la mayor parte del proceso de despacho. Una vez intercambiadas las rutas, los tokens restantes se transfieren mediante una única escritura por par remoto. El intercambio de información de enrutamiento y los búferes de tamaño fijo, junto con las transferencias que los rellenan, se ilustran en la figura siguiente.

Ampliar

Los kernels de combinación son sustancialmente más sencillos, ya que la información de enrutamiento ya está disponible para ellos desde la etapa de despacho. Pueden copiar los tokens en una sola escritura por par, después de lo cual leen desde el búfer de recepción para calcular la media ponderada de los tokens en los tensores de salida.

Para NVLink, los búferes de envío y recepción dentro de un nodo se asignan a los pares para que puedan empujar o extraer tokens directamente de los rangos vecinos.

Despacho y combinación

Las figuras siguientes detallan los tiempos y las interacciones entre la GPU, la CPU y la NIC, mostrando una iteración de despacho y combinación intercalada con los kernels del modelo.

El trabajo comienza en el envío de despacho, con el kernel agregando la información de enrutamiento y calculando los recuentos de tokens por experto mediante acumulación atómica en memoria compartida dentro de un único bloque. Una vez que los recuentos de tokens están disponibles, se copian a la memoria del host y se notifica inmediatamente al proxy para que comience a despachar la información de enrutamiento a todos los demás rangos pares. En paralelo, el despacho continúa calculando la suma acumulada de los recuentos de tokens, identificando los desplazamientos necesarios para empaquetar los tokens en búferes de envío contiguos por par. Una vez disponibles los desplazamientos, todos los tokens que se enviarán a través de EFA se empaquetan en los búferes de envío. Se notifica al proxy para que publique las escrituras en los búferes de envío de los pares entre nodos. Dentro del nodo, los kernels también empujan inmediatamente los tokens a búferes privados dentro del mismo nodo a través de NVLink.

Ampliar

Ampliar

Una vez que un rango recibe información de enrutamiento de todos los pares, la centraliza y calcula varios desplazamientos necesarios para identificar los tokens para el despacho y la combinación. Este trabajo se realiza en la CPU mientras las transferencias están ya pendientes, por lo que no es crítico para la latencia. Posteriormente, una vez conocidos todos los desplazamientos, se publican las escrituras del resto de los tokens que superan el tamaño del búfer privado y la fase de envío del despacho se completa.

En el lado del receptor, el kernel primero extrae tokens del mismo nodo que exceden la capacidad de los búferes privados a través de NVLink. A continuación, se bloquea hasta que el proxy del host señale la finalización de todas las transferencias. Una vez recibida la señal, los tokens se mezclan desde los búferes de recepción hacia el tensor de salida pasado al GEMM de grupo. Dado que los búferes se reutilizan entre el despacho y la combinación, tanto las barreras NVLink como RDMA garantizan que los búferes del receptor se hayan leído antes de permitir que el siguiente kernel los sobrescriba.

Los kernels de combinación siguen un patrón similar, mezclando tokens, señalizando al proxy y esperando a que se completen. Sin embargo, los metadatos para publicar sus escrituras están fácilmente disponibles y se conocen todos los desplazamientos remotos, por lo que todos los tokens transferidos dentro del nodo se envían a través de NVLink a los receptores y los tokens entre nodos se transfieren mediante una única escritura. Antes de devolver los resultados, el kernel acumula los tokens de los búferes del receptor, escribiéndolos en los tensores de salida.

TransferEngine para enrutamiento MoE

Para prototipar nuestra implementación, inicialmente reutilizamos el Transfer Engine creado para el prellenado y la decodificación desagregados con el fin de emitir y rastrear escrituras RDMA. Sin embargo, a diferencia de la decodificación desagregada, donde cada nodo publica un gran número de escrituras en el mismo par, el despacho de MoE publica pequeñas transferencias a todos los pares en una sola operación. De entrada, el rendimiento no era óptimo y nuestra implementación iba por detrás de DeepEP en ConnectX-7 en unos 20us. En consecuencia, especializamos el TransferEngine para soportar dos operaciones, dispersión (scatter) y barrera (barrier), para optimizar las interacciones con un gran número de pares.

Tanto las operaciones de dispersión como las de barrera exponen grupos de pares que pueden ser pre-registrados para pre-poblar las estructuras de datos utilizadas por la implementación de transporte subyacente (plantillado de solicitud de trabajo (WR)). La dispersión copia diferentes fragmentos del mismo búfer de origen en diferentes desplazamientos en todos los miembros de un grupo de pares. La barrera señala a todos los pares a través de un inmediato. Poner en cola y manejar estas operaciones a granel ahorra microsegundos valiosos en la ruta crítica. En EFA, donde dos NIC de 200 Gbps se agregan para proporcionar un ancho de banda de 400 Gbps, fragmentamos los pares del grupo a través de las dos NIC en lugar de fragmentar los bytes de cada transferencia.

Aunque inicialmente desarrollamos el TransferEngine exclusivamente en EFA utilizando libfabric, los kernels de MoE funcionan en ConnectX-7 sin cambiar una sola línea de código tras añadir soporte para ConnectX-7 en el TransferEngine utilizando libibverbs. En comparación con EFA basado en Datagrama Fiable Escalable (SRD), ConnectX-7 requiere trabajo adicional para la configuración de la conexión y la gestión de pares. Utilizamos un par de colas de Datagrama No Fiable (UD) para arrancar la Conexión Fiable (RC). Creamos dos pares de colas para cada par remoto, uno para ENVÍO/RECEPCIÓN de dos lados y otro para ESCRITURA_INMEDIATA de un lado, porque tanto la RECEPCIÓN como la ESCRITURA_INMEDIATA consumen la operación de RECEPCIÓN en orden, mientras que las dos operaciones requieren diferentes tamaños de búfer. Además, relajamos el orden de las transacciones PCIe entre ConnectX-7 y la memoria de la GPU. Similar a EFA, implementamos la optimización de plantillado WR. Además, cuando enviamos múltiples mensajes pequeños al mismo par, utilizamos el encadenamiento WR para mejorar la velocidad de los mensajes.

Dado que el trabajo en el lado de la GPU es idéntico en ambos adaptadores, llevamos la implementación de transporte subyacente a sus límites para superar el rendimiento de DeepEP y resaltar que la brecha entre el enrutamiento MoE de EFA y ConnectX-7 está determinada por las capacidades del hardware subyacente.

Evaluación

Latencia de despacho/combinación

Evaluamos las latencias de despacho y combinación en sistemas H200 equipados con adaptadores ConnectX-7 y EFA en la configuración estándar DeepSeek-V3 y Kimi-K2 con una dimensión oculta de 7168, despacho fp8 escalado por bloques y combinación bf16. Consideramos un tamaño de lote de 128 para el despacho y 4096 para la combinación. En ConnectX-7, comparamos con DeepEP y nuestros anteriores pplx-kernels basados en NVSHMEM, utilizando tanto GPUDirect Async (IBGDA) como un proxy de host (IBRC). Logramos latencias combinadas de 459us, 582us y 692us en EP16, EP32 y EP64, respectivamente. Superamos a otros kernels basados en EFA que dependen de proxies genéricos, como UCCL-EP, que logran 519us, 966us y 1159us.

Ampliar

En la decodificación, a pesar de que nuestros kernels son unos microsegundos más lentos que DeepEP, logramos latencias generales más bajas ya que nuestros kernels de combinación son sustancialmente más rápidos. Aunque utilizamos un proxy de host, debido a nuestras transferencias masivas superamos el rendimiento de todas las implementaciones sin proxy basadas en IBGDA. También mostramos que la implementación del proxy importa, ya que nuestro proxy altamente especializado es un orden de magnitud más rápido que la implementación genérica de NVSHMEM. A pesar de que EFA es dos veces más lento que ConnectX-7 al transferir paquetes de 256 KB, que están cerca del tamaño típico de paquete intercambiado por el enrutamiento MoE, la brecha entre los dos adaptadores es en general menos sustancial.

Ampliar

Aunque nuestros kernels optimizados para la decodificación pueden escalar a tamaños de lote de prellenado sin ajustes adicionales, el rendimiento va por detrás de DeepEP. Esto se debe principalmente al hecho de que acumulamos tokens solo en el lado del receptor de la combinación, mientras que DeepEP acumula en el emisor para reducir la cantidad de datos transferidos. Sin embargo, esta optimización puede tener implicaciones en la precisión. Además, dado que los kernels transfieren datos en una sola escritura, utilizamos más memoria para los búferes de envío y recepción.

Rendimiento de Kimi-K2 / DeepSeek-V3

Para probar la viabilidad de nuestros kernels, realizamos pruebas comparativas del rendimiento de extremo a extremo de nuestro motor de inferencia interno que sirve a DeepSeek-V3 (671B parámetros) y Kimi-K2 (1T parámetros) en múltiples instancias p5en equipadas con Elastic Fabric Adapters (EFA) y GPU H200. Ejecutamos el modelo en modo de verificación de borrador e informamos la latencia asumiendo una tasa de aceptación del 80% utilizando una única capa MTP.

En DeepSeek-V3, comparamos con una línea base intranodo EP=8 que ejecuta kernels de despacho y combinación eficientes solo con NVLink. Las implementaciones entre nodos en general igualan o superan el rendimiento de la implementación de un solo nodo, particularmente en lotes de tamaño mediano, lo que demuestra que escalar entre nodos sí aporta mejoras al rendimiento general.

En Kimi-K2, los kernels permiten servir. Debido a su tamaño, el modelo no puede encajar dentro del mismo nodo utilizando paralelismo de datos (DP) en un nodo H200, por lo que se requiere una implementación entre nodos para servir el modelo. Como el modelo tiene menos cabezales de atención, logramos mejores latencias que DeepSeek-V3 y podemos servir el modelo en EFA a una latencia viable.

Ampliar

Ampliar

También hemos considerado el micro-batching, aunque solo mejoró las latencias en tamaños de lote muy altos o cuando se usan más de un token de borrador. Para beneficiarse de dividir los problemas de GEMM agrupados en dos, los tamaños de lote de tokens deben ser suficientemente densos. La mayor parte del tiempo todavía se gasta en el despacho y la combinación de MoE, lo que indica que cualquier mejora adicional en los kernels mejorará en consecuencia las latencias y facilitará el micro-batching.

Trabajo futuro

Estamos trabajando estrechamente con los ingenieros de AWS para mejorar el rendimiento en EFA. Por nuestra parte, seguiremos las actualizaciones de libfabric que reducirán la sobrecarga del plano de datos. Además, experimentaremos con efa-direct para reducir aún más la sobrecarga del espacio de usuario e interactuar directamente con el hardware.

Referencias

Informe técnico de DeepSeek R1

Informe técnico de Kimi K2

DeepEP

UCCL-EP

Mejora del rendimiento de red de sistemas HPC utilizando NVIDIA Magnum IO NVSHMEM y GPUDirect Async

Citación

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