Speculating Experts Accelerates Inference for Mixture-of-Experts

O artigo propõe um esquema de pré-busca de especialistas que utiliza representações internas do modelo para prever e buscar antecipadamente os pesos de especialistas em MoE, permitindo a sobreposição de transferência de memória e computação e reduzindo o tempo de inferência em até 14% em ambientes com restrições de memória.

O essencial

Imagine que você tem um cozinheiro genial (o modelo de IA) que precisa preparar pratos complexos. Para ser eficiente, esse cozinheiro não tem uma cozinha gigante cheia de todos os ingredientes do mundo na bancada. Em vez disso, ele tem uma despensa enorme no porão (a memória do computador/CPU) e apenas uma pequena bancada de trabalho na cozinha (a memória rápida do chip/GPU).

O modelo "Mixture-of-Experts" (MoE) funciona assim: ele é composto por muitos "especialistas" (como um especialista em matemática, outro em culinária, outro em história). Mas, para cada frase que você escreve, o cozinheiro só precisa de 3 ou 4 desses especialistas específicos.

O Problema: A Corrida contra o Tempo

No método tradicional de usar esses modelos em computadores comuns (que não têm memória gigante), o cozinheiro precisa:

1. Pensar na frase.
2. Parar tudo.
3. Correr até o porão, pegar os ingredientes (pesos dos especialistas) que ele acha que vai precisar.
4. Voltar para a cozinha e cozinhar.
5. Repetir.

O problema é que ir ao porão é lento. O cozinheiro passa 85% do tempo correndo pelo corredor e apenas 15% cozinhando. Isso deixa a IA muito lenta.

A Solução: O "Palpite" do Cozinheiro (Speculating Experts)

Os autores deste paper tiveram uma ideia brilhante: E se o cozinheiro pudesse adivinhar o que vai precisar antes de parar para pegar?

Eles criaram um sistema de "Palpites Inteligentes" (Speculating Experts):

1. O Palpite: Enquanto o cozinheiro está cozinhando o prato atual (processando a frase atual), ele olha para os ingredientes que já estão na bancada e faz uma previsão: "Pelo jeito que estou cozinhando agora, no próximo passo eu vou precisar do Especialista em Matemática e do Especialista em História".
2. A Corrida Paralela: Assim que ele faz esse palpite, ele manda um ajudante (o sistema de memória) correr ao porão e trazer esses ingredientes enquanto ele continua cozinhando o prato atual.
3. O Resultado: Quando ele termina o prato atual e precisa do próximo, os ingredientes já estão na bancada! Não há tempo perdido correndo. O tempo de "cozinha" e o tempo de "corrida" acontecem ao mesmo tempo.

O Truque Secreto: O "Estado Quase-Oculto"

Como o cozinheiro faz esse palpite com tanta precisão?
Eles descobriram que existe um "cheiro" ou uma "vibe" no prato que está sendo feito (chamado no paper de quasi-hidden state). Ao analisar essa vibe, o sistema consegue prever com muita precisão quais especialistas serão necessários no próximo passo.

• Analogia: É como um maestro que, ao ouvir a nota que o violinista está tocando agora, já sabe qual nota o trompetista vai tocar no compasso seguinte e avisa o trompetista para se preparar.

E se o Palpite estiver errado?

Às vezes, o cozinheiro erra o palpite. Ele acha que precisa de Matemática, mas na verdade precisava de História.

• No método antigo: Se ele erra, ele tem que parar, voltar ao porão, pegar o ingrediente certo e recomeçar. Isso quebra a eficiência.
• Neste novo método: O cozinheiro continua usando o ingrediente que ele trouxe por palpite (Matemática), mesmo que não fosse o ideal.
  • A Surpresa: Eles descobriram que, na maioria das vezes, usar o "palpite errado" não estraga o prato final! A IA continua entendendo o contexto e respondendo bem. Isso é incrível porque significa que eles podem manter a velocidade alta sem precisar parar para corrigir o erro.

Quando o Palpite Falha (O "Estimador Leve")

