Acelerando o Sonar através da especulação

A decodificação especulativa acelera a velocidade de geração de Grandes Modelos de Linguagem (LLMs) ao utilizar um modelo de rascunho pequeno e rápido para produzir candidatos de conclusão que são verificados pelo modelo alvo maior. Sob este esquema, em vez de uma execução do modelo alvo dispendioso produzir um único token, múltiplos são emitidos em uma única etapa. Aqui, apresentamos os detalhes de implementação de vários tipos de decodificação especulativa, aplicados na Perplexity para reduzir a latência entre tokens nos modelos Sonar.

Decodificação Especulativa

A Decodificação Especulativa aproveita a estrutura das linguagens naturais e a natureza autorregressiva dos transformadores para acelerar a geração de tokens. Embora modelos maiores, como o Llama-70B, possuam mais conhecimento do que modelos menores, como o Llama-1B, em algumas tarefas mais simples eles apresentam desempenho semelhante. Essa sobreposição sugere que certas sequências são melhor geradas por modelos menos dispendiosos, deixando problemas complexos para os maiores. O desafio reside em determinar quais conclusões são melhores e se a geração do modelo menor tem a mesma qualidade que a do modelo maior.

Felizmente, os LLMs são transformadores autorregressivos: quando recebem uma sequência de tokens, eles emitem a distribuição de probabilidade do próximo token. Além disso, os logits derivados das características intermediárias associadas aos tokens na sequência de entrada também indicam a probabilidade de o modelo emitir esses tokens exatos. Essa propriedade permite a especulação: se uma sequência de tokens é gerada por um modelo menor a partir de um prefixo de entrada, ela pode ser executada pelo modelo maior para determinar quão bem ela se alinha com o modelo alvo. Cada prefixo dos candidatos é pontuado com uma probabilidade e o mais longo acima de um limite de aceitação é escolhido. Como bônus, o modelo alvo também fornece um token subsequente gratuitamente: se um modelo de rascunho gera n tokens, até n + 1 podem ser emitidos em uma etapa.

Expandir

No momento da inferência, o processo de amostragem especulativa pode ser dividido em aproximadamente 4 etapas:

Prefill: tanto o modelo alvo quanto o de rascunho devem ser executados na sequência de entrada para preencher as entradas do cache KV. Embora alguns esquemas, como o Medusa, utilizem camadas densas mais simples para previsão, nesta publicação focamos em rascunhos baseados em transformadores que precisam de seus próprios caches KV.

Geração de rascunho: o modelo de rascunho itera para produzir um número de tokens fixos. A sequência de rascunho pode ser linear ou o modelo pode explorar uma estrutura em árvore até uma determinada profundidade (EAGLE, Medusa). Aqui, focamos em sequências lineares.

Aceitação: o modelo alvo é executado na sequência de rascunho, construindo logits correspondentes a cada token de rascunho. O comprimento da sequência aceitável mais longa é determinado.

Geração do alvo: como o alvo gerou logits, na posição incompatível ou no final da sequência, os logits correspondem a um token subsequente. Esses logits podem ser amostrados para fornecer um token robusto do alvo, finalizando a sequência.

Existem vários métodos para implementar a decodificação especulativa. Nesta publicação, focaremos nos esquemas que utilizamos para acelerar os modelos Sonar usando um modelo 1B interno, bem como nos mecanismos de previsão que estamos desenvolvendo para acelerar modelos na escala do DeepSeek.

Alvo-Rascunho

A decodificação especulativa pode ser alcançada acoplando um LLM pequeno existente como modelo de rascunho a um modelo alvo para gerar sequências candidatas. Em produção, aceleramos o Sonar usando um modelo Llama-1B ajustado no mesmo conjunto de dados que o alvo. Embora essa abordagem não tenha exigido o treinamento de um rascunho do zero, o modelo pequeno ainda utiliza capacidade significativa de cache KV e introduz uma ligeira sobrecarga de prefill, aumentando o TTFT.

Sob este esquema, o decodificador apenas especula sobre lotes de apenas decodificação, gerando tokens através de amostragem padrão durante o prefill ou em lotes mistos de prefill-decodificação. No estágio de prefill, os logits do alvo são imediatamente amostrados para também preencher o token recém-gerado no cache KV do rascunho. O rascunho ainda não é amostrado, mas os logits que ele produz são transferidos para o estágio de decodificação.

Expandir

