Google libera código do TPU Sync para sincronizar cache e pesos entre TPUs

O Google Cloud↳Google Cloud13 conteúdosPrograma da Google Cloud no Brasil projeta formar de forma gratuita 30 mil universitáriosDev (Back & Front) · mar 2026Google Cloud OnBoard capacita estudantes e desenvolvedores de TIGestão Dev & TI · mai 2019Desenvolvedores poderão participar de treinamento gratuito do Google CloudGestão Dev & TI · mai 2019Ver tudo em DevSecOps → abriu, em 12 de novembro de 2026, o código do TPU Sync, uma biblioteca nativa em C++↳C++6 conteúdosSão Paulo recebe encontro de Programadores C & C++Dev (Back & Front) · mai 2019Melhores linguagens de programação para desenvolver jogosMarketing Tech · nov 2024Desmistificando Rust: a linguagem segura e rápida que você precisa conhecerDev (Back & Front) · out 2024Ver tudo em Dev (Back & Front) → que move blocos de KV cache e pesos de modelo entre a memória HBM das TPUs, a DRAM do host, armazenamento compartilhado e a rede do data center. O anúncio veio junto com os novos sistemas TPU 8tx e TPU 8i Max, mas o foco da ferramenta não é o silício: é resolver o gargalo de dados que faz aceleradores caros ficarem ociosos esperando cache ou pesos chegarem, segundo reportagem da InfoQ.

O projeto mira dois cenários específicos. O primeiro é a inferência de contexto longo com prefill e decode desacoplados (disaggregated serving), onde o KV cache precisa viajar entre nós diferentes do cluster. O segundo é o loop de aprendizado por reforço (RL), em que pesos atualizados precisam sair do treinador e chegar rápido aos workers que geram as respostas (rollout samplers).

Os números que o Google divulgou

Para a disaggregação de prefill e decode, o Google relata que a latência mediana de transferência de KV cache do modelo Qwen3 235B caiu de vários segundos para cerca de 230 ms. Isso se traduziu em reduções médias de 60% no TTFT (time to first token, o tempo até o primeiro token de resposta) para esse modelo, e de 90% para um modelo de 480B em baixa concorrência. Para prefix caching em cargas com janela de 64 mil tokens, a empresa reporta quedas de até 90% no TTFT médio.

Esses ganhos de latência não vêm só de banda maior. O desenvolvedor pseudônimo @lgoodbookd fez as contas públicas sobre o novo TPU 8i Max:

O 8i Max dobra a banda de HBM para cerca de 17,2 TB/s, mas uma única sessão de 64K em um modelo MoE de 235B ainda consome cerca de 12,6 GB de KV cache. É por isso que abrir o código do TPU Sync (com tiering de KV em DRAM do host, a cerca de 223 GB/s por host v7x) pode importar tanto quanto o novo hardware. Banda ajuda; capacidade ainda é o muro.

@lgoodbookd, desenvolvedor8i Max doubles HBM bandwidth to ~17.2TB/s, yet one 64K session on a 235B MoE still eats ~12.6GB of KV cache. That's why open-sourcing TPU Sync (host-DRAM KV tiering, ~223GB/s per v7x host) may matter as much as the new package. Bandwidth helps; capacity is still the wall.

Em outras palavras: o problema não é só mover dados mais rápido entre chips, é ter para onde descarregar o que não cabe na HBM sem perder velocidade. É aí que entra o tiering de KV cache para a DRAM do host que o TPU Sync implementa.

O que o desenvolvedor ganha na prática

A biblioteca sai com wheels Python para JAX e PyTorch, expondo três conjuntos de APIs: gerenciador de KV cache, armazenamento de KV (KV store) e sincronizador de pesos. Na prática, ela já se conecta a projetos que boa parte de quem roda inferência em escala conhece:

• Atua como conector de KV para o vLLM-TorchTPU, através da interface KVConnector do próprio vLLM.
• Funciona como backend de disaggregação e cache hierárquico para o SGLang-TorchTPU.
• Integra com o llm-d 1.0 rodando no GKE para roteamento consciente de prefix cache.
• Tem um módulo experimental de sincronização de pesos para RL, voltado a Pathways e JAX multi-controller.

O repositório no GitHub absorveu o código do projeto anterior, chamado TPU Raiden, e está validado para TPU v7x e v6e. Quem for testar precisa saber que o wheel de PyTorch exige uma build específica de torch_tpu que não está disponível em índice público, e que só builds de desenvolvimento estão publicadas no PyPI: o próprio projeto recomenda fixar (pin) no último commit conhecido como estável, já que as APIs ainda mudam rápido.

Como isso se compara ao que já existe em GPU

Quem trabalha com inferência em GPU já lida com variações do mesmo problema, só que com outras ferramentas. O NVIDIA Dynamo e sua biblioteca de transferência NIXL atacam a disaggregação de prefill e decode nesse ecossistema. O LMCache faz offload de KV cache para memória de CPU, disco ou armazenamento remoto. O Mooncake é a arquitetura de serving centrada em KVCache por trás do Kimi, da Moonshot AI↳Inteligência artificial440 conteúdosUX e IA: Transformando Experiências Digitais com Inteligência ArtificialProduto & UX · jan 2025MCP: O que é e por que você vai ouvir falar disso em breve?AI · jul 2025IA generativa e a urgência de reconstruir nossa relação com a verdadeAI · jun 2025Ver tudo em AI →. E o próprio vLLM tem seus conectores nativos de prefill desacoplado.

A diferença do TPU Sync para esse grupo, segundo a InfoQ, está em dois pontos: ele é construído em cima do PJRT e dos layouts de tile específicos de TPU, e cobre no mesmo pacote a sincronização de pesos para RL, algo que as ferramentas de GPU citadas tratam separadamente. Para quem já roda uma stack híbrida GPU e TPU, isso significa decidir se vale manter duas camadas de transferência de dados ou migrar as cargas de TPU para a nova biblioteca.

A leitura de quem já está testando

A reação inicial de parte da comunidade técnica foi menos sobre o hardware novo e mais sobre o que a abertura do código resolve de fato. O desenvolvedor Anand Butani resumiu assim:

O TPU Sync cortar 90% da latência de transferência de KV importa mais para mim do que os nomes dos chips; serving de contexto longo basicamente morre na hora de mover cache entre nós.

Anand Butani, desenvolvedorTPU sync cutting kv transfer latency 90% matters more to me than the chip names, long-context serving mostly dies on moving cache between nodes.

Esse tipo de comentário aponta para quem o TPU Sync realmente serve: times que já sentem na pele o custo de manter contexto longo vivo em produção, seja em chatbots com memória extensa, seja em pipelines de RAG que reusam prefixos grandes entre chamadas.

O que ainda falta amadurecer

Nem tudo está pronto para produção. O próprio repositório avisa que as APIs mudam rápido, que apenas builds de desenvolvimento estão no PyPI, e que a camada de armazenamento em disco e o módulo de sincronização de pesos para RL seguem experimentais. Para times que avaliam adotar agora, isso significa tratar o TPU Sync como peça de laboratório, não como dependência de produção, até que um release estável seja cortado.

O código está sob licença Apache 2.0 e o repositório traz documentação de arquitetura, APIs em C++ e Python, e exemplos de uso tanto em host único quanto em serving desacoplado multi-host, além de exemplos de offloading de KV para o host. Para quem roda inferência de LLMs grandes em TPU na GCP, é o primeiro sinal público de que o Google está tratando a movimentação de dados entre aceleradores como um problema separado do hardware, e disposto a abrir essa camada para a comunidade testar e contribuir.

Fonte: InfoQ