Comunicación punto a punto RDMA para sistemas de LLM

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

Nuestro último artículo de investigación presenta TransferEngine, una biblioteca de comunicación basada en RDMA portátil que funciona en diferentes Controladores de Interfaces de Red (NIC), incluidos NVIDIA ConnectX-7 y AWS Elastic Fabric Adapter (EFA). Demostramos su eficacia en tres sistemas de producción en Perplexity AI:

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

Transferencia de pesos para post-entrenamiento de RL en menos de 2 segundos: sincronización de pesos ultrarrápida entre nodos de entrenamiento e inferencia.

Habilitación de modelos de un billón de parámetros en AWS EFA: operaciones de distribución y combinación de MoE de baja latencia mediante primitivas de dispersión y barrera, logrando una latencia de vanguardia en ConnectX-7 y la primera implementación viable compatible con AWS EFA.

Este artículo se centra en TransferEngine en sí, 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 cargas de trabajo emergentes:

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

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

Orden de operaciones: Las API colectivas imponen semánticas sincronizadas, incluso cuando las aplicaciones solo necesitan saber que un conjunto de operaciones ha concluido. Por ejemplo, en la transferencia de KvCache, todas las páginas de la instancia de prefill a la instancia de decodificación pueden llegar en cualquier orden. De manera similar, en las actualizaciones de pesos de RL, todos los tensores de pesos 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 aún heredan semánticas y limitaciones colectivas.

Diversidad de hardware: Las NIC NVIDIA ConnectX con transporte de Conexión Confiable (RC) son la configuración de RDMA más común, garantizando la entrega en orden. Los proveedores de la nube como AWS utilizan hardware diferente; Elastic Fabric Adapter (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 de software: Las bibliotecas de RDMA existentes sufren de dependencia del proveedor (vendor lock-in). DeepEP requiere RDMA iniciado por GPU (IBGDA) disponible solo en ConnectX. NVSHMEM exhibe una degradación severa del rendimiento en EFA. Bibliotecas como Mooncake y NIXL carecen de soporte para EFA o se encuentran en desarrollo temprano. Como resultado, no existe una solución viable entre proveedores para la inferencia de LLM.

TransferEngine aborda estas barreras mediante cuatro capacidades principales:

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

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

Soporte de múltiples 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 suficientes 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, GPUs y NICs en dos dominios GPU/NUMA.

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

Cada instancia de TransferEngine expone una única dirección principal para la identificación y el descubrimiento de pares. Los pares remotos intercambian estructuras NetAddr que codifican direcciones de red y se pueden serializar 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 de carga automático en 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 en la NIC antes de la transferencia de datos. La función reg_mr toma un puntero, una longitud y un dispositivo, y devuelve dos identificadores:

MrHandle: Un identificador local utilizado como origen para las transferencias

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

La estructura MrDesc contiene una lista de direcciones de red remotas y una clave remota (una por NIC) que otorgan acceso remoto a la región de memoria. Este diseño maneja de forma natural configuraciones de múltiples 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 (callback) con los datos recibidos. El búfer se devuelve automáticamente al grupo después de que se completa la devolución de llamada.

Para transferencias de datos masivos, TransferEngine proporciona operaciones de escritura unidireccional sin copias (zero-copy):

submit_single_write: Transfiere una región contigua del origen al destino

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

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 transferencias grandes en múltiples NIC para maximizar el ancho de banda.

Para operaciones dirigidas a múltiples pares (comunes en el enrutamiento de 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 porciones de datos a cada par del grupo

submit_barrier: Envía una notificación exclusiva de valor inmediato a todos los pares (transferencia de cero bytes)

Estas operaciones están optimizadas para una 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 de 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 UVM Watchers: ubicaciones de memoria virtual unificada que pueden ser actualizadas por los núcleos de la GPU (incluso dentro de gráficos CUDA) y sondeadas eficientemente por la CPU mediante GDRCopy. Cuando se asigna un observador de 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 del transformador, la GPU incrementa el observador, lo que desencadena una transferencia RDMA de la KvCache de esa capa mientras las capas posteriores continúan computando.

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 en el lado del iniciador: Cuando se envía una operación de escritura, el llamador puede elegir 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 exitosamente desde el búfer de origen.

Notificación en el 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 transporta un valor inmediato, el ImmCounter del receptor se incrementa de forma atómica al entregar 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 recuento objetivo. Los contadores se pueden sincronizar de forma transparente con la GPU a través de GDRCopy, sondear directamente o monitorear mediante devoluciones de llamada.

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

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

La API reenvía las solicitudes al DomainGroup que atiende al dispositivo correspondiente, y el primero también atiende al host. En un bucle ajustado (tight loop), el trabajador de dominio busca trabajo nuevo, priorizando el envío de nuevas solicitudes. Dependiendo del hardware y la configuración, las solicitudes se fragmentan y se equilibran en 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 canalización (pipeline) de hardware. Posteriormente, se sondean las colas de finalización para consultar las transferencias finalizadas y los incrementos de contadores inmediatos. Los eventos se agregan para ofrecer 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 dirigirse 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 fábrica por NIC dentro del DomainGroup. Las decisiones de implementación clave y las optimizaciones incluyen:

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

Creación de plantillas de solicitudes de trabajo: Pre-puebla y retiene los campos comunes de las estructuras de libfabric antes de publicar, reduciendo la sobrecarga por operación para transferencias masivas y operaciones de grupos 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 (handshakes): Cada Domain utiliza un par de colas de Datagramas No Confiables (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 Send/Recv bidireccionales y otro para operaciones Write y WriteImm unidireccionales. 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 una semántica de Recv de alto nivel al tiempo que admite WriteImm sin interferencias.

Creación de plantillas de solicitudes de trabajo: Pre-puebla los campos comunes de las estructuras de solicitudes 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 timbres de puerta (doorbell rings) a la NIC.

Orden 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 garantías en 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 sistemas de LLM. La biblioteca es de código abierto y está disponible en GitHub, y alentamos a la comunidad a explorarla 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 la API

Registro de memoria

Send/Recv bidireccional

Escritura unidireccional

Grupos de pares y operaciones masivas

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

Notificación de finalización

Implementación

Optimización específica de hardware

Conclusión

Citación