AWS mostra como especulação de decodificação no Trainium2 acelera geração no vLLM

A AWS mostrou uma forma prática de acelerar e reduzir o custo da inferência de LLM no Trainium2 para cenários onde o modelo gera significativamente mais tokens do que recebe como entrada. Trata-se de speculative decoding: em vez de forçar um modelo grande a produzir sequencialmente um token por vez, o sistema conecta um pequeno draft-model que rapidamente propõe vários tokens seguintes de uma vez, enquanto o main target-model os verifica em uma única passagem. Se as previsões coincidem, o serviço gasta menos etapas sequenciais caras, reduz a latência entre tokens e utiliza melhor o acelerador.

Isso é especialmente importante para cargas decode-heavy — assistentes de escrita, coding agents, geração de relatórios, documentos em template e outras tarefas com respostas longas. Na geração autorregressiva padrão, cada novo token é calculado separadamente, então o acelerador constantemente lê KV-cache da memória e realiza relativamente pouco trabalho útil por etapa. Por causa disso, a inferência frequentemente bate no limite de largura de banda da memória em vez do puro cálculo.

Speculative decoding aponta exatamente para esse gargalo: o target-model executa etapas de decode sequencial com menos frequência, e a verificação em lote torna a carga mais densa. No entanto, a abordagem tem requisitos. Os draft e target models devem usar o mesmo tokenizador e vocabulário, e idealmente pertencer à mesma família arquitetônica para que o modelo pequeno adivinhe com mais frequência a continuação do principal.

Um parâmetro-chave é o número de speculative tokens. Se a janela for muito pequena, o ganho é quase imperceptível; se muito grande, rejeições antecipadas e verificação desnecessária consomem o benefício. Em seu teste, a AWS usou o target-model Qwen3-32B e draft-model Qwen3-1.

7B, executados através do vLLM em uma instância trn2.48xlarge. Para speculative decoding, eles escolheram fused speculation no NeuronX Distributed Inference, onde ambos os modelos são compilados juntos para melhor desempenho.

As configurações baseline e speculative foram implantadas em um único cluster Amazon EKS com tudo mantido idêntico: alocação de acelerador, tensor parallelism, comprimento do contexto, batch limits e imagem Neuron. A única diferença foi a adição do draft-model e do parâmetro num_speculative_tokens. A carga foi aplicada a ambos os serviços via llmperf, e TTFT, inter-token latency e latência de ponta a ponta foram enviados ao CloudWatch para comparação.

A AWS também testou o mais compacto Qwen3-0.6B, mas sua taxa de aceitação foi aproximadamente 60 por cento menor, o que foi suficiente para perder a maior parte do benefício. Na faixa de 5 a 15 speculative tokens, o ponto ótimo nestes testes foi uma configuração com sete tokens, embora os autores enfatizem que o valor ótimo depende fortemente da estrutura do prompt.

Em última análise, a estrutura da requisição determinou o resultado. Em cenários previsíveis — texto repetido, sequências numéricas, código simples — speculative decoding mostrou benefícios claros. Em tais casos, o draft-model frequentemente adivinha o que o target-model produziria de qualquer forma, então o sistema pula uma porção significativa de etapas sequenciais.

Nos testes, inter-token latency caiu para aproximadamente 15 milissegundos por token, e a curva de latência de ponta a ponta se manteve consistentemente abaixo da baseline. Em requisições abertas e menos determinísticas, o quadro é diferente: o draft-model mais frequentemente diverge do target-model, tokens são rejeitados e o ganho potencial desaparece. Para esses prompts, inter-token latency flutuou em torno de 45 milissegundos por token, e as configurações speculative e baseline mostraram latência de ponta a ponta quase idêntica.

TTFT — tempo até o primeiro token — mudou pouco porque speculative decoding não acelera o estágio prefill, onde o modelo codifica o contexto de entrada. O principal benefício aparece depois, na fase de decode, reduzindo o número de etapas sequenciais caras do target-model. A conclusão prática do artigo é simples: speculative decoding no Trainium2 não é um botão de aceleração universal, mas uma otimização direcionada para um tipo de carga específico.

Se seu produto frequentemente gera saída estruturada e previsível — código, extração de dados, relatórios em template, configs — este modo pode reduzir diretamente o custo do token de saída e aumentar a taxa de transferência sem perda de qualidade. Se você tem principalmente chat aberto com geração de forma livre, o efeito pode ser mínimo. Portanto, implementar esse esquema é recomendado apenas após benchmarking em seus próprios prompts, selecionando um draft-model compatível e uma janela de speculative tokens adequada a cenários reais, em vez de confiar em benchmarks isolados do seu produto.