Artigo
Acelerando o Sonar Através da Especulação
Decodificação especulativa acelera a velocidade de geração de Modelos de Linguagem de Grande Porte (LLMs) ao usar um modelo rápido e pequeno para produzir candidatos a conclusão que são verificados pelo modelo alvo maior.
Neste esquema, em vez de uma execução cara do alvo produzindo um único token, múltiplos são emitidos em um único passo. Aqui, apresentamos os detalhes de implementação de diversos 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 transformers para acelerar a geração de tokens. Mesmo que modelos maiores, como o Llama-70B, contenham mais conhecimento que os menores, como o Llama-1B, em algumas tarefas mais simples eles desempenham de forma semelhante. Essa sobreposição sugere que certas sequências são melhor geradas pelos modelos menos caros, deixando problemas complexos para os maiores. O desafio é determinar quais conclusões são melhores e se a geração do modelo menor é de qualidade igual à do maior.
Felizmente, os LLMs são transformers autorregressivos: quando fornecidos com uma sequência de tokens, eles fornecem 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 quão provável é que o modelo emita exatamente esses tokens. Esta propriedade permite especulação: se uma sequência de tokens é gerada por um modelo menor a partir de um prefixo de entrada, ela pode ser rodada no modelo maior para determinar quão bem se alinha com o modelo alvo. Cada prefixo dos candidatos é avaliado com uma probabilidade e o mais longo acima de um limite de aceitação é escolhido. Como um 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 um passo.
No momento da inferência, o processo de amostragem especulativa pode ser dividido em aproximadamente 4 etapas:
Pré-preenchimento: tanto os modelos alvo quanto o de rascunho devem ser executados na sequência de entrada para preencher as entradas do cache KV. Enquanto alguns esquemas, como Medusa, usam camadas densas mais simples para previsão, neste post focamos em rascunhos baseados em transformers que precisam de seus próprios caches KV.
Geração de Rascunhos: o modelo de rascunho itera para produzir um número fixo de tokens. A sequência de rascunho pode ser linear ou o modelo pode explorar uma estrutura em forma de á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 de Alvo: como os logits gerados pelo alvo, 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. Neste post, focaremos nos esquemas que usamos para acelerar os modelos Sonar usando um modelo interno de 1B, assim 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 ao acoplar um LLM pequeno existente como um 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. Enquanto essa abordagem não exigiu treinar um rascunho do zero, o modelo pequeno ainda usa uma capacidade significativa de cache KV e introduz uma leve sobrecarga de pré-preenchimento, aumentando o TTFT.
Neste esquema, o decodificador só especula em lotes apenas de decodificação, gerando tokens através da amostragem padrão durante o pré-preenchimento ou em lotes mistos de pré-preenchimento-decodificação. Na etapa de pré-preenchimento, os logits alvo são imediatamente amostrados para também pré-preencher o token recém-gerado no cache KV do rascunho. O rascunho ainda não é amostrado, mas os logits que produz são transferidos para a etapa de decodificação.
Na decodificação, o modelo de rascunho é avançado, amostrando o token superior em cada etapa. Após atingir o comprimento de rascunho desejado, os tokens são executados pelo modelo alvo para produzir os logits baseados nos quais o amostrador identifica o comprimento da sequência aceita. A aceitação é determinada pela comparação das 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 esse 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 várias sequências de rascunho, geradas através de uma travessia em forma de árvore de tokens de rascunho prováveis. Uma árvore fixa (EAGLE) ou dinamicamente moldada (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 avaliadas e a mais longa adequada é selecionada para continuar, também anexando um token adicional do alvo.
Para alcançar uma previsão mais precisa, um modelo de rascunho EAGLE prevê não apenas com base em tokens, mas também usando os recursos do alvo (últimas camadas de estados ocultos) do modelo alvo. A desvantagem do EAGLE é a necessidade de treinar modelos de rascunho pequenos e personalizados que são suficientemente precisos para gerar candidatos adequados dentro de um orçamento de baixa latência. Tipicamente, um modelo de rascunho é uma única camada de transformer idêntica a uma camada de decodificador do modelo original, que é estreitamente acoplada ao alvo por amarrage 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 forma de á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 desacelera significativamente a atenção para comprimentos de entrada realistas (em até 50%), anulando parte do aumento de velocidade alcançável através da especulação. Ainda não implementamos totalmente a exploração de árvores em produção por essa razão, focando em vez disso no caso especial de previsão de um único token via esquemas tipo MTP apresentados no Relatório Técnico DeepSeek-V3.
MTP
Este esquema é semelhante à decodificação de rascunho-alvo, com a exceção de que estados ocultos são usados juntamente com tokens para previsão. É necessário um pouco mais de trabalho tanto nas etapas de pré-preenchimento quanto de decodificação em comparação com a especulação regular de rascunho-alvo. O modelo de rascunho usa tanto tokens quanto 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 tokens de entrada devem ser deslocados um passo para a esquerda em relação aos vetores de estado oculto de saída do alvo. A figura abaixo marca as correspondências usadas para treino, assim como o deslocamento durante a inferência.
O fluxo de decodificação é bastante similar à decodificação de rascunho-alvo, com a exceção de que tanto estados ocultos quanto logits são transferidos. Nossa implementação compartilha todos os logíticos de amostragem e processamento de logit associados, especializando apenas as invocações forward do modelo. Quando múltiplos tokens são previstos, o modelo de rascunho usa os estados ocultos de rascunho para previsão, também populando entradas de cache KV baseadas em seus próprios recursos. A longo prazo, isso pode degradar a precisão. Subsequentemente, ao rodar o modelo de rascunho para preencher a entrada de cache KV para a previsão alvo, o rodamos 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 extras é desprezível.
Treinamento de Cabeçalhos MTP
Para se beneficiar do MTP, construímos a infraestrutura necessária para treinar cabeçalhos MTP anexados aos nossos modelos ajustados nos conjuntos de dados da Perplexity, rodando em um nó com dispositivos 8xH100. Em cerca de um dia, podemos construir cabeçalhos para modelos que variam do Llama-1B ao Llama-70B e DeepSeek V2-Lite. Para modelos maiores, contamos com cabeçalhos MTP construídos durante o processo de ajuste fino.
O alvo do treinamento MTP é alinhar os estados ocultos de rascunho e os logits extrapolados dos estados ocultos do alvo aos próximos logits de token e estados ocultos do alvo. Como a inferência de estados ocultos é cara, nós 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 da inferência MTP e garantir que diferenças numéricas devido à quantização ou otimizações não prejudiquem os resultados, para perda e estimativa de precisão de validação reutilizamos completamente a implementação de inferência de ambos os modelos alvo e de rascunho.
Ao aumentar a partir do conjunto de dados ShareGPT usado no artigo original para amostras maiores, notamos que a arquitetura de cabeçalho MTP delineada e implementada no artigo EAGLE falhou em treinar para modelos do tamanho 70B. Ao contrário 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 os cabeçalhos do EAGLE originais divergiram ligeiramente na estrutura de um transformer típico, reintroduzimos algumas camadas de Normalização RMS que foram removidas. Descobrimos que isso não apenas permitiu que o treinamento convergisse, mas também aumentou a precisão dos cabeçalhos em alguns pontos percentuais.
Não apenas as normas de camada facilitam o treinamento, mas reintroduzir as normas também é matematicamente intuitivo. Cabeçalhos MTP reutilizam as embeddings e as projeções de logit do modelo alvo, já que podem ser substanciais em tamanho (cerca de 2 GB para o Llama 70B). Durante o treinamento, esses são congelados 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 retirar as normas, espera-se que uma única MLP aprenda a mesma função de uma MLP seguida de uma norma, o que dificulta o alinhamento entre os estados ocultos dos rascunhos e do alvo.
Inferência com Decodificação Especulativa
No mecanismo de inferência, para gerar tokens para sequências de entrada, eles precisam primeiro ser 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 todos os ranks paralelos rodando o modelo. Finalmente, os metadados são copiados para a memória GPU e o modelo é executado para produzir os logíts 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 invólucro que orquestra solicitações entre eles, nossos pares de rascunho-alvo são estreitamente acoplados e avançam na geração em uníssono. O agendamento de lotes e a alocação de páginas KV são compartilhados entre os modelos para todas as formas de decodificação especulativa: isso unifica a lógica que conecta um modelo ao servidor de inferência global, já que todos eles expõem a mesma interface.
O tempo de execução da 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 pré-preenchimento, decodificação ou verificação, trabalho no lado da CPU deve ser realizado para alocar buffers intermediários e preencher certos buffers constantes usados na atenção. Este trabalho é adicional ao custo do agendamento de lotes e alocação de páginas KV, que também incorrem em latências que devem ser ocultadas para maximizar a utilização da GPU.
Enquanto totalmente paralelizamos o trabalho do lado da CPU e da GPU para inferência sem especulação, descobrimos que o equilíbrio CPU-GPU para decodificação especulativa é mais intricado. 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 CPU-GPU difícil de evitar. Experimentamos diferentes esquemas de agendamento para melhor ocultar a latência do trabalho da CPU.
Agendamento Rascunho-Alvo
Apesar de ser menor que um modelo alvo, quando um LLM inteiro é usado como o rascunho, ainda introduz considerável latência na GPU, fornecendo algum espaço para ocultar operações caras da CPU. Como modelos menores não se beneficiam do paralelismo de tensor, há uma discrepância entre o número de ranks que um alvo e um rascunho são divididos. Em nossa implementação, o modelo de rascunho é executado apenas no rank líder de um grupo TP.
Conforme indicado anteriormente, um passo de decodificação transporta 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 para o amostrador e o 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 alvo são verificados e amostrados para determinar os comprimentos das sequências aceitas. Neste ponto, a sincronização CPU-GPU é necessária para determinar os comprimentos de sequência subsequentes. 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 este rascunho executado 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 Único Token
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 produzidos por um modelo do tamanho de uma única camada de decodificador de transformer. Este esquema pode ser usado para previsão de rascunho usando os pesos de código aberto do DeepSeek R1. O subcaso de previsão de um único token é interessante, já que grandes camadas MTP alcançam taxas de aceitação suficientemente altas para justificar seu overhead.
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 é dividido 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 remanescentes, de forma semelhante ao esquema anterior. Em seguida, o alvo é executado para validar os tokens, processando 2 * D tokens, onde D é o tamanho do lote de decodificação. Isso é ideal para micro-loteamento em modelos Mixture-of-Experts (MoE) sobre conexõ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.
Ao realizar uma quantidade limitada de trabalho adicional na GPU, evitamos a sincronização CPU-GPU após a aceitação da sequência de rascunho. Após os tokens de entrada dos alvos serem deslocados, um kernel insere os próximos tokens alvo em seus locais correspondentes. O rascunho é então re-executado com as mesmas informações de lote que o alvo, populando 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 é pouco mensurável devido ao pequeno tamanho do modelo de rascunho. Paralelamente à execução do rascunho, os comprimentos das sequências são determinados na CPU e o agendamento do próximo lote é iniciado, sem a necessidade de esperar pelo término do trabalho na GPU.
O overhead de trabalho adicional na camada de rascunho não é notável na atenção, mas as camadas MLP são mais problemáticas. Como as instruções de multiplicação de matrizes preenchem até um limite de 64 na dimensão do número de tokens, se dobrar não exigir significativamente mais blocos, o overhead é ocultado. Para sequências de rascunho mais longas, o overhead é mais caro e o esquema usado para modelos rascunho-alvo regulares funciona melhor.
Referências
Inferência Rápida de Transformers via Decodificação Especulativa
EAGLE: Amostragem Especulativa Requer Repensar a Incerteza dos Recursos
EAGLE-2: Inferência Mais Rápida de Modelos de Linguagem com Árvores de Rascunho Dinâmicas
FlashInfer: Motor de Atenção Eficiente e Personalizável para Servir Inferência de LLM