Habilitando modelos de trilhões de parâmetros no AWS EFA

Na Perplexity, utilizamos os melhores modelos para o nosso produto, as nossas APIs e as nossas equipes de pesquisa. Grandes modelos de Mistura de Especialistas (Mixture-of-Experts) de código aberto, como o Kimi-K2, apresentam desafios específicos, uma vez que os maiores nós de inferência com 8 GPUs NVIDIA H200 não conseguem acomodá-los de forma eficiente, necessitando de implementações em vários nós. Apresentamos um conjunto de kernels para paralelismo de especialistas que atingem latências de ponta no ConnectX-7, superando o desempenho do DeepEP. Os mesmos kernels também são os primeiros a atingir latências viáveis no AWS Elastic Fabric Adapter (EFA), permitindo implementações de modelos de trilhões de parâmetros.

Experimente os nossos kernels no GitHub e leia o artigo de pesquisa completo no arXiv.

Introdução

A Mistura de Especialistas (MoE) surgiu como a arquitetura padrão para escalar para centenas de bilhões ou trilhões de parâmetros, oferecendo simultaneamente uma latência razoável no momento da inferência. Os modelos MoE substituem a camada densa de um transformador por um conjunto de especialistas e uma pequena camada de roteamento linear que decide com quais especialistas multiplicar um token. Ao agrupar os tokens associados ao mesmo especialista, o problema torna-se esparso e trivialmente paralelizado, uma vez que diferentes especialistas podem agora ser hospedados em diferentes GPUs distribuídas por vários nós.

Ao contrário de outras formas de paralelismo, como o Paralelismo de Tensores (TP) e o Paralelismo de Dados (DP), que podem ser facilmente implementados usando bibliotecas de comunicação coletiva como torch.distributed e NCCL, o roteamento MoE envolve comunicação esparsa ponto a ponto. Um kernel de despacho deve dividir o conjunto de tokens recebidos e despachá-los para os ranks que hospedam os especialistas, enquanto o seu problema dual, o kernel de combinação, deve receber as suas contrapartes processadas e calcular a sua média ponderada. Estas transferências não mapeiam para coletivos e requerem kernels personalizados para atingir baixa latência.

Na Perplexity, temos desenvolvido kernels inter-nó portáteis e kernels intra-nó especializados, apresentando o nosso trabalho com maior profundidade num artigo anterior sobre comunicação MoE. Mostramos também noutra publicação no blog que a implementação em vários nós pode trazer menor latência e maior rendimento para grandes modelos MoE.

Neste artigo, aprofundamos a implementação dos nossos novos kernels inter-nó, que melhoram significativamente o desempenho do nosso trabalho anterior baseado em NVSHMEM. Além de mostrar resultados de ponta no ConnectX-7, demonstramos também que, dentro do nosso tempo de execução de inferência, os kernels podem ser usados para servir eficazmente os maiores modelos de código aberto.

Inferência sobre InfiniBand

O paralelismo de tensores e o paralelismo de especialistas podem escalar eficazmente dentro de um único nó, uma vez que o NVLink oferece alto rendimento (900 GB/s) com latência na ordem dos microssegundos. Embora os sistemas NVIDIA GB200 e GB300 estejam equipados com memória suficiente para armazenar os pesos e caches KV de modelos de trilhões de parâmetros, as implementações na instância AWS p5en com até 8 GPUs H200 estão limitadas a 1120 GB de HBM, que deve ser dividida entre pesos do modelo e caches KV. Consequentemente, certos modelos necessitam de implementações inter-nó sobre InfiniBand.

Embora o InfiniBand seja ainda extremamente rápido, oferece apenas 400 Gbps (50 GB/s) de rendimento, acrescentando dezenas a centenas de microssegundos de latência para o roteamento MoE. Felizmente, as latências adicionadas podem ser ocultadas através da utilização de especialistas compartilhados, micro-batching e sobreposição de computação-comunicação para fornecer um rendimento viável a um custo competitivo. A DeepSeek foi pioneira nestas técnicas através do DeepSeek-R1, implementado em clusters de GPUs H800 e adaptadores NVIDIA ConnectX-7 InfiniBand.

