Nova técnica promete acelerar inferência de LLMs em TPUs do Google

Pesquisadores da UCSD implementaram uma técnica de decodificação especulativa em estilo de difusão em TPUs do Google e alcançaram aumento médio de 3,13 vezes em tokens por segundo no TPU v5p. O resultado foi destacado no Google Developers Blog como um marco open-source para aceleração de inferência de grandes modelos de linguagem.

A implementação usa DFlash, uma abordagem baseada em block diffusion. Em vez de prever tokens candidatos um a um, como em métodos autoregressivos tradicionais, a técnica gera um bloco inteiro de tokens candidatos em uma única passada.

O problema da inferência autoregressiva

Modelos de linguagem tradicionais geram texto de forma autoregressiva. Isso significa que cada token depende do token anterior, exigindo uma nova etapa de processamento para cada avanço da resposta.

Esse modelo funciona, mas cria gargalos de latência. Mesmo quando aceleradores como TPUs têm grande capacidade de computação paralela, a geração token por token pode subutilizar o hardware.

A decodificação especulativa tenta resolver parte desse problema usando um mecanismo menor para propor tokens candidatos. O modelo principal verifica esses tokens em paralelo e aceita vários deles quando estão corretos.

O que o DFlash muda

O DFlash muda a etapa de rascunho. Em métodos especulativos tradicionais, o modelo que propõe tokens também pode operar de forma sequencial, o que mantém um gargalo.

Com block diffusion, o DFlash propõe vários tokens de uma vez. A ideia é reduzir a complexidade da fase de rascunho e aproveitar melhor a capacidade paralela das TPUs.

Segundo o Google Developers Blog, a implementação foi integrada ao ecossistema open-source vLLM TPU Inference. O trabalho exigiu adaptações em cache, gerenciamento de contexto e sincronização de metadados para que a técnica funcionasse de forma eficiente em JAX e TPUs.

Resultados dos testes

Nos testes em TPU v5p, a equipe da UCSD observou aumento médio de 3,13 vezes em tokens por segundo, com picos próximos de 6 vezes em tarefas complexas de matemática.

Em comparação direta com EAGLE-3, outro método de decodificação especulativa, o DFlash alcançou speedup end-to-end de 2,29 vezes, contra 1,30 vez do EAGLE-3.

Em tarefas de programação como MBPP, o DFlash reduziu o tempo de geração de 9,81 milissegundos por token para 3,48 milissegundos por token. Em avaliações como HumanEval, as velocidades de geração melhoraram em mais de 3,5 vezes.

Por que isso importa

A aceleração de inferência é um dos principais desafios econômicos da IA generativa. Modelos maiores e mais capazes custam caro para servir em escala, especialmente quando precisam responder a milhões de usuários com baixa latência.

Técnicas como DFlash apontam para uma direção diferente da simples expansão de hardware. Em vez de depender apenas de mais chips, a indústria também busca formas de usar melhor o hardware disponível.

Isso tem impacto direto em produtos de IA, custos de API, agentes autônomos, assistentes corporativos e sistemas que precisam gerar respostas longas ou executar tarefas complexas.

Análise

A notícia reforça que a corrida de IA não está apenas no treinamento de modelos. A fase de inferência, onde os modelos realmente atendem usuários, virou campo central de disputa.

Reduzir latência e aumentar throughput pode melhorar experiência, reduzir custos e permitir novas aplicações. Para empresas que operam modelos em escala, ganhos de 2 vezes ou 3 vezes não são apenas melhorias técnicas. Eles podem alterar a viabilidade econômica de produtos inteiros.

O avanço do DFlash em TPUs também fortalece o ecossistema open-source de serving de IA. À medida que modelos, aceleradores e frameworks se tornam mais especializados, a capacidade de otimizar a pilha completa será uma vantagem estratégica.