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 grandes modelos de mezcla de expertos de código abierto, como Kimi-K2, plantean desafíos 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 de múltiples nodos. Presentamos un conjunto de núcleos para el paralelismo de expertos que logran latencias de vanguardia en ConnectX-7, superando el rendimiento de DeepEP. Los mismos núcleos también son los primeros en lograr latencias viables en el adaptador de tejido elástico (EFA) de AWS, lo que permite implementaciones de modelos de un billón de parámetros.

Prueba nuestros núcleos en GitHub y lee el artículo de investigación completo en arXiv.

Introducción

La mezcla de expertos (MoE) se ha convertido en la arquitectura estándar para escalar a cientos de miles de millones o billones de parámetros, al tiempo que ofrece una latencia razonable en el momento de la inferencia. Los modelos MoE reemplazan la capa densa de un transformador con un conjunto de expertos y una pequeña capa de enrutamiento lineal que decide con qué experto multiplicar un token. Al agrupar los tokens asociados con el mismo experto, el problema se vuelve disperso y paralelizable de forma trivial, ya que ahora los diferentes expertos se pueden alojar en diferentes GPU distribuidas en múltiples nodos.

A diferencia de otras formas de paralelismo, como el paralelismo de tensores (TP) y el paralelismo de datos (DP), que se pueden implementar fácilmente utilizando bibliotecas de comunicación colectiva como torch.distributed y NCCL, el enrutamiento de MoE implica una comunicación dispersa de igual a igual (peer-to-peer). Un núcleo de envío debe dividir el conjunto de tokens entrantes y enviarlos a los rangos que alojan a los expertos, mientras que su problema dual, el núcleo de combinación, debe recibir sus homólogos procesados y calcular su promedio ponderado. Estas transferencias no se asignan a colectivos y requieren núcleos personalizados para lograr una baja latencia.

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

En este artículo, analizamos en profundidad la implementación de nuestros nuevos núcleos entre nodos, lo que mejora enormemente el rendimiento de nuestro trabajo anterior basado en NVSHMEM. Además de mostrar resultados de vanguardia en ConnectX-7, también demostramos que, dentro de nuestro tiempo de ejecución de inferencia, los núcleos se pueden utilizar 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 de manera efectiva dentro de un solo nodo, ya que NVLink ofrece un alto rendimiento (900 GB/s) con una latencia de microsegundos. Si bien los sistemas NVIDIA GB200 y GB300 están equipados con suficiente memoria para contener los pesos y los cachés KV de modelos de un billón de parámetros, las implementaciones en la instancia p5en de AWS con hasta 8 GPU H200 se limitan a 1120 GB de HBM, que deben dividirse entre los pesos del modelo y los cachés KV. En consecuencia, ciertos modelos requieren implementaciones entre nodos a través de InfiniBand.

Aunque InfiniBand sigue siendo increíblemente rápido, solo ofrece un rendimiento de 400 Gbps (50 GB/s), lo que añade de decenas a cientos de microsegundos de latencia para el enrutamiento de MoE. Afortunadamente, las latencias adicionales se pueden ocultar mediante el uso de expertos compartidos, micro lotes (micro-batching) y la superposición de computación y comunicación para ofrecer un rendimiento viable a un costo 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 de tejido elástico (EFA) personalizados. A pesar de que estos ofrecen un rendimiento máximo de 400 Gbps en los colectivos que se usan típicamente en las cargas de trabajo de entrenamiento, se quedan ligeramente cortos en comparación con ConnectX-7 en los tamaños de los mensajes intercambiados durante el envío y la combinación de MoE. Además, EFA no es compatible con GPUDirect Async, lo que requiere un hilo proxy de CPU que conecte la GPU con la tarjeta de red (NIC) para iniciar las transferencias. Además de la complejidad añadida, las transacciones PCIe adicionales agregan una sobrecarga de varios microsegundos a todas las transferencias. Los marcos como NVSHMEM permiten a los creadores de núcleos ofrecer implementaciones independientes del dispositivo; sin embargo, las API genéricas generan sobrecargas que hacen que las implementaciones sean ineficaces tanto en términos de costo como de latencia.

pplx-kernels basados en NVSHMEM

Nuestros núcleos anteriores abstrajeron la implementación de transporte subyacente a través de NVSHMEM. Se basaron en nvshmemx_putmem_signal_nbi_warp para transferir tokens uno por uno, utilizando contadores atómicos para sincronizar pares. Si bien esto ofreció un rendimiento razonable en las NIC ConnectX-7 con IBGDA, la implementación IBRC de ConnectX-7 basada en proxy fue significativamente más lenta, mientras que la implementación de EFA no logró ofrecer latencias inferiores a un milisegundo.

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

