Habilitando modelos de um trilhão de parâmetros no AWS EFA

Na Perplexity, usamos os melhores modelos para o nosso produto, nossas APIs e nossas equipes de pesquisa. Grandes modelos Mixture-of-Experts de código aberto, como o Kimi-K2, apresentam desafios específicos, já que os maiores nós de inferência com GPUs 8x NVIDIA H200 não conseguem acomodá-los de forma eficiente, exigindo implantações em vários nós. Apresentamos um conjunto de kernels para paralelismo de especialistas que alcançam latências de última geração no ConnectX-7, superando o desempenho do DeepEP. Os mesmos kernels também são os primeiros a alcançar latências viáveis no AWS Elastic Fabric Adapter (EFA), viabilizando implantações de modelos de um trilhão de parâmetros.

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

Introdução

O Mixture-of-Experts (MoE) emergiu como a arquitetura padrão para escalar para centenas de bilhões ou trilhões de parâmetros, oferecendo também 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 qual especialista multiplicar um token. Ao agrupar os tokens associados ao mesmo especialista, o problema torna-se esparso e trivialmente paralelizável, uma vez que diferentes especialistas agora podem ser hospedados em GPUs diferentes espalhadas por vários nós.

Ao contrário de outras formas de paralelismo, como o Parallelismo de Tensor (TP) e o Parallelismo 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 de entrada e despachá-los para as classificações que hospedam os especialistas, enquanto seu problema dual, o kernel de combinação, deve receber seus equivalentes processados e computar sua média ponderada. Essas transferências não são mapeadas para coletivos e exigem kernels personalizados para atingir baixa latência.

Na Perplexity, temos desenvolvido tanto kernels portáteis entre nós quanto kernels especializados intra-nó, apresentando nosso trabalho com mais profundidade em um artigo anterior sobre comunicação MoE. Também mostramos em outra postagem de blog que a implantação em vários nós pode trazer menor latência e maior rendimento para grandes modelos MoE.

Neste artigo, mergulhamos profundamente na implementação de nossos novos kernels entre nós, o que melhora significativamente o desempenho do nosso trabalho anterior baseado em NVSHMEM. Além de exibir resultados de última geração no ConnectX-7, também demonstramos que, em nosso tempo de execução de inferência, os kernels podem ser usados para atender efetivamente aos maiores modelos de código aberto.

Inferência sobre InfiniBand

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

Embora o InfiniBand ainda seja extremamente rápido, ele oferece apenas 400 Gbps (50 GB/s) de rendimento, adicionando dezenas a centenas de microssegundos de latência para o roteamento MoE. Felizmente, as latências adicionadas podem ser ocultadas por meio do uso de especialistas compartilhados, micro-lotes e sobreposição de computação e comunicação para entregar rendimento viável a um custo competitivo. O DeepSeek foi pioneiro nessas técnicas por meio do DeepSeek-R1, implantado em clusters de GPUs H800 e adaptadores NVIDIA ConnectX-7 InfiniBand.

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

pplx-kernels baseados em NVSHMEM

Nossos kernels anteriores abstraíram a implementação de transporte subjacente por meio do NVSHMEM. Eles confiavam em nvshmemx_putmem_signal_nbi_warp para transferir tokens um por um, usando contadores atômicos para sincronizar pares. Embora isso tenha oferecido um desempenho razoável em NICs ConnectX-7 com IBGDA, a implementação IBRC do ConnectX-7 baseada em proxy foi significativamente mais lenta, enquanto a implementação do EFA falhou em entregar latências abaixo de um milissegundo.

Usamos os kernels como uma bancada de testes para o hardware subjacente, ao mesmo tempo em que determinamos a interface mais adequada para expor aos nossos modelos. Formulamos a hipótese de que as transferências frequentes e as interações entre os kernels e as NICs ou os tópicos de proxy prejudicavam o desempenho. Além disso, 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 das primitivas de sincronização.

Seguindo essas observações, criamos nossos kernels baseados em EFA expondo a mesma interface, otimizando as interações entre o kernel do dispositivo e o proxy do host nos bastidores para obter o pico de desempenho.

MoE all-to-all sobre EFA

Expandir

Os novos kernels seguem uma arquitetura híbrida de CPU-GPU, onde os kernels de GPU interagem com o modelo no dispositivo e um tópico de proxy executado 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 é executado por um kernel de despacho, que recebe os tokens e os especialistas por token para roteá-los a partir do modelo, enviando-os pela rede para as classificações 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 suas classificações de origem e calcula sua média ponderada antes de retorná-los ao modelo.

