Comunicación RDMA punto a punto para sistemas LLM

Los sistemas LLM tradicionalmente han dependido de comunicaciones colectivas para patrones paralelos estructurados en entrenamiento e inferencia. Sin embargo, las cargas de trabajo emergentes, tales como la inferencia desagregada, el enrutamiento de Mezcla de Expertos (MoE) y el ajuste fino de aprendizaje por refuerzo asíncrono, exigen la flexibilidad de la comunicación punto a punto.

Nuestro artículo de investigación más reciente presenta TransferEngine, una biblioteca de comunicación basada en RDMA portátil que funciona a través de diferentes Controladores de Interfaz de Red (NIC), incluyendo NVIDIA ConnectX-7 y el Adaptador de Estructura Elástica de AWS (EFA). Demostramos su eficacia en tres sistemas de producción en Perplexity AI:

Prefill y Decodificación desagregados: Transferencia de KvCache capa por capa utilizando escrituras paginadas y observadores UVM, lo que permite que las instancias de prefill y decodificación escalen de forma independiente.

Transferencia de pesos para post-entrenamiento de RL en menos de 2 segundos: Sincronización de pesos sumamente rápida entre los nodos de entrenamiento y de inferencia.

Habilitación de modelos de un billón de parámetros en AWS EFA: Operaciones de envío y combinación de MoE de baja latencia mediante primitivas de dispersión y barrera, logrando una latencia de última generación en ConnectX-7 y la primera implementación viable compatible con AWS EFA.

Este artículo se centra en TransferEngine, la biblioteca de comunicación que impulsa estos sistemas. Lea el artículo de investigación completo en arXiv y explore el código en GitHub.

Las comunicaciones colectivas destacan en patrones paralelos estructurados como el Paralelismo de Tensores y el Paralelismo de Datos, pero imponen restricciones que limitan su aplicabilidad a las cargas de trabajo emergentes:

Membresía fija: Los colectivos requieren establecer un entorno de comunicación global en el momento de la inicialización, con una configuración sincronizada entre todos los participantes. Para la inferencia desagregada, esto es problemático ya que el número de instancias de prefill y decodificación debe escalar dinámicamente en respuesta a los patrones de tráfico.

Formas de tensor uniformes: Los colectivos asumen que todos los participantes intercambian tensores con formas y tipos de datos idénticos. En el enrutamiento MoE, el número de tokens enviados a cada experto depende de los datos y se descubre durante el tiempo de ejecución, lo que obliga a los colectivos densos de todos a todos a rellenar hasta los tamaños de búfer en el peor de los casos.

Orden de las operaciones: Las API colectivas imponen semánticas sincronizadas, incluso cuando las aplicaciones solo necesitan saber que un conjunto de operaciones se ha completado. Por ejemplo, en la transferencia de KvCache, todas las páginas desde la instancia de prefill a la de decodificación pueden llegar en cualquier orden. De manera similar, en las actualizaciones de pesos de RL, todos los tensores de peso se pueden transferir de forma independiente sin restricciones de orden.

La comunicación punto a punto basada en RDMA aborda naturalmente estas limitaciones, pero tres barreras han impedido su adopción:

Limitaciones de la biblioteca: Aunque torch.distributed y NCCL proporcionan API punto a punto como Send/Recv, estas operaciones todavía heredan semánticas y limitaciones colectivas.

Diversidad de hardware: Las NIC NVIDIA ConnectX con transporte de Conexión Confiable (RC) son la configuración RDMA más común, lo que garantiza la entrega en orden. Los proveedores de la nube como AWS utilizan hardware diferente, el Adaptador de Estructura Elástica (EFA) con transporte de Datagrama Confiable Escalable (SRD) entrega mensajes fuera de orden. EFA también requiere agregar cuatro NIC de 100 Gbps (o dos NIC de 200 Gbps) para lograr un ancho de banda de 400 Gbps.

Portabilidad del software: Las bibliotecas RDMA existentes sufren de dependencia de proveedores. DeepEP requiere RDMA iniciado por GPU (IBGDA) disponible solo en ConnectX. NVSHMEM exhibe una grave degradación del rendimiento en EFA. Bibliotecas como Mooncake y NIXL carecen de soporte para EFA o permanecen en desarrollo temprano. Como resultado, no existe una solución viable entre proveedores para la inferencia de LLM.

TransferEngine aborda estas barreras a través de cuatro capacidades principales:

Supuestos de transporte comunes: En lugar de apuntar a una NIC específica, TransferEngine asume únicamente que el transporte subyacente es confiable pero desordenado, es decir, la intersección de las garantías proporcionadas por los transportes RC y SRD.

Sincronización de ImmCounter: Una primitiva novedosa para la notificación de finalización que funciona sin depender del orden de los mensajes, permitiendo que el mismo código se ejecute tanto en transportes ordenados (RC) como desordenados (SRD).

Soporte multi-NIC: Gestión transparente de múltiples NIC por GPU, fragmentando automáticamente las transferencias para maximizar el ancho de banda.

Superficie de API mínima: Un pequeño conjunto de operaciones suficiente para construir patrones de comunicación complejos.

Diseño del sistema TransferEngine, que muestra las rutas de solicitud y devolución de llamada, trabajadores de CPU, buffers UVM, GPU y NIC en dos dominios GPU/NUMA.

TransferEngine genera un hilo de trabajo por GPU. Cada trabajador gestiona un DomainGroup que coordina de 1 a 4 NIC RDMA, dependiendo de la plataforma de hardware (una sola ConnectX-7 proporciona 400 Gbps, mientras que EFA requiere cuatro adaptadores de 100 Gbps o dos de 200 Gbps). Dentro de un DomainGroup, cada Domain es específico del hardware y gestiona una sola NIC, manejando la gestión de pares de colas, el envío de trabajo y el sondeo de finalización.

Cada instancia de TransferEngine expone una dirección principal única para la identificación y el descubrimiento de pares. Los pares remotos intercambian estructuras NetAddr que codifican direcciones de red y pueden serializarse para la comunicación fuera de banda. Todos los pares en un sistema deben usar el mismo número de NIC por GPU. Esta configuración uniforme permite a TransferEngine fragmentar las transferencias de forma inteligente: al iniciar una transferencia, tanto el origen como el destino saben exactamente qué NIC están disponibles, lo que permite el equilibrio automático de carga entre todos los adaptadores.

TransferEngine expone una API mínima que abstrae la complejidad de las interfaces RDMA subyacentes.

rust

#[serde] struct NetAddr(Bytes);
#[serde] struct MrDesc{ ptr: u64, rkeys: Vec<(NetAddr, u64)> }
struct MrHandle(NonNull<c_void>); type Offset = u64;
struct Pages{ indices: Vec<u32>, stride: u64, offset: Offset }
struct PeerGroupHandle(u64);
struct ScatterDst{ len: u64, src: Offset, dst: (MrDesc,Offset)}
enum OnDone { Callback(fn () -> ()), Flag(Atomic<bool>) }

trait TransferEngine {
  fn main_address() -> NetAddr;
  // Memory Region Management
  fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc);
  // Two-sided Send/Recv
  fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ());
  fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ());
  // One-sided Write
  fn expect_imm_count(imm: u32, count: u32, cb: fn () -> ());
  fn submit_single_write(len: u64, imm: Option<u32>
      src: (MrHandle, Offset), dst: (MrDesc, Offset), OnDone);
  fn submit_paged_writes(page_len: u64, imm: Option<u32>,
      src: (MrHandle, Pages), dst: (MrDesc, Pages), OnDone);
  // One-sided Write to a group of peers
  fn add_peer_group(addrs: Vec<NetAddr>) -> PeerGroupHandle;
  fn submit_scatter(h: Option<PeerGroupHandle>, OnDone,
      imm: Option<u32>, src: MrHandle, dst: Vec<ScatterDst>);
  fn submit_barrier(h: Option<PeerGroupHandle>, OnDone,
      imm: u32, dst: Vec<MrDesc>);
  // Watcher for CPU-GPU synchronization
  fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>;
}

La API consta de varios componentes clave:

Las operaciones RDMA requieren que las regiones de memoria se registren con la NIC antes de la transferencia de datos. La función reg_mr toma un puntero, una longitud y un dispositivo, devolviendo dos manejadores:

MrHandle: Un manejador local utilizado como origen para las transferencias

MrDesc: Un descriptor serializable que puede intercambiarse con pares remotos para permitirles escribir en la memoria registrada

La estructura MrDesc contiene una lista de direcciones de red remotas y claves remotas (una por NIC) que otorgan acceso remoto a la región de memoria. Este diseño maneja naturalmente las configuraciones multi-NIC.