Nos clusters AWS, a situação é diferente devido ao uso dos seus Elastic Fabric Adapters (EFA) personalizados. Embora estes forneçam um rendimento de pico de 400 Gbps nos coletivos normalmente usados em cargas de trabalho de treinamento, eles ficam ligeiramente aquém do ConnectX-7 nos tamanhos de mensagem trocados durante o despacho e combinação MoE. Além disso, o EFA não suporta GPUDirect Async, exigindo um tópico (thread) de proxy da CPU que faça a ponte entre a GPU e a NIC para iniciar as transferências. Além da complexidade acrescida, as transações PCIe adicionais adicionam microssegundos de sobrecarga a todas as transferências. Frameworks como o NVSHMEM permitem que os escritores de kernel forneçam implementações independentes de dispositivo, contudo, as APIs genéricas incorrem em sobrecargas que tornam as implementações ineficazes em termos de custo e latência.

pplx-kernels baseados em NVSHMEM

Os nossos kernels anteriores abstraíram a implementação de transporte subjacente através do NVSHMEM. Eles baseavam-se no nvshmemx_putmem_signal_nbi_warp para transferir tokens um a um, usando contadores atômicos para sincronizar os pares. Embora isto oferecesse um desempenho razoável em NICs ConnectX-7 com IBGDA, a implementação IBRC do ConnectX-7 baseada em proxy era significativamente mais lenta, enquanto a implementação EFA falhava em entregar latências abaixo de um milissegundo.

Usamos os kernels como uma bancada de testes para o hardware subjacente, enquanto também definimos a interface mais adequada para expor aos nossos modelos. Levantamos a hipótese de que as transferências e interações frequentes entre os kernels e as NICs ou os tópicos (threads) de proxy prejudicavam o desempenho. Além disso, a nossa implementação não dependia de nenhuma garantia do hardware além da confiabilidade, sugerindo que poderíamos acelerar as coisas reduzindo a sobrecarga de primitivas de sincronização.

Seguindo estas observações, concebemos os nossos kernels baseados em EFA expondo a mesma interface, otimizando as interações kernel do dispositivo-proxy do host internamente para obter o desempenho máximo.

MoE all-to-all sobre EFA

Expandir

Os novos kernels seguem uma arquitetura híbrida CPU-GPU, onde os kernels da GPU interagem com o modelo no dispositivo e um tópico (thread) de proxy em execução na CPU do host lida com as interações com a NIC. Os novos kernels reutilizam o mesmo TransferEngine que construímos para transferências de cache KV.

O roteamento MoE é realizado por um kernel de despacho, que recebe os tokens e os especialistas por token para roteá-los a partir do modelo, enviando-os através da rede para os ranks que hospedam os especialistas. Após os kernels GEMM agrupados multiplicarem os tokens com os pesos dos especialistas, o kernel de combinação transfere os tokens para os seus ranks de origem e calcula a sua média ponderada antes de devolvê-los ao modelo.

Tanto o kernel de despacho quanto o de combinação são divididos em uma parte de envio e uma de recepção para facilitar a sobreposição e o micro-batching. O emissor baralha os tokens em buffers de envio para que possam ser escritos nos pares através de uma única operação de escrita RDMA. Após executar o emissor, a GPU pode realizar outro trabalho (como a multiplicação de especialistas compartilhados) enquanto as transferências pela rede estão pendentes em segundo plano. Os kernels receptores bloqueiam até que todas as transferências tenham sido recebidas. Em segundo plano, um tópico (thread) de trabalho do lado do host lida com as transferências: ele é sinalizado pela GPU quando os buffers de envio foram preenchidos, após o que ele despacha as escritas RDMA. Assim que a NIC sinaliza a conclusão, ele desbloqueia o receptor, que retorna os tokens ao modelo.

Os kernels estão fortemente acoplados ao host, comunicando entre si através de memória unificada e GDRCopy. A memória unificada CUDA permite que a memória do dispositivo ou a memória do host sejam mapeadas nos espaços de endereçamento de ambos, manipulando transparentemente a comunicação via PCIe.

Embora a memória unificada seja adequada para transferências em massa, o polling de baixa latência é feito através do GDRCopy. O GDRCopy utiliza GPUDirect RDMA para ler e escrever na memória da GPU com latência de microssegundos.

