P-EAGLE: Inferência de LLM Mais Rápida com Decodificação Especulativa Paralela no vLLM

O Desafio da Decodificação Especulativa em LLMs

EAGLE estabeleceu-se como o método de referência para decodificação especulativa na inferência de modelos de linguagem de grande escala (LLM). Sua abordagem inovadora permite gerar múltiplos tokens especulados simultaneamente, reduzindo a latência de forma significativa. Porém, sua estratégia de geração autorregressiva esconde um gargalo crítico: quanto mais tokens são especulados, mais passagens sequenciais pela rede neural o modelo auxiliar precisa realizar. Eventualmente, esse overhead consome os ganhos conquistados.

A AWS identificou essa limitação e desenvolveu P-EAGLE, uma evolução que remove esse teto de desempenho ao gerar todos os K tokens de rascunho em uma única passagem pela rede. O resultado? Até 1,69× de aceleração em relação ao EAGLE-3 vanilla em cenários reais utilizando GPUs NVIDIA B200.

Para começar, basta fazer o download (ou treinar) de uma cabeça de modelo auxiliar compatível com paralelismo e adicionar "parallel_drafting": true à sua configuração de pipeline no vLLM. Modelos P-EAGLE pré-treinados já estão disponíveis no HuggingFace para GPT-OSS 120B, GPT-OSS 20B e Qwen3-Coder 30B.

Como P-EAGLE Funciona

Eliminando o Gargalo Sequencial

EAGLE consegue acelerações de 2–3× sobre decodificação autorregressiva padrão e está amplamente implementado em estruturas de inferência em produção, incluindo vLLM, SGLang e TensorRT-LLM. No entanto, sua abordagem gera tokens especulados de forma sequencial — para produzir K tokens de rascunho, são necessárias K passagens completas pela rede do modelo auxiliar. À medida que esses modelos melhoram sua capacidade de geração, o overhead dessa abordagem se torna cada vez mais significativo: a latência do modelo auxiliar escala linearmente com a profundidade de especulação, limitando agressividade possível.

A Abordagem P-EAGLE

P-EAGLE transforma EAGLE de geração de rascunho autorregressiva para paralela. Em GPUs B200, o método alcança 1,05×–1,69× de aceleração em relação ao EAGLE-3 vanilla no GPT-OSS 20B, testado em MT-Bench, HumanEval e SpeedBench. Agora integrado ao vLLM desde a versão 0.16.0, está pronto para acelerar implantações no mundo real.

Dois Passos Fundamentais

Passo 1: Prefenchimento. O modelo alvo processa o prompt e gera um novo token, como faria durante inferência normal. Durante esse processo, P-EAGLE captura os estados ocultos internos do modelo: h_prompt para cada posição do prompt e h_context para o token recém-gerado. Esses estados codificam o que o modelo alvo “compreende” em cada posição e guiam as predições do modelo auxiliar. Este passo é idêntico ao EAGLE autorregressivo.

Passo 2: Modelo Auxiliar P-EAGLE. O modelo auxiliar constrói entradas para cada posição em paralelo. Cada entrada consiste em um embedding de token concatenado a um estado oculto. Para posições do prompt, cada embedding de token do prompt (emb(p)) é pareado com seu correspondente h_prompt do modelo alvo. Seguindo a mesma convenção do EAGLE autorregressivo, as posições são deslocadas por uma unidade: a posição i recebe o token e estado oculto da posição i-1, permitindo que ela prediga o token na posição i.

Para a posição 1 (Predição do Próximo Token — NTP), a entrada pareia o embedding do token recém-gerado (emb(new)) com h_context. Essa posição funciona de forma idêntica ao EAGLE autorregressivo padrão.

Para as posições 2 a K (Predição de Múltiplos Tokens — MTP), as entradas necessárias — embedding de token e estado oculto — ainda não existem. P-EAGLE preenche esses espaços com dois parâmetros treináveis: um embedding de máscara compartilhado (emb(mask)) e um estado oculto compartilhado (h_shared). Esses são vetores fixos aprendidos durante o treinamento que atuam como placeholders neutros. Todas as posições passam simultaneamente por N camadas de transformador e, em seguida, pela cabeça do modelo de linguagem para predizer tokens de rascunho t1, t2, t3 e t4 em uma única passagem.

Treinamento em Sequências Longas

Modelos de raciocínio moderno geram saídas extensas. No conjunto de dados UltraChat, GPT-OSS 120B produz sequências (incluindo prompts) com comprimento mediano de 3.891 tokens e percentil 90 de 10.800 tokens. Modelos auxiliares devem ser treinados em comprimentos de contexto correspondentes para serem efetivos na inferência.

Um desafio-chave é que a decodificação especulativa paralela amplifica requisitos de memória durante o treinamento. Treinar K grupos paralelos em uma sequência de comprimento N cria N × K posições totais. Com N = 8.192 e K = 8, um único exemplo de treinamento contém 65.536 posições. Atenção requer que cada posição atenda a todas as posições válidas — 65K × 65K significa mais de 4 bilhões de elementos, consumindo 8GB em bf16.