Para mensajes pequeños y mensajes de plano de control estilo RPC, TransferEngine proporciona submit_send y submit_recvs:

submit_send: Envía el mensaje a un par remoto identificado por NetAddr. El mensaje se copia internamente, por lo que el llamador puede reutilizar o liberar inmediatamente el búfer.

submit_recvs: Publica un grupo de búferes de recepción. Cuando llega un mensaje, se invoca una devolución de llamada con los datos recibidos. El búfer se devuelve automáticamente al grupo después de que se complete la devolución de llamada.

Para transferencias de datos masivos, TransferEngine proporciona operaciones de escritura de un lado sin copia:

submit_single_write: Transfiere una región contigua desde el origen al destino

submit_paged_writes: Transfiere páginas no contiguas especificadas por índices, pasos y desplazamientos

Ambas operaciones aceptan un valor inmediato opcional de 32 bits que se entrega al receptor de forma atómica al finalizar.

El motor fragmenta automáticamente las grandes transferencias a través de múltiples NIC para maximizar el ancho de banda.

Para operaciones dirigidas a múltiples pares (comunes en el enrutamiento MoE), TransferEngine proporciona primitivas optimizadas:

add_peer_group: Pre-registra un conjunto de direcciones de pares para una búsqueda rápida

submit_scatter: Envía diferentes fragmentos de datos a cada par en el grupo

submit_barrier: Envía una notificación solo inmediata a todos los pares (transferencia de cero bytes)

Estas operaciones están optimizadas para baja latencia mediante el almacenamiento en caché de estructuras de datos internas y la creación de plantillas de solicitudes de trabajo RDMA, evitando la sobrecarga de configuración repetitiva.

Un desafío clave en los sistemas LLM es coordinar las transferencias RDMA con el progreso de la GPU. Idealmente, la transferencia de datos debería comenzar tan pronto como la GPU produzca resultados, pero sondear la GPU desde la CPU es ineficiente.

TransferEngine aborda esto con observadores UVM: ubicaciones de memoria virtual unificada que pueden ser actualizadas por núcleos de GPU (incluso dentro de gráficos CUDA) y sondeadas eficientemente por la CPU utilizando GDRCopy. Cuando se asigna un observador UVM, se registra una devolución de llamada que se invoca cada vez que la GPU actualiza el valor.

Esto permite patrones como la transferencia de KvCache capa por capa: después de que se completa cada capa de transformador, la GPU incrementa el observador, lo que desencadena una transferencia RDMA de la KvCache de esa capa mientras las capas posteriores continúan calculando.

TransferEngine proporciona mecanismos de notificación de finalización tanto para el lado del iniciador (el llamador de submit_write) como para el lado del receptor (el destino de las operaciones de escritura).

Notificación del lado del iniciador: Cuando se envía una operación de escritura, el llamador puede optar por ser notificado al finalizar mediante una función de devolución de llamada o una bandera atómica. Esta notificación confirma que los datos se han transferido correctamente desde el búfer de origen.

Notificación del lado del receptor: La innovación clave es el ImmCounter, un componente que rastrea contadores indexados por números inmediatos de 32 bits. Cuando una operación de escritura conlleva un valor inmediato, el ImmCounter del receptor se incrementa de forma atómica al momento de la entrega de los datos. La función expect_imm_count permite al receptor registrar una devolución de llamada que se activa cuando un valor inmediato específico alcanza un conteo objetivo. Los contadores se pueden sincronizar de forma transparente con la GPU a través de GDRCopy, sondearse directamente o monitorearse a través de devoluciones de llamada.

Este mecanismo del lado del receptor proporciona una sincronización confiable sin requerir el orden de los mensajes. Funciona igual de bien en las NIC ConnectX (que garantizan la entrega en orden sobre el transporte RC) y las NIC EFA (que entregan mensajes fuera de orden sobre el transporte SRD).

TransferEngine está implementado en Rust, optimizando cuidadosamente las asignaciones, los subprocesos y la sincronización para minimizar la latencia y lograr un alto rendimiento. El motor genera un hilo de trabajo por DomainGroup, cada uno fijado a un núcleo de CPU en el nodo NUMA al que están conectados los dispositivos, minimizando tanto la programación como la latencia de acceso a la memoria. Las estructuras de datos específicas de un dominio se asignan después de la fijación para garantizar que la memoria se reserve en el nodo NUMA correcto. Un hilo de trabajo maneja hasta 4 Domains (cada uno gestionando una sola NIC), mientras que otro hilo dedicado sondea la GPU para actualizar los observadores UVM. La comunicación entre hilos se realiza a través de colas sin bloqueo.