Em alguns modelos (como o Qwen-30B), os primeiros passos da receita são muito caóticos, e o palpite fica ruim. Para resolver isso, eles treinaram um "Pequeno Ajudante" (um estimador neural leve).

• Esse ajudante é um especialista em prever os primeiros passos difíceis. Ele é tão pequeno e rápido que não atrapalha a cozinha, mas garante que os ingredientes certos cheguem nos momentos críticos.

O Resultado Final

Ao usar essa técnica de "palpites" e "corridas paralelas":

• A IA ficou 5% a 14% mais rápida.
• Isso parece pouco, mas em termos de tempo de espera, é como se você tivesse comprado um novo motor para o seu carro sem gastar dinheiro extra.
• Isso permite que modelos gigantes e inteligentes rodem em computadores domésticos (como seu PC ou notebook) sem travar, pois o gargalo da memória foi reduzido.

Resumo em uma frase:
Em vez de esperar para descobrir o que precisa e depois buscar, o sistema adivinha o que vai precisar e busca enquanto trabalha, transformando uma corrida lenta em uma dança eficiente.

Resumo técnico

Resumo Técnico: Aceleração de Inferência em Modelos MoE via Adivinhação de Especialistas

1. O Problema

Os modelos de Mistura de Especialistas (MoE) tornaram-se a arquitetura predominante para Grandes Modelos de Linguagem (LLMs) de última geração (como Qwen, GPT-OSS, GLM), permitindo escalar a capacidade de parâmetros mantendo a ativação esparsa e reduzindo o custo computacional por token.

No entanto, em cenários de inferência com restrições de memória (comuns em hardware de consumidor ou GPUs únicas), os pesos dos especialistas não cabem inteiramente na memória da GPU (HBM). Isso exige o offloading (transferência) da maioria dos pesos para a memória RAM da CPU, transferindo apenas um subconjunto para a GPU durante a execução.

• Gargalo: A latência de transferência de dados entre CPU e GPU (via PCIe) torna-se o fator limitante, dominando o tempo por token de saída (TPOT). Em alguns casos, as transferências de memória representam 84-88% do tempo total de inferência, enquanto o cálculo computacional ocupa uma fração mínima.
• Limitação das Soluções Atuais: Métodos existentes de "caching" ou "prefetching" (busca antecipada) tratam previsões erradas como "misses" de cache, exigindo que o sistema pare e recarregue o especialista correto sob demanda. Isso interrompe a sobreposição entre computação e transferência de dados, limitando os ganhos de desempenho.

2. Metodologia Proposta

Os autores propõem um esquema de prefetching (busca antecipada) de especialistas que utiliza representações internas do modelo para prever quais especialistas serão necessários na próxima camada, permitindo que as transferências de memória ocorram em paralelo com a computação atual.

Componentes Chave:

• Estado Quase-Oculto (Quasi-Hidden State - $q_l$):

  • Em vez de usar apenas o fluxo residual normalizado ($s_l$) que alimenta o roteador, os autores definem um estado $q_l$ que combina o fluxo residual pós-atenção com um vetor padrão ($d_l$).
  • O vetor padrão representa a ativação média associada a cada especialista, calculada offline.
  • A fórmula é: $q_l = LN_{l+1}(d_l + r_l)$, onde $r_l$ é o resíduo pós-atenção.
  • Objetivo: $q_l$ atua como um sinal condicional ao especialista que melhora a precisão da previsão da próxima camada, especialmente em arquiteturas onde há "deriva" (drift) significativa entre as camadas.

• Execução Especulativa (Speculative Execution):

  • Diferente de abordagens anteriores que descartam previsões erradas, este método executa os especialistas previstos (os "especialistas especulados").
  • Se a previsão estiver correta, a transferência de dados já ocorreu e a computação avança sem atraso.
  • Se estiver incorreta, o sistema assume que a execução do especialista errado não degrada significativamente a precisão da tarefa downstream, evitando o custo de parar e recarregar o especialista "verdadeiro".

• Estimadores Neurais Leves (Opcional):

  • Para arquiteturas onde a deriva de representação é alta (ex: camadas iniciais do Qwen3-30B-A3B), os autores treinam um estimador neural leve (feed-forward network) para prever os logits do roteador diretamente a partir de $q_l$.
  • Isso melhora a taxa de acerto (hit rate) em camadas críticas sem adicionar um gargalo serial significativo.