Na decodificação, o modelo de rascunho avança, amostrando o token principal em cada estágio. Após o comprimento de rascunho desejado ser atingido, os tokens são executados através do modelo alvo para produzir os logits com base nos quais o amostrador identifica o comprimento da sequência aceita. A aceitação é determinada comparando as distribuições de probabilidade completas do rascunho e do alvo. Como o alvo sempre gera um conjunto de logits após a sequência de rascunho aceita, isso é amostrado para produzir uma saída adicional. Como o modelo de rascunho ainda não viu aquele token aceito, ele é reexecutado para preencher suas entradas de cache KV correspondentes em preparação para a próxima etapa de decodificação, transferindo os logits novamente.

EAGLE

O EAGLE é um esquema de decodificação especulativa que explora múltiplas sequências de rascunho, geradas através de um percurso em árvore de prováveis tokens de rascunho. Uma árvore fixa (EAGLE) ou de forma dinâmica (EAGLE-2) é explorada usando execuções consecutivas dos tokens de rascunho, considerando os candidatos Top-K em cada nó em vez de seguir o token com maior pontuação em uma sequência linear. As sequências são então pontuadas e a mais longa adequada é selecionada para continuar, acrescentando também um token adicional do alvo.

Expandir

Para obter uma previsão mais precisa, um modelo de rascunho EAGLE prevê não apenas com base em tokens, mas também usando as características alvo (estados ocultos da última camada) do modelo alvo. A desvantagem do EAGLE é a necessidade de treinar modelos de rascunho pequenos e personalizados que sejam precisos o suficiente para gerar candidatos adequados dentro de um baixo orçamento de latência. Normalmente, um modelo de rascunho é uma única camada de transformador idêntica a uma camada de decodificador do modelo original, que é fortemente acoplada ao alvo ligando-se às suas projeções de embeddings e lm_head. Como isso requer menos capacidade de cache KV, o EAGLE tem uma pegada de memória menor.

Para verificar sequências em árvore no modelo alvo, máscaras de atenção personalizadas devem ser usadas. Infelizmente, usar uma máscara de atenção personalizada para uma sequência inteira retarda significativamente a atenção para comprimentos de entrada realistas (em até 50%), anulando parte da aceleração alcançável através da especulação. Ainda não implantamos a exploração completa em árvore em produção por esse motivo, focando em vez disso no caso especial da previsão de token único via esquemas tipo MTP apresentados no Relatório Técnico do DeepSeek-V3.

MTP

Este esquema é semelhante à decodificação rascunho-alvo, com a exceção de estados ocultos sendo usados junto com tokens para previsão. Um pouco mais de trabalho deve ser feito tanto nos estágios de prefill quanto de decodificação em comparação com a especulação rascunho-alvo regular. O modelo de rascunho usa tokens e estados ocultos: o token t_{i+1} é amostrado dos logits L_i correspondentes ao token t_i, que por sua vez são derivados dos estados ocultos H_i. Consequentemente, os buffers de token de entrada devem ser deslocados um passo para a esquerda em relação aos vetores de estado oculto emitidos pelo alvo. A figura abaixo marca as correspondências usadas para treinamento, bem como o deslocamento durante a inferência.

Expandir

O fluxo de decodificação é bastante semelhante à decodificação rascunho-alvo, com a exceção de estados ocultos e logits sendo transferidos. Nossa implementação compartilha todos os logits de processamento de amostragem e logit associados, especializando apenas as invocações de avanço do modelo. Quando múltiplos tokens são previstos, o modelo de rascunho usa estados ocultos de rascunho para previsão, também preenchendo entradas de cache KV com base em suas próprias características. A longo prazo, isso pode degradar a precisão. Posteriormente, ao executar o modelo de rascunho para preencher a entrada do cache KV para a previsão do alvo, nós o executamos em toda a sequência, tomando os estados ocultos do alvo mais precisos como entrada. Como esses modelos de rascunho são pequenos, o custo adicional de processar os tokens adicionais é desprezível.

Treinando Cabeças MTP

Para nos beneficiarmos do MTP, construímos a infraestrutura necessária para treinar cabeças MTP anexadas aos nossos modelos ajustados nos conjuntos de dados da Perplexity, executando em um nó com 8 dispositivos H100. Em cerca de um dia, podemos construir cabeças para modelos que variam do Llama-1B ao Llama-70B e DeepSeek V2-Lite. Para modelos maiores, contamos com cabeças MTP construídas durante o processo de ajuste fino.