La API reenvía las solicitudes al DomainGroup que sirve al dispositivo apropiado, siendo el primero el que también sirve al host. En un bucle cerrado, el trabajador del dominio sondea nuevo trabajo, priorizando el envío de nuevas solicitudes. Dependiendo del hardware y la configuración, las solicitudes se fragmentan y se equilibran la carga entre los Domains disponibles. La primera escritura de una solicitud compuesta se publica inmediatamente en la cola de envío de la NIC. Una vez que se agotan las nuevas solicitudes, el trabajador procede a avanzar en las operaciones pendientes, publicando escrituras para llenar la tubería de hardware. Posteriormente, se sondean las colas de finalización para consultar las transferencias finalizadas y los incrementos de contador inmediatos. Los eventos se agregan para entregar notificaciones por transferencia, entregando la transferencia a un hilo de devolución de llamada dedicado compartido por todos los grupos.

La fragmentación dentro de un DomainGroup es flexible. Las transferencias pueden apuntar a NIC específicas por índice. Una sola escritura se puede dividir. Las transferencias paginadas, la dispersión y las operaciones de barrera, que se traducen en múltiples escrituras, pueden fragmentarse en todas las NIC.

Los Domains dentro de TransferEngine están especializados y optimizados para el hardware bajo su control:

AWS EFA: La implementación de EFA utiliza libfabric, gestionando un dominio de tejido por NIC dentro del DomainGroup. Las decisiones de implementación clave y las optimizaciones incluyen:

Escrituras de tamaño cero: EFA diverge de la especificación RDMA al requerir descriptores de destino válidos incluso para escrituras de tamaño cero. Por lo tanto, TransferEngine impone descriptores válidos para todas las transferencias.

Plantillas de solicitud de trabajo: Pre-puebla y conserva los campos comunes de las estructuras libfabric antes de publicar, reduciendo la sobrecarga por operación para transferencias masivas y operaciones de grupo de pares

NVIDIA ConnectX-7: La implementación de ConnectX-7 utiliza libibverbs con transporte RC. Las decisiones de implementación clave y las optimizaciones incluyen:

Pares de colas UD para apretones de manos: Cada Domain utiliza un par de colas de Datagrama No Confiable (UD) para intercambiar apretones de manos de conexión RC.

Pares de colas RC separados: Crea 2 pares de colas RC por par, uno para operaciones de envío/recepción de dos lados y otro para operaciones de escritura y WriteImm de un lado. Esta separación es necesaria porque tanto las finalizaciones de Recv como de WriteImm consumen solicitudes de trabajo en orden de publicación, lo que permite semánticas de Recv de alto nivel mientras se admite WriteImm sin interferencias.

Plantillas de solicitud de trabajo: Pre-puebla campos comunes de estructuras de solicitud de trabajo para transferencias masivas y grupos de pares.

Encadenamiento de solicitudes de trabajo: Vincula hasta 4 solicitudes de trabajo a través del puntero next de ibv_send_wr, reduciendo el número de timbrazos a la NIC.

Ordenamiento relajado: Habilita IBV_ACCESS_RELAXED_ORDERING para permitir transacciones PCIe fuera de orden entre la NIC y la memoria de la GPU.

Al apuntar a la intersección de las garantías a través de diferentes transportes RDMA e introducir una sincronización independiente del orden a través de ImmCounter, TransferEngine proporciona una base portátil para la comunicación punto a punto en los sistemas LLM. La biblioteca es de código abierto y está disponible en GitHub, y alentamos a la comunidad a explorar y construir sobre ella para patrones de comunicación emergentes en sistemas de aprendizaje automático a gran escala.

BibTeX

@online{perplexity2025RdmaPointToPointCommunicationForLlmSystems,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{RDMA point-to-point communication for LLM systems}},
  date = {2025-11-05},
  url = {https://www.perplexity.ai/hub/blog/rdma-point-to-point-communication-for-llm-systems}
}

Motivación

TransferEngine

Diseño de API

Registro de memoria

Envío/Recepción de dos lados

Escritura de un lado

Grupos de pares y operaciones masivas

Observador UVM: Señalización de GPU a CPU

Notificación de finalización

Implementación

Optimizaciones específicas de hardware

Conclusión

Citación