• Implementação no Motor YALIS:

  • O esquema foi integrado ao motor de inferência open-source YALIS.
  • Utiliza double buffering e streams assíncronos CUDA para sobrepor a cópia de dados (CPU $\to$ GPU) da camada $l+1$ com a computação da camada $l$.

3. Principais Contribuições

1. Prefetching sem Parâmetros: Identificação de que representações internas (especificamente o estado quase-oculto) contêm sinais suficientes para prever decisões de roteamento futuras em grandes pools de especialistas, sem necessidade de re-treinamento do modelo base.
2. Execução Especulativa que Preserva Precisão: Demonstração de que executar especialistas previstos (em vez de tratá-los como erros de cache) mantém a precisão em tarefas complexas (raciocínio, matemática, código), eliminando a necessidade de re-buscar especialistas corretos.
3. Implementação Otimizada: Integração completa em um motor de inferência, demonstrando ganhos reais de latência em configurações de hardware limitado.
4. Estimadores Híbridos: Proposta de uma estratégia híbrida que usa estimadores neurais apenas em camadas com alta deriva de representação, equilibrando precisão e desempenho.

4. Resultados Experimentais

Os experimentos foram realizados em modelos como Qwen3-30B-A3B, GPT-OSS-20B/120B e GLM-4.7-Flash em GPUs como A6000, A100 e GH200.

• Redução de TPOT:

  • O método alcançou uma redução de 5% a 14% no Tempo por Token de Saída (TPOT) em comparação com o carregamento sob demanda (on-demand).
  • Os ganhos foram maiores em GPUs com menor throughput computacional (ex: A6000), onde o tempo de cópia domina o tempo total. Em GPUs mais potentes (A100/GH200), o ganho foi de 5-8%, pois o cálculo se torna mais rápido, reduzindo a janela de sobreposição útil.
  • Em sequências longas (65k tokens), os ganhos foram mais pronunciados devido ao aumento do tempo de computação que pode ser sobreposto à transferência.

• Precisão (Accuracy):

  • Para modelos como GPT-OSS, a execução especulativa baseada apenas no roteador manteve a precisão em benchmarks de raciocínio (HumanEval, GSM8k, StrategyQA) quase idêntica à linha de base.
  • Para Qwen3-30B-A3B, houve uma degradação inicial em tarefas matemáticas pesadas (AIME, GSM8k) devido à alta deriva nas camadas iniciais.
  • Solução: O uso do Estimador Híbrido (Hybrid-PF) recuperou a maior parte da precisão perdida (ex: recuperou ~37% da lacuna de precisão no GSM8k), aplicando o estimador neural apenas nas camadas críticas.

• Análise de Overlap:

  • Perfis de Nsight Systems mostraram que o prefetching efetivamente esconde a latência de transferência de CPU-GPU, transformando um gargalo de I/O em um processo quase totalmente sobreposto à computação.

5. Significado e Impacto

• Viabilidade de Deslocamento Local: Este trabalho torna a execução de modelos MoE massivos (com centenas de bilhões de parâmetros) mais viável em hardware de consumidor (GPUs de desktop com memória limitada), reduzindo a dependência de clusters de servidores caros.
• Paradigma de Execução: A descoberta de que a execução de especialistas "errados" (mas previstos) é tolerável para a precisão do modelo abre caminho para novas arquiteturas de inferência que priorizam o fluxo de dados contínuo em vez da correção estrita em tempo real.
• Caminho para Roteamento Livre: A técnica sugere que, no futuro, o roteador tradicional por camada poderia ser substituído por estimadores leves, eliminando a sobrecarga serial do roteamento e permitindo paralelismo ainda maior.

Em resumo, o artigo apresenta uma solução prática e eficiente para o gargalo de I/O em modelos MoE offloaded, utilizando previsões inteligentes de representações internas para sobrepor transferência de dados e computação, resultando em inferências mais rápidas sem sacrificar significativamente a qualidade do modelo.