Amostragem de posições reduz memória saltando aleatoriamente posições, mas pular de forma muito agressiva degrada a qualidade do rascunho. Acumulação de gradientes é a solução padrão para treinamento com memória limitada, mas divide entre diferentes exemplos de treinamento. Quando uma sequência única excede a memória disponível, não há nada para dividir.

P-EAGLE introduz um algoritmo de particionamento de sequência para divisão intra-sequência. O algoritmo divide a sequência de posições N × K em pedaços contíguos, mantém dependências de atenção corretas através das fronteiras dos pedaços e acumula gradientes através de pedaços da mesma sequência. Para detalhes técnicos, consulte o paper de P-EAGLE.

Integração no vLLM

Desafios da Decodificação Especulativa Paralela

Em muitas configurações de decodificação especulativa, rascunho e verificação compartilham o mesmo layout de token por requisição. Isso é geralmente verdadeiro para EAGLE: o modelo auxiliar consome uma janela que já corresponde ao que o verificador checará — K tokens especulados e um token adicional amostrado.

Decodificação paralela quebra essa consistência. Para predizer K tokens em uma passagem do modelo auxiliar, placeholders de MÁSCARA são anexados (por exemplo, [token, MÁSCARA, MÁSCARA, …]). Essas posições extras existem apenas para rascunho, então a forma do lote de rascunho não corresponde mais à forma do lote de verificação. Como metadados de verificação não podem ser reutilizados, o metadata do lote deve ser reconstruído. Os IDs de token de entrada, estados ocultos e posições são expandidos para inserir espaços para tokens/embeddings de máscara, posições por requisição são incrementadas, então o mapeamento de slot e índices de início por requisição são recomputados a partir das posições atualizadas.

Kernel Triton Otimizado

Para compensar o overhead da reconstrução de metadata do lote, um kernel Triton fundido popula o lote de entrada do modelo auxiliar na GPU copiando e expandindo o lote do modelo alvo. Em uma única passagem, o kernel copia os IDs de token e posições anteriores do lote alvo para novos slots de destino, insere o token bônus por requisição amostrado pelo modelo alvo e preenche os slots adicionais de decodificação paralela com um ID de token MÁSCARA especial. Por fim, gera metadata leve: máscara de token rejeitado, máscara de token mascarado para slots de decodificação paralela, índices de novo-token para amostragem de tokens de rascunho e mapeamento de estado oculto. Essa lógica, de outra forma, demandaria muitas operações GPU (cópia/espalhamento + inserção + preenchimento + máscara + remapeamento). Fundindo em um kernel único reduz overhead de lançamento e tráfego extra de memória, mantendo a configuração de rascunho eficiente.

Gestão de Estados Ocultos

Para métodos baseados em EAGLE que passam estados ocultos ao modelo auxiliar, decodificação paralela trata esse preenchimento separadamente. Como estados ocultos são significativamente maiores que o resto do lote de entrada, o trabalho é dividido: o kernel Triton produz um mapeamento, e um kernel de cópia dedicado transmite o placeholder de estado oculto aprendido para slots de token mascarado.

# Copiar estados ocultos alvo para suas novas posições
self.hidden_states[out_hidden_state_mapping] = target_hidden_states

# Preencher posições mascaradas com o estado oculto aprendido de Decodificação Paralela
mask = self.is_masked_token_mask[:total_num_output_tokens]
torch.where(
    mask.unsqueeze(1),
    self.parallel_drafting_hidden_state_tensor,
    self.hidden_states[:total_num_output_tokens],
    out=self.hidden_states[:total_num_output_tokens],
)

O parallel_drafting_hidden_state_tensor é carregado do buffer mask_hidden do modelo, uma representação aprendida que informa ao modelo que essas posições devem predizer tokens futuros. Para mapeamento de slot de cache KV, tokens válidos recebem atribuição de slot normal enquanto tokens rejeitados são mapeados para PADDING_SLOT_ID (-1) para evitar escritas spurias em cache. Para grafos CUDA, o intervalo de captura é estendido por K × max_num_seqs para acomodar o lote de rascunho maior introduzido pela decodificação paralela.

Resultados de Benchmark

P-EAGLE foi treinado em GPT-OSS-20B e avaliado em três benchmarks: MT-Bench para instrução multi-turno, SPEED-Bench para geração de código longo e HumanEval para síntese de código em nível de função. P-EAGLE entrega 55–69% de throughput superior em baixa concorrência (c=1), com ganhos de 5–25% sustentados em alta concorrência (c=64), comparado ao checkpoint público vanilla EAGLE-3.

O modelo auxiliar P-EAGLE é um modelo leve de 4 camadas treinado para predizer até 10 tokens em paralelo. Para avaliação de desempenho, profundidades de especulação K ∈ {3, 5, 7} foram testadas em níveis de concorrência C ∈ {1, 2, 4, 8, 16, 32, 64}. P-EAGLE alcança TPS (Tokens Por Segundo) máximo em K=7 em todos os níveis de concorrência. Em contraste, EAGLE-3 vanilla atinge seu TPS máximo em K=3, com profundidade melhorada ocasionalmente deslocando para valores maiores dependendo da concorrência. Esse padrão reflete a vantagem fundamental da decodificação paralela.

