AWS mostra como o speculative decoding no Trainium2 acelera a geração no vLLM

A AWS apresentou uma forma prática de acelerar e baratear a inferência de LLMs no Trainium2 em cenários nos quais o modelo gera muito mais tokens do que recebe na entrada. A técnica é o speculative decoding. Em vez de obrigar o modelo grande a produzir um token por vez, em sequência, o sistema usa um pequeno modelo draft, que propõe rapidamente vários tokens seguintes.

O modelo principal, o target, verifica todos eles em uma única passagem. Quando as previsões batem, o serviço executa menos etapas sequenciais caras, reduz a latência entre tokens e aproveita melhor o acelerador. Isso é especialmente relevante para cargas decode-heavy: assistentes de escrita, coding agents, geração de relatórios, documentos baseados em templates e outras tarefas com respostas longas.

Na geração autorregressiva convencional, cada novo token é calculado separadamente. Por isso, o acelerador lê o KV-cache da memória o tempo todo e faz relativamente pouco trabalho útil a cada etapa. Com isso, a inferência costuma esbarrar na largura de banda de memória, e não na capacidade de computação.

O speculative decoding ataca justamente esse gargalo: o modelo target executa menos etapas sequenciais de decode, e a verificação de um lote de tokens torna a carga mais densa. A abordagem, porém, tem requisitos. Os modelos draft e target precisam usar o mesmo tokenizador e o mesmo vocabulário.

O ideal é que também pertençam à mesma família de arquitetura, para que o modelo menor acerte com mais frequência a continuação do principal. O parâmetro-chave é o número de speculative tokens. Se a janela for pequena demais, o ganho é quase imperceptível; se for grande demais, rejeições precoces e verificações extras anulam o efeito.

No teste, a AWS usou o Qwen3-32B como modelo target e o Qwen3-1.7B como draft, rodando via vLLM em uma instância Trn2 trn2.48xlarge.

Para o speculative decoding, foi escolhida a fused speculation do NeuronX Distributed Inference, em que os dois modelos são compilados juntos para obter melhor desempenho. As configurações baseline e speculative foram implantadas no mesmo cluster Amazon EKS, com todo o resto idêntico: alocação de aceleradores, tensor parallelism, tamanho de contexto, batch limits e imagem Neuron. A única diferença foi a adição do modelo draft e do parâmetro num_speculative_tokens.

A carga nos dois serviços foi gerada com o llmperf, e as métricas de TTFT, inter-token latency e latência ponta a ponta foram enviadas ao CloudWatch para comparação. A AWS também testou separadamente o Qwen3-0.6B, mais compacto, mas a acceptance rate dele ficou cerca de 60% menor, o suficiente para perder a maior parte do ganho.

Na faixa de 5 a 15 speculative tokens, o melhor ponto nesses testes foi a configuração com sete tokens, embora os autores ressaltem que o valor ideal depende muito da estrutura dos prompts.

Foi justamente a estrutura da requisição que definiu o resultado. Em cenários previsíveis — texto repetitivo, sequências numéricas, código simples —, o speculative decoding trouxe benefício claro. Nesses casos, o modelo draft frequentemente acerta o que o target produziria de qualquer forma, e o sistema pula boa parte das etapas sequenciais.

Nos testes, a inter-token latency caiu para cerca de 15 milissegundos por token, e a curva de latência ponta a ponta ficou consistentemente abaixo do baseline. Em requisições abertas e menos determinísticas, o quadro é outro: o draft diverge do target com mais frequência, os tokens são rejeitados e o ganho potencial desaparece. Para esses prompts, a inter-token latency ficou em torno de 45 milissegundos por token, e as configurações speculative e baseline apresentaram latência total praticamente igual.

Já o TTFT, o tempo até o primeiro token, quase não mudou, porque o speculative decoding não acelera a etapa de prefill, em que o modelo codifica o contexto de entrada. O principal ganho aparece depois, na fase de decode, graças à redução do número de etapas sequenciais caras do modelo target.

A conclusão prática do artigo é simples: o speculative decoding no Trainium2 não é um botão universal de aceleração, e sim uma otimização pontual para um tipo específico de carga. Se o produto gera com frequência saídas estruturadas e previsíveis — código, extração de dados, relatórios baseados em templates, arquivos de configuração —, esse modo pode reduzir diretamente o custo por token de saída e aumentar o throughput sem perda de qualidade. Se o uso for principalmente chat aberto com geração livre, o efeito pode ser mínimo.

Por isso, vale adotar o esquema só depois de medir com os próprios prompts, escolhendo um modelo draft compatível e uma janela de speculative tokens ajustada aos cenários reais — e não com base em benchmarks desconectados do produto.