Além do RDMA, dentro de um nó, os kernels também aproveitam o NVLink para trocar tokens. Isso pode descarregar uma fração significativa do tráfego, até um quarto para EP32 e até um oitavo para EP64, alcançando latências gerais menores.

Buffers de Envio e Recebimento

O design dos buffers de envio para o kernel de despacho é crucial para alcançar latências viáveis envolvendo o proxy do host. Se houver N ranks, cada um despachando T tokens, cada um roteado para R especialistas de um total de E, um rank pode receber até N * T * max(R, E / N) tokens no pior caso possível, se todos forem roteados para ele. Além disso, este limite mínimo só pode ser alcançado quando os emissores dispõem as suas escritas de forma contígua, uma vez que buffers de recepção privados consumiriam demasiada memória. Para que os emissores possam dispor as suas escritas consecutivamente sem conflitos, o conhecimento total das informações de roteamento deve ser distribuído para cada rank. Consequentemente, os nossos kernels de despacho trocam primeiro as contagens de tokens por especialista, após o que cada um pode determinar para onde vão as suas escritas e de qual offset recebem tokens de cada rank.

Simplesmente trocar informações de roteamento resultaria em latências desnecessariamente altas, uma vez que cada rank teria que esperar que as escritas chegassem antes de despachar qualquer escrita sua. Em vez disso, reservamos uma pequena quantidade de Espaço privado por emissor em cada rank receptor e despachamos alguns dos tokens juntamente com as escritas das informações de roteamento. Isto garante a utilização total da largura de banda durante a maior parte do processo de despacho. Após as rotas serem trocadas, os tokens restantes são transferidos através de uma única escrita por par remoto. A troca de informações de roteamento e os buffers de tamanho fixo, juntamente com as transferências que os preenchem, são ilustrados na figura abaixo.

Expandir

Os kernels de combinação são substancialmente mais simples, uma vez que as informações de roteamento já estão disponíveis para eles a partir da fase de despacho. Eles podem copiar os tokens através de uma única escrita por par, após o que leem a partir do buffer de recebimento para calcular a média ponderada dos tokens nos tensores de saída.

Para o NVLink, os buffers de envio e recebimento dentro de um nó são mapeados nos pares para que eles possam empurrar ou puxar tokens diretamente dos ranks vizinhos.

Despacho e Combinação

As figuras abaixo detalham o tempo e as interações entre a GPU, CPU e NIC, mostrando uma iteração de despacho e combinação intercalada com os kernels do modelo.

O trabalho começa no envio de despacho, com o kernel agregando as informações de roteamento e calculando as contagens de tokens por especialista através de acumulação atômica na memória compartilhada dentro de um único bloco. Assim que as contagens de tokens estão disponíveis, elas são copiadas para a memória do host e o proxy é notificado imediatamente para começar a despachar as informações de roteamento para todos os outros ranks pares. Em paralelo, o despacho continua a calcular a soma cumulativa das contagens de tokens, identificando os offsets necessários para empacotar tokens em buffers de envio contíguos por par. Assim que os offsets estão disponíveis, todos os tokens a serem enviados via EFA são empacotados nos buffers de envio. O proxy é notificado para publicar as escritas nos buffers de envio de pares inter-nó. Intra-nó, os kernels também empurram tokens imediatamente para buffers privados dentro do mesmo nó através do NVLink.

Expandir

Expandir

Assim que um rank recebe informações de roteamento de todos os pares, ele centraliza-as e calcula vários offsets necessários para identificar tokens para despacho e combinação. Este trabalho é feito na CPU enquanto as transferências já estão pendentes, portanto, não é crítico para a latência. Subsequentemente, assim que todos os offsets são conhecidos, as escritas para o restante dos tokens, que excedem o tamanho do buffer privado, são publicadas e a fase de envio do despacho está completa.

