Google apresenta o TPU sync: ferramenta de código aberto para cache kv e transferência de pesos eficientes em tpus

A Google lançou o TPU Sync, uma biblioteca C++ de código aberto projetada para gerenciar o movimento de pesos e caches KV em várias camadas de memória dentro dos TPUs na nuvem. Esta biblioteca trata o HBM do TPU, DRAM do host, armazenamento compartilhado e a rede do centro de dados como uma hierarquia de memória unificada, abordando o crescente gargalo do movimento de estado do modelo, em vez de computação. Testes internos conduzidos pela Google em julho de 2026 no TPU v7x (2x2x1) demonstraram que o TPU Sync alcançou uma latência de transferência KV superior a 90% menor ao lidar com modelos MoE com parâmetros variando de 235B a 480B. O TPU Sync facilita duas funções principais: gerenciar o estado do cache KV durante a inferência e transmitir pesos de treinadores para amostradores em aprendizado por reforço. A arquitetura envolve um armazenamento KV global, um controlador de transferência e um gerenciador de cache KV, onde os nós de preenchimento descarregam blocos para a RAM do host fixada via DMA assíncrono antes de transmitir para os nós de decodificação do HBM. A biblioteca é licenciada sob a Apache-2.0 e integra-se com PyTorch via TorchTPU e nativamente com JAX. As integrações já estão estabelecidas com vLLM-TorchTPU como o conector KV do TPU e SGLang-TorchTPU como o backend de transferência P/D. Além disso, o llm-d 1.0 no GKE o utiliza para roteamento de afinidade de cache de prefixo e para despejo experimental no Google Cloud Managed Lustre. Meta e Wayfair estão relatadamente avaliando a arquitetura, enquanto Inferact e RadixArk co-engenharam as integrações vLLM e SGLang.