O alvo do treinamento MTP é combinar os estados ocultos e logits do rascunho extrapolados dos estados ocultos do alvo com os logits e estados ocultos do próximo token do alvo. Como a inferência para estados ocultos é cara, nós os pré-computamos usando nossa implementação otimizada para inferência do modelo alvo, para serem usados durante o treinamento. No entanto, para validar a implementação MTP de inferência e garantir que as diferenças numéricas devido à quantização ou otimizações não prejudiquem os resultados, para perda de validação e estimativa de precisão, reutilizamos totalmente a implementação de inferência de ambos os modelos, alvo e rascunho.

Ao escalar do conjunto de dados ShareGPT usado no artigo original para amostras maiores, notamos que a arquitetura de cabeça MTP delineada e implementada no artigo do EAGLE falhou ao treinar para modelos de tamanho 70B. Diferente do ShareGPT, que continha um número maior de sequências mais curtas, treinamos em um número ligeiramente menor de prompts substancialmente mais longos. Como as cabeças EAGLE originais divergiam ligeiramente em estrutura de um transformador típico, reintroduzimos algumas camadas de Normalização RMS que foram removidas. Descobrimos que isso não só permitiu que o treinamento convergisse, mas também aumentou a precisão das cabeças em alguns pontos percentuais.

Expandir

As normas de camada não apenas facilitam o treinamento, mas reintroduzir as normas também é matematicamente intuitivo. As cabeças MTP reutilizam os embeddings e as projeções de logit do modelo alvo, pois podem ser substanciais em tamanho (cerca de 2 GB para o Llama 70B). Durante o treinamento, elas são congeladas e a expectativa é que a camada MTP aprenda a incorporar previsões no mesmo Espaço vetorial que a camada de projeção do modelo original aprendeu durante o treinamento. Ao remover as normas, espera-se que um único MLP aprenda a mesma função que um MLP seguido por uma norma, o que dificulta a correspondência entre os estados ocultos do rascunho e dos modelos alvo.

Inferência com Decodificação Especulativa

No mecanismo de inferência, para gerar tokens para sequências de entrada, eles precisam ser primeiro agrupados em lotes de tamanho razoável, então páginas devem ser alocadas no cache KV para os próximos tokens. Os tokens de entrada e as informações da página KV são então empacotados em um buffer transmitido para todas as classificações paralelas executando o modelo. Finalmente, os metadados são copiados para a memória da GPU e o modelo é executado para produzir os logits a partir dos quais o próximo token é amostrado.

Ao contrário de certas implementações que acoplam frouxamente um servidor de inferência de rascunho e alvo via um wrapper que orquestra solicitações entre eles, nossos pares rascunho-alvo são rigidamente acoplados e passam pela geração em uníssono. O agendamento de lotes e a alocação de página KV são compartilhados entre os modelos para todas as formas de decodificação especulativa: isso unifica a lógica que une um modelo ao servidor de inferência abrangente, pois todos expõem a mesma interface.

O tempo de execução de inferência na Perplexity é moldado em torno do FlashInfer, que determina os metadados que precisam ser construídos para configurar e agendar o kernel de atenção. Dadas algumas sequências de entrada formando um lote, para prefill, decodificação ou verificação, o trabalho no lado da CPU deve ser feito para alocar buffers intermediários e preencher certos buffers constantes usados na atenção. Este trabalho é adicional ao custo do agendamento de lotes e da alocação de página KV, que também incorrem em latências que devem ser ocultadas para maximizar a utilização da GPU.

Embora tenhamos paralelizado totalmente o trabalho no lado da CPU e da GPU para inferência sem especulação, descobrimos que o equilíbrio CPU-GPU para decodificação especulativa é mais complexo. O principal desafio surge do fato de que o número de tokens aceitos determina o comprimento da sequência para uma execução subsequente, introduzindo um ponto de sincronização GPU-para-CPU difícil de evitar. Experimentamos diferentes esquemas de agendamento para ocultar melhor a latência do trabalho da CPU.

Agendamento Rascunho-Alvo

Apesar de ser menor que um modelo alvo, quando um LLM inteiro é usado como rascunho, ele ainda introduz uma latência considerável na GPU, proporcionando alguma margem para ocultar operações de CPU dispendiosas. Como modelos menores não se beneficiam de paralelismo de tensor, há uma incompatibilidade entre o número de classificações em que um alvo e um rascunho são fragmentados. Em nossa implementação, o modelo de rascunho é executado apenas na classificação líder de um grupo TP.

Expandir