No lado do receptor, o kernel primeiro puxa tokens do mesmo nó que estão além da capacidade dos buffers privados através do NVLink. Em seguida, ele bloqueia até que o proxy do host sinalize a conclusão de todas as transferências. Assim que o sinal é recebido, os tokens são baralhados dos buffers de recebimento para o tensor de saída passado para o grupo GEMM. Uma vez que os buffers são reutilizados entre o despacho e a combinação, tanto as barreiras NVLink quanto RDMA garantem que os buffers do receptor foram lidos antes de permitir que o próximo kernel os sobrescreva.

Os kernels de combinação seguem um padrão aproximadamente semelhante, baralhando tokens, sinalizando o proxy e aguardando a conclusão. No entanto, os metadados para publicar as suas escritas estão prontamente disponíveis e todos os offsets remotos são conhecidos, portanto, todos os tokens transferidos intra-nó são empurrados via NVLink para os receptores e os tokens inter-nó são transferidos via uma única escrita. Antes de retornar os resultados, o kernel acumula os tokens dos buffers do receptor, escrevendo-os nos tensores de saída.

TransferEngine para roteamento MoE

Para prototipar a nossa implementação, reutilizamos inicialmente o Transfer Engine construído para prefill e decodificação desagregada a fim de emitir e rastrear escritas RDMA. No entanto, ao contrário da decodificação desagregada onde cada nó publica um grande número de escritas para o mesmo par, o despacho MoE publica pequenas transferências para todos os pares numa única operação. O desempenho imediato não foi ideal e a nossa implementação estava atrasada em relação ao DeepEP no ConnectX-7 em cerca de 20us. Consequentemente, especializamos o TransferEngine para suportar duas operações, scatter e barreira, para otimizar as interações com um grande número de pares.

Tanto as operações scatter quanto de barreira expõem grupos de pares que podem ser pré-registrados para pré-preencher as estruturas de dados usadas pela implementação de transporte subjacente (templating de Requisição de Trabalho (WR)). O scatter copia diferentes fatias do mesmo buffer de origem para diferentes offsets em todos os membros de um grupo de pares. A barreira sinaliza todos os pares através de um imediato. Colocar e manipular estas operações em massa economiza microssegundos valiosos no caminho crítico. No EFA, onde duas NICs de 200 Gbps agregam para fornecer 400 Gbps de largura de banda, fragmentamos os pares do grupo pelas duas NICs em vez de fragmentar os bytes de cada transferência.

Embora inicialmente tenhamos desenvolvido o TransferEngine apenas no EFA usando libfabric, os kernels MoE funcionam no ConnectX-7 sem alterar uma única linha de código após adicionar suporte ao ConnectX-7 no TransferEngine usando libibverbs. Comparado ao EFA baseado em Scalable Reliable Datagram (SRD), o ConnectX-7 requer trabalho adicional para a configuração da conexão e gerenciamento de pares. Usamos um par de filas Unreliable Datagram (UD) para inicializar a Reliable Connection (RC). Criamos dois pares de filas para cada par remoto, um para SEND/RECV bidirecional e um para WRITE_IMM unidirecional, porque tanto o RECV quanto o WRITE_IMM consomem a operação RECV em ordem enquanto as duas operações requerem tamanhos de buffer diferentes. Além disso, relaxamos a ordenação de transações PCIe entre o ConnectX-7 e a memória da GPU. Semelhante ao EFA, implementamos a otimização de templating WR. Além disso, ao enviar várias mensagens pequenas para o mesmo par, usamos o encadeamento WR para melhorar a taxa de mensagens.

Uma vez que o trabalho do lado da GPU é idêntico em ambos os adaptadores, levamos a implementação de transporte subjacente aos seus limites para superar o desempenho do DeepEP, a fim de destacar que a lacuna entre o roteamento MoE EFA e ConnectX-7 é determinada pelos recursos do hardware subjacente.

Avaliação

Latência de Despacho/Combinação

Avaliamos as latências de despacho e combinação em sistemas H200 equipados com adaptadores ConnectX-7 e EFA na configuração padrão DeepSeek-V3 e Kimi-K2 com uma dimensão oculta de 7168, despacho fp8 escalonado por bloco e combinação bf16. Consideramos um tamanho de batch de 128 para despacho e 4096 para combinação. No ConnectX-7, comparamos com o DeepEP e os nossos pplx-kernels anteriores baseados em NVSHMEM, usando tanto GPUDirect Async (IBGDA) quanto um proxy de host (IBRC). Atingimos latências combinadas de 459us, 582us e 692us em EP16, EP32 e EP64, respectivamente. Superamos outros kernels baseados em EFA que dependem de proxies genéricos, como o UCCL-EP, que atingem 519us, 966us e 1159us.