Siguiendo estas observaciones, diseñamos nuestros núcleos basados en EFA que exponen la misma interfaz, optimizando las interacciones entre el núcleo 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 núcleos siguen una arquitectura híbrida de CPU-GPU, donde los núcleos de GPU interactúan con el modelo en el dispositivo y un hilo proxy que se ejecuta en la CPU del host maneja las interacciones con la NIC. Los nuevos núcleos reutilizan el mismo TransferEngine que construimos para las transferencias de caché KV.

El enrutamiento MoE lo realiza un núcleo de envío, que recibe del modelo los tokens y los expertos por token para enrutarlos, enviándolos a través de la red a los rangos que alojan a los expertos. Después de que los núcleos GEMM agrupados multiplican los tokens con los pesos de los expertos, el núcleo de combinación transfiere los tokens a sus rangos de origen y calcula su promedio ponderado antes de devolverlos al modelo.

Tanto los núcleos de envío como los de combinación se dividen en una mitad emisora y una mitad receptora para facilitar la superposición y los micro lotes. El emisor reorganiza los tokens en búferes de envío para que puedan escribirse en los pares mediante una sola operación de escritura RDMA. Después de ejecutar el emisor, la GPU puede realizar otro trabajo (como la multiplicación de expertos compartidos) mientras las transferencias a través de la red están pendientes en segundo plano. Los núcleos receptores se bloquean hasta que se hayan recibido todas las transferencias. En segundo plano, un hilo de trabajo del lado del host maneja las transferencias: la GPU le avisa cuando se han poblado los búferes de envío, después de lo cual envía las escrituras RDMA. Una vez que la NIC señala la finalización, desbloquea al receptor, el cual devuelve los tokens al modelo.

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

Si bien 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 de RDMA, dentro de un nodo los núcleos 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 parte para EP64, logrando en general menores latencias.

Búferes de envío y recepción

El diseño de los buffers de envío para el kernel de despacho es crucial para lograr latencias viables que involucren al proxy del host. Si hay N ranks y cada uno despacha T tokens, cada uno enrutado a R expertos de un total de E, un rank puede recibir hasta N * T * max(R, E / N) tokens en el peor de los casos, si todos se enrutan hacia él. Además, este límite mínimo solo puede lograrse cuando los emisores organizan sus escrituras de forma contigua, ya que los buffers de recepción privados ocuparían demasiada memoria. Para que los emisores puedan organizar sus escrituras de manera consecutiva sin conflictos, es necesario distribuir a cada rank el conocimiento completo de la información de enrutamiento. En consecuencia, nuestros kernels de despacho primero intercambian los conteos de tokens por experto; después, cada uno puede determinar a dónde van sus escrituras y qué desplazamiento recibe los tokens de cada rank.

Intercambiar solo la información de enrutamiento daría como resultado latencias innecesariamente altas, ya que cada rank tendría que esperar a que llegaran las escrituras antes de despachar cualquiera de sus propias escrituras. En su lugar, reservamos una pequeña cantidad de espacio privado por emisor en cada rank receptor y despachamos algunos de los tokens junto con las escrituras de la información de enrutamiento. Esto garantiza una utilización total del ancho de banda durante la mayor parte del proceso de despacho. Después de intercambiar las rutas, los tokens restantes se transfieren mediante una sola escritura por cada par remoto. El intercambio de la información de enrutamiento y los buffers de tamaño fijo, junto con las transferencias que los llenan, se ilustran en la figura siguiente.

Ampliar

Los núcleos de combinación son sustancialmente más simples, ya que la información de enrutamiento ya está disponible para ellos desde la etapa de envío. Pueden copiar los tokens en una sola escritura por par, después de lo cual leen del búfer de recepción para calcular el promedio ponderado de los tokens en los tensores de salida.

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

Envío y combinación (Dispatch and Combine)

Las siguientes figuras detallan la sincronización y las interacciones entre la GPU, la CPU y la NIC, mostrando una iteración de envío y combinación intercalada con los núcleos del modelo.

El trabajo comienza en el envío (dispatch send), con el núcleo agregando la información de enrutamiento y calculando los recuentos de tokens por experto mediante acumulación atómica en la memoria compartida dentro de un solo 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 enviar la información de enrutamiento a todos los demás rangos pares. En paralelo, el envío 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 que los desplazamientos están disponibles, todos los tokens que se van a enviar 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 (intra-node), los núcleos también empujan inmediatamente los tokens a los 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 envío y la combinación. Este trabajo se realiza en la CPU mientras las transferencias ya están pendientes, por lo que no es crítico para la latencia. Posteriormente, una vez que se conocen todos los desplazamientos, se publican las escrituras para el resto de los tokens que exceden el tamaño del búfer privado y la fase de envío del encauzamiento se completa.

