Comunicação ponto a ponto RDMA para sistemas de LLM
Ferramenta elegante para abordar padrões de comunicação de LLM
Tradicionalmente, os sistemas de LLM têm confiado em comunicações coletivas para padrões paralelos estruturados em treinamento e inferência. No entanto, cargas de trabalho emergentes, como inferência desagregada, roteamento de Mistura de Especialistas (MoE) e ajuste fino de aprendizado por reforço assíncrono, exigem a flexibilidade da comunicação ponto a ponto.
Nosso artigo de pesquisa mais recente apresenta o TransferEngine, uma biblioteca de comunicação baseada em RDMA portátil que funciona em diferentes Controladores de Interface de Rede (NICs), incluindo NVIDIA ConnectX-7 e AWS Elastic Fabric Adapter (EFA). Demonstramos sua eficácia em três sistemas de produção na Perplexity AI:
Prefill e Decode Desagregados: Transferência de KvCache camada por camada usando gravações paginadas e observadores UVM, permitindo que instâncias de prefill e decode sejam dimensionadas de forma independente.
Transferência de Pesos para Pós-Treinamento de RL em menos de 2 segundos: Sincronização ultrarrápida de pesos entre nós de treinamento e inferência.
Habilitando Modelos de Trilhão de Parâmetros no AWS EFA: Operações de despacho e combinação de MoE de baixa latência usando primitivas de dispersão e barreira, alcançando latência de ponta no ConnectX-7 e a primeira implementação viável compatível com AWS EFA.
Este artigo foca no próprio TransferEngine, a biblioteca de comunicação que alimenta esses sistemas. Leia o artigo de pesquisa completo no arXiv e explore o código no GitHub.
Motivação
As comunicações coletivas se destacam em padrões paralelos estruturados, como Paralelismo de Tensores e Paralelismo de Dados, mas impõem restrições que limitam sua aplicabilidade a cargas de trabalho emergentes:
Associação fixa: As coletivas exigem o estabelecimento de um mundo de comunicação global no momento da inicialização, com configuração sincronizada entre todos os participantes. Para inferência desagregada, isso é problemático, pois o número de instâncias de prefill e decode deve ser dimensionado dinamicamente em resposta aos padrões de tráfego.
Formatos de tensores uniformes: As coletivas assumem que todos os participantes trocam tensores com formatos e tipos de dados idênticos. No roteamento de MoE, o número de tokens enviados a cada especialista depende dos dados e é descoberto em tempo de execução, forçando coletivas all-to-all densas a preencher até tamanhos de buffer de pior caso.
Ordenação de operações: As APIs coletivas impõem semântica sincronizada, mesmo quando as aplicações precisam apenas saber que um conjunto de operações foi concluído. Por exemplo, na transferência de KvCache, todas as páginas da instância de prefill para a instância de decode podem chegar em qualquer ordem. Da mesma forma, nas atualizações de pesos de RL, todos os tensores de pesos podem ser transferidos independentemente, sem restrições de ordenação.
A comunicação ponto a ponto baseada em RDMA aborda naturalmente essas limitações, mas três barreiras impediram sua adoção:
Limitações da biblioteca: Embora torch.distributed e NCCL forneçam APIs ponto a ponto como Send/Recv, essas operações ainda herdam semânticas e limitações coletivas.
Diversidade de hardware: Os NICs NVIDIA ConnectX com transporte de Conexão Confiável (RC) são a configuração RDMA mais comum, garantindo entrega em ordem. Provedores de nuvem como a AWS usam hardware diferente; o Elastic Fabric Adapter (EFA) com transporte de Datagrama Confiável Escalável (SRD) entrega mensagens fora de ordem. O EFA também exige a agregação de quatro NICs de 100 Gbps (ou dois NICs de 200 Gbps) para atingir largura de banda de 400 Gbps.
Portabilidade de software: As bibliotecas RDMA existentes sofrem com o aprisionamento tecnológico (vendor lock-in). DeepEP exige RDMA iniciado por GPU (IBGDA) disponível apenas no ConnectX. NVSHMEM exibe degradação severa de desempenho no EFA. Bibliotecas como Mooncake e NIXL carecem de suporte ao EFA ou permanecem em desenvolvimento inicial. Como resultado, não existe uma solução viável entre provedores para inferência de LLM.
TransferEngine
O TransferEngine aborda essas barreiras por meio de quatro capacidades principais:
Suposições de transporte comuns: Em vez de direcionar um NIC específico, o TransferEngine assume apenas que o transporte subjacente é confiável, mas não ordenado, ou seja, a interseção de garantias fornecidas pelos transportes RC e SRD.
Sincronização ImmCounter: Uma nova primitiva para notificação de conclusão que funciona sem depender da ordenação de mensagens, permitindo que o mesmo código seja executado em transportes ordenados (RC) e não ordenados (SRD).
Suporte a múltiplos NICs: Gerenciamento transparente de múltiplos NICs por GPU, dividindo automaticamente as transferências para maximizar a largura de banda.
Superfície de API mínima: Um pequeno conjunto de operações suficientes para construir padrões complexos de comunicação.
O TransferEngine gera um tópico de trabalho por GPU. Cada trabalhador gerencia um DomainGroup que coordena de 1 a 4 NICs RDMA, dependendo da plataforma de hardware (um único ConnectX-7 fornece 400 Gbps, enquanto o EFA exige quatro adaptadores de 100 Gbps ou dois adaptadores de 200 Gbps). Dentro de um DomainGroup, cada Domain é específico de hardware e gerencia um único NIC, lidando com gerenciamento de pares de filas, submissão de trabalho e polling de conclusão.
Cada instância do TransferEngine expõe um único endereço principal para identificação e descoberta de pares. Pares remotos trocam estruturas NetAddr que codificam endereços de rede e podem ser serializados para comunicação fora de banda. Todos os pares em um sistema devem usar o mesmo número de NICs por GPU. Esta configuração uniforme permite que o TransferEngine divida as transferências de forma inteligente: ao iniciar uma transferência, tanto a origem quanto o destino sabem exatamente quais NICs estão disponíveis, permitindo o balanceamento de carga automático em todos os adaptadores.
Design da API
O TransferEngine expõe uma API mínima que abstrai a complexidade das interfaces RDMA subjacentes.
#[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>;
}A API consiste em vários componentes principais:
Registro de Memória
As operações RDMA exigem que as regiões de memória sejam registradas no NIC antes da transferência de dados. A função reg_mr aceita um ponteiro, comprimento e dispositivo, retornando dois identificadores:
MrHandle: Um identificador local usado como origem para transferências
MrDesc: Um descritor serializável que pode ser trocado com pares remotos para permitir que eles escrevam na memória registrada
A estrutura MrDesc contém uma lista de endereços de rede remotos e chaves remotas (uma por NIC) que concedem acesso remoto à região de memória. Este design lida naturalmente com configurações de múltiplos NICs.
Send/Recv Bilateral
Para mensagens pequenas e mensagens de plano de controle no estilo RPC, o TransferEngine fornece submit_send e submit_recvs:
submit_send: Envia a mensagem para um par remoto identificado por NetAddr. A mensagem é copiada internamente, de modo que o chamador pode reutilizar ou liberar imediatamente o buffer.
submit_recvs: Publica um pool de buffers de recebimento. Quando uma mensagem chega, um callback é invocado com os dados recebidos. O buffer é retornado automaticamente ao pool após a conclusão do callback.
Gravação Unilateral
Para transferências de dados em massa, o TransferEngine fornece operações de gravação unilateral sem cópia (zero-copy):
submit_single_write: Transfere uma região contígua da origem para o destino
submit_paged_writes: Transfere páginas não contíguas especificadas por índices, passos (strides) e desvios (offsets)
Ambas as operações aceitam um valor imediato opcional de 32 bits que é entregue ao receptor de forma atômica após a conclusão.
O mecanismo divide automaticamente grandes transferências em múltiplos NICs para maximizar a largura de banda.
Grupos de Pares e Operações em Massa
Para operações direcionadas a múltiplos pares (comuns no roteamento de MoE), o TransferEngine fornece primitivas otimizadas:
add_peer_group: Pré-registra um conjunto de endereços de pares para busca rápida
submit_scatter: Envia diferentes fatias de dados para cada par no grupo
submit_barrier: Envia uma notificação estritamente imediata para todos os pares (transferência de zero bytes)
Estas operações são otimizadas para baixa latência armazenando em cache estruturas de dados internas e criando modelos de solicitações de trabalho RDMA, evitando a sobrecarga de configuração repetitiva.
Observador UVM: Sinalização de GPU para CPU
Um desafio fundamental em sistemas de LLM é coordenar as transferências RDMA com o progresso da GPU. Idealmente, a transferência de dados deve começar assim que a GPU produzir saída, mas fazer o polling da GPU a partir da CPU é ineficiente.
O TransferEngine resolve isso com UVM Watchers: locais de memória virtual unificada que podem ser atualizados por kernels de GPU (mesmo dentro de grafos CUDA) e verificados eficientemente pela CPU usando o GDRCopy. Quando um observador UVM é alocado, um callback é registrado e invocado sempre que a GPU atualiza o valor.
Isso permite padrões como transferência de KvCache camada por camada: após a conclusão de cada camada do transformador, a GPU incrementa o observador, disparando uma transferência RDMA do KvCache dessa camada enquanto as camadas subsequentes continuam computando.
Notificação de Conclusão
O TransferEngine fornece mecanismos de notificação de conclusão tanto para o lado do iniciador (o chamador de submit_write) quanto para o lado do receptor (o alvo das operações de gravação).
Notificação do lado do iniciador: Quando uma operação de gravação é submetida, o chamador pode optar por ser notificado após a conclusão por meio de uma função de callback ou de uma flag atômica. Esta notificação confirma que os dados foram transferidos com sucesso do buffer de origem.
Notificação do lado do receptor: A inovação principal é o ImmCounter, um componente que rastreia contadores indexados por números imediatos de 32 bits. Quando uma operação de gravação carrega um valor imediato, o ImmCounter do receptor é incrementado de forma atômica após a entrega dos dados. A função expect_imm_count permite que o receptor registre um callback que é acionado quando um valor imediato específico atinge uma contagem alvo. Os contadores podem ser sincronizados transparentemente com a GPU via GDRCopy, consultados diretamente por polling ou monitorados por meio de callbacks.
Este mecanismo do lado do receptor fornece sincronização confiável sem exigir a ordenação de mensagens. Ele funciona igualmente bem em NICs ConnectX (que garantem entrega em ordem sobre o transporte RC) e NICs EFA (que entregam mensagens fora de ordem sobre o transporte SRD).
Implementação
O TransferEngine é implementado em Rust, otimizando cuidadosamente alocações, tópicos e sincronização para minimizar a latência e alcançar alto rendimento. O mecanismo gera um tópico de trabalho por DomainGroup, cada um fixado a um núcleo de CPU no nó NUMA ao qual os dispositivos estão conectados, minimizando a latência de agendamento e de acesso à memória. As estruturas de dados específicas de um domínio são alocadas após a fixação para garantir que a memória seja reservada no nó NUMA correto. Um tópico de trabalho gerencia até 4 Domains (cada um gerenciando um único NIC), enquanto outro tópico dedicado faz o polling da GPU para atualizar os observadores UVM. A comunicação entre tópicos é feita por meio de filas livres de bloqueio (lock-free).
A API encaminha solicitações para o DomainGroup que atende ao dispositivo apropriado, sendo que o primeiro também atende ao host. Em um loop restrito, o trabalhador de domínio faz o polling por novos trabalhos, priorizando a submissão de novas solicitações. Dependendo do hardware e da configuração, as solicitações são divididas e balanceadas em termos de carga entre os Domains disponíveis. A primeira gravação de uma solicitação composta é enviada imediatamente para a fila de envio do NIC. Uma vez esgotadas as novas solicitações, o trabalhador passa a progredir nas operações pendentes, publicando gravações para preencher o pipeline de hardware. Subsequentemente, as filas de conclusão são consultadas para verificar transferências concluídas e incrementos de contadores imediatos. Os eventos são agregados para fornecer notificações por transferência, repassando a transferência para um tópico de callback dedicado compartilhado por todos os grupos.
A divisão (sharding) dentro de um DomainGroup é flexível. As transferências podem ser direcionadas a NICs específicos por índice. Uma única gravação pode ser dividida. Transferências paginadas, dispersão (scatter) e operações de barreira, que se traduzem em múltiplas gravações, podem ser divididas entre todos os NICs.
Otimizações Específicas de Hardware
Os Domains dentro do TransferEngine são especializados e otimizados para o hardware sob seu controle:
AWS EFA: A implementação do EFA usa libfabric, gerenciando um domínio de fabric por NIC dentro do DomainGroup. Decisões de implementação e otimizações importantes incluem:
Gravações de tamanho zero: O EFA diverge da especificação RDMA ao exigir descritores de destino válidos mesmo para gravações de tamanho zero. Portanto, o TransferEngine impõe descritores válidos para todas as transferências.
Criação de modelos de solicitações de trabalho (Work request templating): Pré-popula e retém os campos comuns das estruturas do libfabric antes da publicação, reduzindo a sobrecarga por operação para transferências em massa e operações de grupos de pares
NVIDIA ConnectX-7: A implementação do ConnectX-7 usa libibverbs com transporte RC. Decisões de implementação e otimizações importantes incluem:
Pares de filas UD para handshakes: Cada Domain usa um par de filas de Datagrama Não Confiável (UD) para trocar handshakes de conexão RC.
Pares de filas RC separados: Cria 2 pares de filas RC por par, um para operações Send/Recv bilaterais e outro para operações Write e WriteImm unilaterais. Essa separação é necessária porque as conclusões de Recv e WriteImm consomem solicitações de trabalho na ordem de postagem, permitindo semântica de Recv de alto nível enquanto suporta WriteImm sem interferência.
Criação de modelos de solicitações de trabalho: Pré-popula campos comuns de estruturas de solicitação de trabalho para transferências em massa e grupos de pares.
Encadeamento de solicitações de trabalho: Encadeia até 4 solicitações de trabalho por meio do ponteiro next de ibv_send_wr, reduzindo o número de disparos de campainha (doorbell rings) para o NIC.
Ordenação relaxada: Habilita IBV_ACCESS_RELAXED_ORDERING para permitir transações PCIe fora de ordem entre o NIC e a memória da GPU.
Conclusão
Ao direcionar a interseção de garantias em diferentes transportes RDMA e introduzir sincronização independente de ordem por meio do ImmCounter, o TransferEngine fornece uma base portátil para comunicação ponto a ponto em sistemas de LLM. A biblioteca é de código aberto e está disponível no GitHub, e encorajamos a comunidade a explorá-la e utilizá-la como base para padrões de comunicação emergentes em sistemas de aprendizado de máquina em larga escala.
Citação
@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}
}