P-EAGLE gera todos os K tokens de rascunho em uma única passagem, permitindo que se beneficie de especulação mais profunda sem incorrer em overhead sequencial adicional. Rascunhadores autorregressivos, por contraste, devem gerar tokens especulativos passo a passo, o que limita sua capacidade de escalar eficientemente para K maiores.

Comprimento de Aceitação Aprimorado

Além de reduzir overhead de rascunho, ganhos de throughput do P-EAGLE também são impulsionados por melhor comprimento de aceitação (AL — do inglês Acceptance Length) — a média de tokens de rascunho aceitos pelo verificador por rodada de especulação. AL maior significa mais do trabalho de rascunho se transforma em saída real, o que impulsiona diretamente OTPS/TPS (Tokens de Saída Por Segundo / Tokens Por Segundo) efetivo.

P-EAGLE alcança consistentemente AL superior ao EAGLE-3 na mesma profundidade de especulação K. Em K=7, P-EAGLE supera EAGLE-3 por 30% em HumanEval (3,94 vs 3,03), 31% em SPEED-Bench (3,38 vs 2,59) e 13% em MT-Bench (3,70 vs 3,27). Notavelmente, P-EAGLE se beneficia mais de especulação mais profunda. De K=3 para K=7, AL do P-EAGLE aumenta 0,92 em HumanEval (3,02 para 3,94), enquanto EAGLE-3 ganha apenas 0,38 (2,65 para 3,03). Essa lacuna crescente em K maior é consistente com decodificação especulativa paralela de passagem única do P-EAGLE, que não incorre em custo adicional de especulação mais profunda.

Como Começar

Configuração Básica

P-EAGLE pode ser ativado com uma mudança simples de configuração na classe SpeculativeConfig:

# vllm/config/speculative.py
parallel_drafting: bool = True

Aqui está um comando de exemplo no vLLM para ativar decodificação paralela com P-EAGLE como modelo auxiliar:

vllm serve openai/gpt-oss-20b \
  --speculative-config '{"method": "eagle3", "model": "amazon/gpt-oss-20b-p-eagle", "num_speculative_tokens": 5, "parallel_drafting": true}'

Executando Benchmarks

Após lançar o servidor, execute benchmarks com vllm bench serve:

#MT-Bench
export MODEL="openai/gpt-oss-20b"
export BASE_URL="http://localhost:8000"
vllm bench serve \
  --dataset-name hf \
  --dataset-path philschmid/mt-bench \
  --num-prompts 80 \
  --max-concurrency 1 \
  --model $MODEL \
  --base-url $BASE_URL \
  --temperature 0.0 \
  --hf-output-len 2048

#HumanEval
# Baixar dataset HumanEval openai/openai_humaneval
vllm bench serve \
  --dataset-name custom \
  --dataset-path  \
  --num-prompts 164 \
  --max-concurrency 1 \
  --model $MODEL \
  --base-url $BASE_URL \
  --temperature 0.0 \
  --custom-output-len 2048

Nota Importante: Servir GPT-OSS-20B com rascunhadores EAGLE atualmente requer um patch vLLM de uma linha (PR#36684). Aplique-o antes de lançar. Essa correção deve chegar em uma versão próxima do vLLM.

Recursos Disponíveis

Para explorar P-EAGLE, a AWS disponibilizou:

• Paper ArXiv com detalhes técnicos completos
• Modelos pré-treinados no HuggingFace: GPT-OSS 120B, GPT-OSS 20B, Qwen3-Coder-30B-A3B-Instruct
• Integração vLLM a partir da versão 0.16.0 com Unified Parallel Drafting
• Discussões técnicas: RFC e PR no repositório vLLM-Speculators

Conclusão

P-EAGLE remove o gargalo sequencial da decodificação especulativa, entregando até 1,69× de aceleração em relação ao EAGLE-3 vanilla em cargas reais. Ao desacoplar contagem de rascunho de contagem de passagens diretas, agora é possível explorar arquiteturas de rascunho maiores, que podem até permitir taxas de aceitação aumentadas comparadas a baselines de camada única.

A implementação lida com cuidado com complexidades de preparação de entrada, gerenciamento de metadata de atenção e mapeamento de slot de cache KV através de kernels fundidos escritos manualmente. Embora requeira modelos especialmente treinados, os benefícios de desempenho a tornam uma adição valiosa às capacidades de decodificação especulativa do vLLM. Conforme mais modelos treinados em paralelo se tornem disponíveis, espera-se que essa abordagem se torne a escolha preferida para implantações de LLM em produção.

Para começar agora: baixe uma cabeça P-EAGLE pré-treinada do HuggingFace, configure "parallel_drafting": true em seu config vLLM para qualquer um dos modelos suportados e veja a aceleração em ação.

Fonte

P-EAGLE: Faster LLM inference with Parallel Speculative Decoding in vLLM (https://aws.amazon.com/blogs/machine-learning/p-eagle-faster-llm-inference-with-parallel-speculative-decoding-in-vllm/)