TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference

1. Declaração do Problema

A inferência distribuída de Modelos de Linguagem de Grande Escala (LLMs) usando Paralelismo de Tensores (TP) enfrenta gargalos de desempenho significativos devido a sobrecargas de comunicação, mesmo em interconexões de alta velocidade como NVLink.

• Sobrecarga de Comunicação: Em implantações de TP, a comunicação (especificamente AllReduce) representa 9–23% da latência de inferência de ponta a ponta.
• O Desafio do "Lote Pequeno": Técnicas existentes para sobrepor comunicação com computação (por exemplo, decompor o trabalho em azulejos ou tokens menores) não estão habilitadas por padrão nos principais sistemas de serviço (vLLM, SGLang, TensorRT-LLM). Isso ocorre porque a decomposição de cargas de trabalho pequenas (típicas em inferência de baixa latência, por exemplo, <2K tokens) frequentemente introduz ineficiências de computação (efeitos de quantização de ondas) que superam os ganhos da sobreposição de comunicação.
• Contenção de Recursos: Operações de comunicação frequentemente consomem um grande número de Multiprocessadores de Fluxo (SMs) (16–20+), deixando os kernels de computação sem recursos.
• Redundância do RMSNorm: A operação RMSNorm, frequentemente negligenciada, adiciona 4–9% de sobrecarga. No TP padrão, ela é realizada de forma redundante em cada GPU após o AllReduce, apesar de todas as GPUs possuírem dados idênticos naquele ponto.

2. Metodologia: TokenWeave

O TokenWeave é um sistema projetado para habilitar uma sobreposição eficiente de computação-comunicação para inferência com paralelismo de tensores, mesmo para lotes pequenos (até 1024 tokens). Ele emprega três inovações principais:

A. Divisão Inteligente de Tokens Consciente de Ondas

Em vez de dividir ingenuamente um lote em duas metades iguais, o TokenWeave utiliza uma estratégia de divisão consciente de ondas.

• Mecanismo: Ele particiona o lote de entrada em dois subconjuntos (prefixo e sufixo) de modo que o número total de "ondas" de execução (grupos de CTAs) necessárias para ambos os subconjuntos não exceda as ondas necessárias para o lote original não dividido.
• Objetivo: Isso minimiza as sobrecargas de quantização de ondas, garantindo que a divisão do trabalho não degrade a eficiência de computação.
• Execução: O sistema usa dois streams CUDA (Computação e Comunicação). Enquanto o AllReduce (fundido com RMSNorm) é executado na primeira divisão, o cálculo de atenção é executado na segunda divisão, e vice-versa.

B. Reordenação e Fusão do RMSNorm

O TokenWeave identifica o RMSNorm como um alvo crítico de otimização e reordena o fluxo de execução.

• Reordenação: Em vez de realizar AllReduce seguido de RMSNorm, o TokenWeave decompõe o AllReduce em ReduceScatter e AllGather.
  1. ReduceScatter: Cada GPU recebe seu fragmento do tensor.
  2. RMSNorm Local: Cada GPU realiza o RMSNorm em seu fragmento local antes da agregação final. Como os dados já estão particionados, isso elimina o cálculo redundante entre as GPUs.
  3. AllGather: Os fragmentos normalizados são reunidos.
• Kernel Fundido: Um novo kernel fundido AllReduce–RMSNorm é implementado. Ele funde ReduceScatter, RMSNorm e AllGather em uma única operação.
  • Otimização de Memória: Ao calcular a variância diretamente no resultado do ReduceScatter nos registradores SM, elimina-se uma leitura de HBM. Ao escrever valores normalizados diretamente nos endereços Multimem para AllGather, elimina-se uma escrita de HBM.
  • Eficiência de SM: Este kernel aproveita o NVSHARP/Multimem (recursos de hardware nas GPUs Hopper/Blackwell) para realizar comunicação e normalização usando apenas 2–8 SMs, deixando a maioria dos SMs livres para computação.

C. Estratégia de Sobreposição Adaptativa

O TokenWeave decide dinamicamente se deve usar divisão/sobreposição completa ou apenas o kernel fundido com base no tamanho do lote:

• Lotes Grandes: Habilita a divisão e sobreposição completa em duas vias.
• Lotes Pequenos: Desabilita a divisão (para evitar sobrecarga), mas ainda aplica o kernel fundido AllReduce–RMSNorm para reduzir o tráfego de memória e eliminar o cálculo redundante.

3. Contribuições Principais

1. Otimização do RMSNorm: Primeiro a identificar e otimizar o RMSNorm na inferência de TP via reordenação e um kernel fundido personalizado, reduzindo o tráfego de HBM e a computação redundante.
2. Divisão Consciente de Ondas: Um algoritmo de divisão de tokens inovador que previne a perda de eficiência de computação devido à quantização de ondas, habilitando a sobreposição mesmo para lotes pequenos de tokens (1K–2K).
3. Eficiência Consciente de Hardware: Aproveita o NVSHARP/Multimem para reduzir o uso de SMs de comunicação de ~20 para 2–8, prevenindo a escassez de recursos.
4. Integração de Sistema: O primeiro sistema a habilitar com sucesso a sobreposição de computação-comunicação para serviço com paralelismo de tensores em configurações de baixa latência, integrado ao framework vLLM.

4. Resultados Experimentais

As avaliações foram conduzidas em sistemas 8×H100 DGX usando modelos como Llama-3.3-70B, Qwen2.5-72B e Mixtral-8x22B.

• Melhoria de Latência:
  • Alcança até 1,28× de aceleração na latência em comparação com uma base otimizada (vLLM com Multimem).
  • Ganhos significativos mesmo em comprimentos de sequência curtos (por exemplo, 1,2× em 1K tokens).
  • Supera soluções de sobreposição de última geração como TileLink, que mostram ganhos apenas em lotes maiores (≥4K) e incorrem em sobrecargas em tamanhos menores.
• Melhoria de Throughput:
  • Entrega até 1,19× de throughput mais alto em rastreamentos do mundo real (ShareGPT, arXiv).
  • Em algumas configurações, o TokenWeave supera uma base teórica com zero comunicação (vLLM-nocomm) porque seu kernel fundido otimiza a operação limitada por memória RMSNorm melhor do que a base lida com a computação pura.
• Eficiência do Kernel:
  • O kernel fundido AllReduce–RMSNorm fornece 1,34–1,39× de aceleração sobre a execução sequencial em todos os tamanhos de tokens.
  • Alcança desempenho quase ideal usando apenas 4–8 SMs.

5. Significado

• Habilita Inferência Distribuída de Baixa Latência: Ao resolver o problema do "lote pequeno", o TokenWeave torna a inferência distribuída de TP viável para aplicações interativas onde a latência é crítica, sem exigir lotes massivos.
• Utilização de Hardware: Maximiza a utilidade de recursos modernos de GPU (NVSHARP, Multimem), deslocando o gargalo da sobrecarga de comunicação para a computação pura.
• Impacto Prático: O código-fonte é de código aberto e integrado ao vLLM, oferecendo benefícios de desempenho imediatos para sistemas de serviço de LLM em produção sem exigir alterações na arquitetura do modelo.
• Mudança de Paradigma: Desafia a suposição de que a sobreposição de comunicação é benéfica apenas para lotes de treinamento grandes ou configurações de paralelismo de especialistas (MoE), provando que é essencial para a inferência padrão com paralelismo de tensores também.