Como indicado anteriormente, uma etapa de decodificação transfere logits para a próxima execução. Isso nos permite sobrepor uma execução do modelo de rascunho com o trabalho de agendamento de lotes no lado da CPU. Após o lote ser montado, chamadas repetidas ao amostrador e ao rascunho produzem os tokens de rascunho. Em paralelo, o lote para verificação é montado para o modelo alvo e sincronizado com os trabalhadores paralelos. Os logits do alvo são verificados e amostrados para determinar os comprimentos da sequência aceita. Neste ponto, a sincronização GPU-para-CPU é necessária para determinar os comprimentos da sequência subsequente. Como o modelo de rascunho é executado apenas no nó líder, seu lote é configurado sequencialmente e sua execução é iniciada para preencher suas entradas de cache KV com o token adicional que o alvo produziu. Os logits produzidos por esta execução de rascunho na execução atual serão usados para amostrar o primeiro token de rascunho na execução subsequente. Mais importante, enquanto o rascunho está sendo executado, o próximo lote pode ser agendado.

Agendamento MTP para um Token Único

Embora o tempo de execução ainda não forneça exploração de árvore de rascunho estilo Eagle, implementamos um caso especial deste esquema, considerando uma sequência linear de tokens de rascunho produzida por um modelo do tamanho de uma única camada de decodificador de transformador. Este esquema pode ser usado para previsão de rascunho usando os pesos de código aberto do DeepSeek R1. O subcaso de prever um único token é interessante, pois as camadas MTP grandes alcançam taxas de aceitação suficientemente altas para justificar sua sobrecarga.

O agendamento MTP é um pouco mais complexo, pois o modelo de rascunho é muito mais rápido, ocultando menos latência do lado da CPU. Além disso, o rascunho é fragmentado junto com o modelo alvo, exigindo transferências de memória compartilhada para informações de lote. Uma execução começa transferindo informações de lote e amostrando o primeiro token de logits transferidos, de forma semelhante ao esquema anterior. Em seguida, o alvo é executado para validar tokens, processando 2 * D tokens, onde D é o tamanho do lote de decodificação. Isso é ideal para micro-agrupamento em modelos de Mistura de Especialistas (MoE) sobre interconexões mais lentas, como InfiniBand, já que o lote se divide uniformemente em duas metades. Os estados ocultos do alvo são transferidos para a próxima execução de rascunho, enquanto os logits são passados para o amostrador para verificação.

Expandir

Ao realizar uma quantidade limitada de trabalho adicional na GPU, evitamos a sincronização CPU-para-GPU após a aceitação da sequência de rascunho. Depois que os tokens de entrada dos alvos são deslocados, um kernel conecta os próximos tokens alvo em suas localizações correspondentes. O rascunho é então reexecutado com as mesmas informações de lote que o alvo, preenchendo entradas de cache KV e construindo os logits e estados ocultos para a próxima execução, fazendo algum trabalho redundante em tokens que não foram aceitos. Nessas situações, a latência do trabalho não utilizado é mal mensurável devido ao tamanho pequeno do modelo de rascunho. Em paralelo com a execução do rascunho, os comprimentos da sequência são determinados na CPU e o agendamento do próximo lote é iniciado, sem ter que esperar a conclusão do trabalho da GPU.

A sobrecarga de trabalho adicional na camada de rascunho não é perceptível na atenção, no entanto, as camadas MLP são mais problemáticas. Como as instruções de multiplicação de matriz preenchem até um limite de 64 ao longo da dimensão do número de tokens, se dobrar não exigir significativamente mais blocos, a sobrecarga é ocultada. Para sequências de rascunho mais longas, a sobrecarga é mais cara e o esquema usado para modelos rascunho-alvo regulares funciona melhor.

Referências

Inferência Rápida de Transformadores via Decodificação Especulativa

EAGLE: A Amostragem Especulativa Exige Repensar a Incerteza de Recursos

EAGLE-2: Inferência Mais Rápida de Modelos de Linguagem com Árvores de Rascunho Dinâmicas

EAGLE-3: Escalonando a Aceleração de Inferência de Grandes Modelos de Linguagem via Teste de Tempo de Treinamento

Medusa: Estrutura de Aceleração de Inferência LLM Simples com Múltiplas Cabeças de Decodificação

Normalização de Camada

FlashInfer: Mecanismo de Atenção Eficiente e Personalizável para Atendimento de Inferência LLM

Relatório Técnico DeepSeek-V3

Citação

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}