Tanto o kernel de despacho quanto o de combinação são divididos em uma metade de envio e uma de recebimento para facilitar a sobreposição e o micro-loteamento. O remetente embaralha os tokens em buffers de envio para que possam ser gravados em pares por meio de uma única operação de gravação RDMA. Após a execução do remetente, a GPU pode realizar outro trabalho (como 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 de trabalhador do lado do host lida com as transferências: ele recebe um sinal da GPU quando os buffers de envio foram preenchidos, após o que despacha as gravações RDMA. Assim que a NIC sinaliza a conclusão, ele desbloqueia o receptor, que retorna os tokens ao modelo.

Os kernels são fortemente acoplados ao host, comunicando-se uns com os outros 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 Espaço de endereçamento de ambos, lidando de forma transparente com a comunicação via PCIe.

Embora a memória unificada seja adequada para transferências em massa, a sondagem de baixa latência é feita por meio do GDRCopy. O GDRCopy usa o GPUDirect RDMA para ler e gravar 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 mais baixas.

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 classificações, cada uma despachando T tokens, cada um roteado para R especialistas de um total de E, uma classificação pode receber até N * T * max(R, E / N) tokens no pior caso possível, se todos forem roteados para ela. Além disso, esse limite mínimo pode ser alcançado apenas quando os remetentes organizam suas gravações de forma contígua, já que buffers de recebimento privados usariam muita memória. Para que os remetentes possam organizar suas gravações consecutivamente sem conflitos, o conhecimento completo das informações de roteamento deve ser distribuído para cada classificação. Consequentemente, nossos kernels de despacho primeiro trocam contagens de tokens por especialista, após o que cada um pode determinar para onde vão suas gravações e de qual classificação recebem tokens e a partir de qual deslocamento.

Simplesmente trocar informações de roteamento resultaria em latências desnecessariamente altas, pois cada classificação teria que esperar as gravações chegarem antes de despachar qualquer gravação própria. Em vez disso, reservamos uma pequena quantidade de Espaço privado por remetente em cada classificação receptora e despachamos alguns dos tokens junto com as gravações das informações de roteamento. Isso garante a utilização total da largura de banda durante a maior parte do processo de despacho. Após a troca das rotas, os tokens restantes são transferidos por meio de uma única gravação 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, pois as informações de roteamento já estão disponíveis para eles a partir do estágio de despacho. Eles podem copiar os tokens em uma única gravação por par, após o que leem do buffer de recebimento para computar 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 para os pares para que eles possam enviar ou puxar tokens diretamente das classificações vizinhas.

Despacho e Combinação

As figuras abaixo detalham o tempo e as interações entre a GPU, a CPU e a 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 computando contagens de tokens por especialista por meio de acúmulo atômico em 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 todas as outras classificações de pares. Em paralelo, o despacho continua a computar a soma cumulativa das contagens de tokens, identificando os deslocamentos necessários para empacotar tokens em buffers de envio contíguos por par. Assim que os deslocamentos 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 gravações nos buffers de envio de pares entre nós. Intra-nó, os kernels também empurram imediatamente tokens para buffers privados dentro do mesmo nó através do NVLink.

Expandir

Expandir

Assim que uma classificação recebe informações de roteamento de todos os pares, ela as centraliza e computa vários deslocamentos necessários para identificar tokens para despacho e combinação. Esse 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 deslocamentos são conhecidos, as gravações para o restante dos tokens, que excedem o tamanho do buffer privado, são publicadas e a fase de envio do despacho é concluída.

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

Os kernels de combinação seguem um padrão aproximadamente semelhante, embaralhando tokens, sinalizando o proxy e aguardando a conclusão. No entanto, os metadados para publicar suas gravações estão prontamente disponíveis e todos os deslocamentos remotos são conhecidos, portanto, todos os tokens transferidos intra-nó são enviados via NVLink para os receptores e os tokens entre nós são transferidos via uma única gravação. Antes de retornar os resultados, o kernel acumula os tokens dos buffers receptores, gravando-os nos tensores de saída.

TransferEngine para roteamento MoE

Para criar um protótipo de nossa implementação, inicialmente reutilizamos o Motor de Transferência construído para pré-preenchimento e decodificação desagregados, a fim de emitir e rastrear gravações RDMA. No entanto, ao contrário da decodificação desagregada, onde cada nó publica um grande número de gravações para o mesmo par, o despacho MoE publica pequenas transferências para todos os pares em uma única operação. Pronta para uso, a performance não era ideal e nossa implementação estava ficando atrás do DeepEP no ConnectX-7 em cerca de 20us. Consequentemente, especializamos o TransferEngine para suportar duas operações, dispersão (scatter) e barreira (barrier), para otimizar as interações com um grande número de pares.

Tanto as operações de dispersão quanto as de barreira expõem grupos de pares que podem ser pré-registrados para pré-popular as estruturas de dados usadas pela implementação de transporte subjacente (modelagem de Solicitação de Trabalho (WR)). A dispersão copia diferentes fatias do mesmo buffer de origem em diferentes deslocamentos em todos os membros de um grupo de pares. A barreira sinaliza todos os pares através de um imediato. Enfileirar e manipular essas operações em massa economiza microssegundos valiosos no caminho crítico. No EFA, onde duas NICs de 200 Gbps se agregam para fornecer largura de banda de 400 Gbps, dividimos os pares do grupo entre as duas NICs em vez de dividir 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 Datagrama Confiável Escalável (SRD), o ConnectX-7 requer trabalho adicional para a configuração de conexão e gerenciamento de pares. Usamos um par de filas de Datagramas Não Confiáveis (UD) para inicializar a Conexão Confiável (RC). Criamos dois pares de filas para cada par remoto, um para SEND/RECV bidirecional e um para WRITE_IMM unidirecional, porque tanto RECV quanto WRITE_IMM consomem a operação RECV em ordem, enquanto as duas operações exigem 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. Assim como no EFA, implementamos a otimização de modelagem de WR. Adicionalmente, ao enviar várias mensagens pequenas para o mesmo par, usamos o encadeamento de WR para melhorar a taxa de mensagens.

Como 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 e destacar que a lacuna entre o roteamento MoE do EFA e do ConnectX-7 é determinada pelas capacidades 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 dimensionado por bloco e combinação bf16. Consideramos um tamanho de lote de 128 para despacho e 4096 para combinação. No ConnectX-7, comparamos com o DeepEP e nossos antigos pplx-kernels baseados em NVSHMEM, usando tanto o 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 alcançam 519us, 966us e 1159us.

Expandir

Na decodificação, apesar de nossos kernels serem alguns microssegundos mais lentos que o DeepEP, alcançamos latências gerais menores, pois 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. Também mostramos que a implementação de proxy importa, pois nosso proxy altamente especializado é uma ordem de magnitude mais rápido do que a implementação NVSHMEM genérica. Apesar do fato de que o EFA é duas vezes mais lento que o ConnectX-7 ao transferir pacotes de 256 KB, que estão próximos ao tamanho típico de pacote trocado pelo roteamento MoE, a lacuna entre os dois adaptadores é globalmente menos substancial.

Expandir

Apesar de nossos kernels otimizados para decodificação poderem escalar para tamanhos de lote de pré-preenchimento sem quaisquer ajustes adicionais, o desempenho está ficando atrás do DeepEP. Isso se deve principalmente ao fato de acumularmos tokens apenas no lado do receptor da combinação, enquanto o DeepEP acumula no remetente para reduzir a quantidade de dados transferidos. No entanto, essa otimização pode ter implicações na precisão. Além disso, como os kernels transferem dados em uma única gravação, utilizamos mais memória para os buffers de envio e recebimento.

Rendimento de Kimi-K2 / DeepSeek-V3

Para provar a viabilidade de nossos kernels, realizamos benchmarks do desempenho de ponta a ponta de nosso motor de inferência interno atendendo DeepSeek-V3 (671B de parâmetros) e Kimi-K2 (1T de parâmetros) em várias instâncias p5en equipadas com Elastic Fabric Adapters (EFA) e GPUs H200. Executamos o modelo no 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 eficientes de despacho e combinação apenas com NVLink. As implantações entre nós geralmente igualam ou superam o desempenho da implantação de nó único, particularmente em lotes de tamanho médio, demonstrando que a escala entre nós traz melhorias para o rendimento geral.

No Kimi-K2, os kernels habilitam a atendimento. Devido ao seu tamanho, o modelo não cabe no mesmo nó usando Paralelismo de Dados (DP) em um nó H200, portanto, uma implantação entre nós é necessária para servir o modelo. Como o modelo tem menos cabeças de atenção, alcançamos latências melhores do que o DeepSeek-V3 e podemos servir o modelo no EFA com uma latência viável.

Expandir

Expandir

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

Trabalho Futuro

Estamos trabalhando em estreita colaboração com os engenheiros da AWS para melhorar o desempenho no EFA. Do nosso lado, seguiremos as atualizações do libfabric, que reduzirão 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}
}