En el extremo receptor, el núcleo primero extrae del mismo nodo los tokens que superan 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 que se recibe la señal, los tokens se reorganizan desde los búferes de recepción al tensor de salida pasado al GEMM de grupo. Dado que los búferes se reutilizan entre el envío y la combinación, tanto las barreras de NVLink como las de RDMA garantizan que se hayan leído los búferes receptores antes de permitir que el siguiente núcleo los sobrescriba.

Los núcleos de combinación siguen un patrón similar, reorganizando los tokens, señalando al proxy y esperando a que se complete. 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 sola escritura. Antes de devolver los resultados, el núcleo acumula los tokens de los búferes receptores y los escribe en los tensores de salida.

TransferEngine para enrutamiento MoE

Para crear un prototipo de nuestra implementación, inicialmente reutilizamos el Transfer Engine construido para la precarga (prefill) y decodificación desagregadas con el fin de emitir y rastrear escrituras RDMA. Sin embargo, a diferencia de la decodificación desagregada, donde cada nodo publica una gran cantidad de escrituras en el mismo par, el envío MoE publica pequeñas transferencias a todos los pares en una sola operación. De forma predeterminada, el rendimiento no era óptimo y nuestra implementación se quedaba atrás de DeepEP en ConnectX-7 en aproximadamente 20 us. En consecuencia, especializamos el TransferEngine para admitir dos operaciones, dispersión (scatter) y barrera (barrier), para optimizar las interacciones con una gran cantidad de pares.

Tanto las operaciones de dispersión como las de barrera exponen grupos de pares que se pueden registrar previamente para rellenar de antemano las estructuras de datos utilizadas por la implementación de transporte subyacente (plantillas 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 de forma masiva ahorra valiosos microsegundos 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 de grupos entre las dos NIC en lugar de fragmentar los bytes de cada transferencia.

Aunque inicialmente desarrollamos TransferEngine únicamente en EFA utilizando libfabric, los núcleos MoE funcionan en ConnectX-7 sin cambiar una sola línea de código después de agregar soporte para ConnectX-7 en TransferEngine utilizando libibverbs. En comparación con EFA basado en Datagramas Confiables Escalables (SRD), ConnectX-7 requiere trabajo adicional para la configuración de la conexión y la gestión de pares. Usamos un par de colas de datagramas no confiables (UD) para iniciar la conexión confiable (RC). Creamos dos pares de colas para cada par remoto, uno para SEND/RECV bidireccional y otro para WRITE_IMM unidireccional, porque tanto RECV como WRITE_IMM consumen la operación RECV 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. Al igual que con EFA, implementamos la optimización de plantillas WR. Además, al enviar múltiples mensajes pequeños al mismo par, usamos el encadenamiento WR para mejorar la tasa de mensajes.

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

Evaluación

Latencia de envío/combinación (Dispatch/Combine Latency)

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

Ampliar

En la decodificación, a pesar de que nuestros núcleos son unos microsegundos más lentos que DeepEP, logramos latencias generales más bajas ya que nuestros núcleos 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 de paquete típico intercambiado por el enrutamiento MoE, la brecha entre los dos adaptadores es en general menos sustancial.

Ampliar

Aunque nuestros núcleos optimizados para la decodificación pueden escalar a tamaños de lote de precarga sin ningún ajuste adicional, el rendimiento se está quedando atrás de DeepEP. Esto se debe principalmente a que acumulamos tokens solo en el lado receptor de la combinación, mientras que DeepEP se 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 núcleos transfieren datos en una sola escritura, utilizamos más memoria para los búferes de envío y recepción.

Rendimiento (Throughput) de Kimi-K2 / DeepSeek-V3

Para demostrar la viabilidad de nuestros núcleos, 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 adaptadores de tejido elástico (EFA) y GPU H200. Ejecutamos el modelo en modo de verificación de borrador (draft verification) y reportamos la latencia asumiendo una tasa de aceptación del 80% usando una sola capa MTP.

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

En Kimi-K2, los núcleos permiten dar servicio. Debido a su tamaño, el modelo no puede caber dentro del mismo nodo utilizando el 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 cabezas de atención, logramos mejores latencias que DeepSeek-V3 y podemos servir el modelo en EFA con una latencia viable.

Ampliar

Ampliar

También hemos considerado los micro lotes, aunque solo mejoraron las latencias en tamaños de lote muy altos o al usar más de un token de borrador. Para beneficiarse de la división de problemas GEMM agrupados en dos, los tamaños de los lotes de tokens deben ser suficientemente densos. La mayor parte del tiempo todavía se dedica al envío y la combinación de MoE, lo que indica que cualquier mejora adicional en los núcleos mejorará en consecuencia las latencias y facilitará los micro lotes.

Trabajo futuro

Estamos trabajando en estrecha colaboración 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 las sobrecargas 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 los 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}
}