Expandir

Na decodificação, apesar dos nossos kernels serem alguns microssegundos mais lentos que o DeepEP, alcançamos latências globais menores, pois os nossos kernels de combinação são substancialmente mais rápidos. Embora usemos um proxy de host, devido às nossas transferências em massa superamos o desempenho de todas as implementações sem proxy baseadas em IBGDA. Mostramos também que a implementação do proxy importa, uma vez que o nosso proxy altamente especializado é uma ordem de magnitude mais rápido que a implementação NVSHMEM genérica. Apesar do fato de o EFA ser duas vezes mais lento que o ConnectX-7 ao transferir pacotes de 256 KB, que estão próximos do tamanho típico de pacote trocado pelo roteamento MoE, a lacuna entre os dois adaptadores é, no geral, menos substancial.

Expandir

Mesmo que os nossos kernels otimizados para decodificação possam escalar para tamanhos de batch de prefill sem quaisquer ajustes adicionais, o desempenho está atrasado em relação ao DeepEP. Isso deve-se principalmente ao fato de acumularmos tokens apenas no lado do receptor da combinação, enquanto o DeepEP acumula no emissor para reduzir a quantidade de dados transferidos. No entanto, esta otimização pode ter implicações na precisão. Além disso, como os kernels transferem dados numa única escrita, utilizamos mais memória para buffers de envio e recebimento.

Rendimento do Kimi-K2 / DeepSeek-V3

Para provar a viabilidade dos nossos kernels, fazemos o benchmark do desempenho ponta a ponta do nosso motor de inferência interno que serve o DeepSeek-V3 (671B parâmetros) e o Kimi-K2 (1T parâmetros) em múltiplas instâncias p5en equipadas com Elastic Fabric Adapters (EFA) e GPUs H200. Executamos o modelo em modo de verificação de rascunho e relatamos a latência assumindo uma taxa de aceitação de 80% usando uma única camada MTP.

No DeepSeek-V3, comparamos com uma linha de base intra-nó EP=8 executando kernels de despacho e combinação eficientes apenas via NVLink. As implementações inter-nó no geral igualam ou excedem o desempenho da implementação de nó único, particularmente em batches de tamanho médio, demonstrando que escalar entre nós traz melhorias ao rendimento global.

No Kimi-K2, os kernels permitem o atendimento. Devido ao seu tamanho, o modelo não cabe dentro do mesmo nó usando Paralelismo de Dados (DP) num nó H200, portanto, uma implementação inter-nó é necessária para servir o modelo. Como o modelo tem menos cabeças de atenção, atingimos latências melhores que o DeepSeek-V3 e podemos servir o modelo no EFA com uma latência viável.

Expandir

Expandir

Também consideramos o micro-batching, embora ele apenas tenha melhorado as latências em tamanhos de batch muito elevados ou ao usar mais do que um token de rascunho. Para beneficiar da divisão de problemas GEMM agrupados em dois, os tamanhos de batch de tokens devem ser suficientemente densos. A maior parte do tempo ainda é gasta no despacho e combinação MoE, indicando que quaisquer outras melhorias nos kernels melhorarão consequentemente as latências e facilitarão o micro-batching.

Trabalho Futuro

Estamos trabalhando em estreita colaboração com engenheiros da AWS para melhorar o desempenho no EFA. Da nossa parte, seguiremos as atualizações do libfabric, o que reduzirá a sobrecarga do plano de dados. Além disso, experimentaremos o efa-direct para reduzir ainda mais as sobrecargas do espaço do usuário e interagir diretamente com o hardware.

Referências

Relatório Técnico do DeepSeek R1

Relatório Técnico do Kimi K2

DeepEP

UCCL-EP

Melhorando o Desempenho de Rede de Sistemas HPC Usando NVIDIA Magnum IO NVSHMEM e GPUDirect Async

Citação

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