Voltar ao blog

DFlash em TPUs: quando a decodificação especulativa ganha um proponente por difusão

IA LLM Inferência MLOps Engenharia de IA

Gerar tokens com um LLM é um processo serial: para produzir K tokens, o modelo executa K passos sequenciais de decodificação. A decodificação especulativa reduz esse gargalo propondo vários tokens com um mecanismo mais barato e usando o modelo principal apenas como verificador, sem alterar a distribuição final quando o algoritmo de aceitação é implementado corretamente. Até agora, os proponentes eram autoregressivos — modelos menores que compartilham o mesmo vocabulário e predizem o próximo token um por vez. O DFlash substitui esse proponente por um modelo de difusão em blocos, que gera múltiplos tokens em paralelo por refinamento iterativo.

Em 4 de maio de 2026, o Google Developers publicou os resultados da integração do DFlash em TPU v5p, feita em colaboração com a UCSD. O Google relata aumento médio de 3,13x em tokens por segundo, com picos próximos de 6x em tarefas matemáticas complexas. Em uma comparação direta com EAGLE-3 no mesmo hardware e com o mesmo modelo alvo, Llama-3.1-8B, o DFlash atingiu 2,29x de speedup ponta a ponta, contra 1,30x do EAGLE-3. Em tarefas de código como MBPP, a latência caiu de 9,81 ms por token para 3,48 ms por token, uma melhoria de 2,83x.

Esse ganho exigiu mais do que portar um modelo. A difusão em blocos do DFlash é não causal, enquanto a inferência de alto desempenho em TPUs usa atenção paginada e cache KV organizado para execução eficiente. Esses dois modelos de memória não encaixam naturalmente. A solução descrita foi uma arquitetura de cache duplo: o modelo alvo continua usando KV cache paginado com kernels Pallas, enquanto o proponente usa um caminho especializado com arrays JAX estáticos no dispositivo. Além disso, como o DFlash é condicionado por estados ocultos do modelo alvo, a passagem de verificação precisa exportar esses estados para o proponente em cada iteração. Esse acoplamento entre os dois modelos — um autoregressivo, outro baseado em difusão — é a contribuição de engenharia que sustenta os números reportados.

A análise dos limites de batch também traz implicações práticas. A equipe observou que, em TPU v5p, verificar 1.024 tokens pode ter custo muito próximo ao de verificar apenas 16 tokens, porque o tempo é dominado pelo carregamento dos pesos do modelo, não pela aritmética incremental de atenção nesses comprimentos. Isso desloca a fronteira de pesquisa: aumentar K indefinidamente tem retorno decrescente. Segundo a análise reportada, K=16 já captura mais de 90% do speedup teórico no ponto operacional testado, enquanto melhorar a probabilidade de aceitação por posição é de 2 a 3 vezes mais valioso do que apenas ampliar o bloco. O gargalo deixa de ser “quantos tokens verificar” e passa a ser “quão bons são os tokens propostos”.

Essa distinção importa para aplicações reais porque o ganho não é uniforme. Matemática e código tendem a produzir sequências mais previsíveis, com menor decaimento de aceitação ao longo do bloco. Conversas abertas têm maior variabilidade lexical e semântica, o que reduz a quantidade de tokens aceitos antes de uma rejeição. Por isso, o DFlash é mais interessante para cargas de trabalho com estrutura: raciocínio matemático, geração de código, agentes que executam passos repetíveis e respostas longas com trajetória relativamente determinada. Em chat genérico, o benefício existe, mas a expectativa técnica deve ser mais conservadora.

Há trade-offs claros. O sistema exige um proponente treinado ou ajustado para o modelo alvo, aumenta a complexidade do serving e introduz novos estados que precisam permanecer consistentes entre iterações. A lista de modelos suportados no repositório do DFlash inclui variantes como Qwen3.6-35B-A3B, Qwen3.6-27B, Qwen3-Coder-Next e Llama-3.1-8B-Instruct, o que indica um ecossistema inicial, mas ainda dependente de checkpoints específicos. O ganho operacional também depende do hardware, do framework, do tamanho do bloco, da tarefa e da taxa de aceitação. A promessa não é transformar qualquer modelo em um sistema 6x mais rápido; é abrir uma camada de otimização que antes estava limitada por proponentes autoregressivos.

A implicação técnica é que escolher um LLM para produção fica menos parecido com escolher apenas um endpoint e mais parecido com projetar um caminho de execução. Modelo alvo, proponente, cache KV, runtime, kernels, hardware e distribuição das tarefas passam a formar um conjunto único. Em sistemas de desenvolvimento de software com IA, isso é especialmente relevante porque a latência de cada iteração afeta ciclos de teste, inspeção de erro, geração de patch e validação. A redução de alguns milissegundos por token se acumula quando um agente produz centenas ou milhares de tokens em um loop de trabalho. A métrica que passa a importar não é apenas tokens por segundo, mas tokens aceitos por passagem cara do modelo alvo.

O DFlash em TPUs aponta para uma direção concreta: parte do avanço em IA virá de modelos melhores, mas outra parte virá de sistemas que extraem mais trabalho útil de cada passagem do modelo. Nesse regime, inferência deixa de ser etapa passiva depois do treinamento. Ela vira uma camada de arquitetura, com algoritmos próprios, invariantes de estado e decisões de hardware que afetam diretamente custo, latência e viabilidade de produto. Para engenharia de IA, esse é um deslocamento relevante: a fronteira prática não está só no próximo modelo maior, mas em como cada token é produzido, proposto, verificado e aceito.

Fontes: