O Google Cloud lançou o TPU Sync, uma biblioteca C++ de código aberto para mover caches KV e pesos de modelo em Cloud TPUs. A biblioteca de transferência de cache KV Google TPU Sync trata a HBM da TPU, a DRAM do host, o armazenamento compartilhado e a rede do data center como uma única hierarquia de memória. Ela visa um gargalo crescente: mover o estado do modelo, e não computá-lo. O Google relata mais de 90% menos latência na transferência de KV em modelos MoE de 235B a 480B parâmetros. Resumo da Biblioteca: Núcleo C++ nativo com pacotes Python para PyTorch e JAX. Testado em modelos MoE de 235B a 480B. Executa em: Cloud TPU v7x (Ironwood) e v6e, conforme o README. Os benchmarks usaram TPU v7x (2x2x1), PCIe Gen5 e NICs de 200 Gbps. Desempenho: Grandes reduções na latência de transferência de KV e no tempo para o primeiro token (TTFT) nos testes de laboratório internos do Google. Melhor: Até 90% menos TTFT médio (serviço P/D de 480B com baixa concorrência e cache de prefixo de contexto de 64K). Pior: Nenhum número de RL publicado. O README afirma que ainda não é recomendado para uso geral. Conclusão: Melhor: A transferência de KV do Qwen3 235B cai de vários segundos para ~230 ms. Pior: todos os números são autorrelatados e estão em fase inicial. O que é TPU Sync? TPU Sync é uma biblioteca C++ de código aberto do Google que move tensores formatados para TPU entre camadas de memória com DMA assíncrona. Ela gerencia duas tarefas: o estado do cache KV durante a inferência e a transmissão de pesos do treinador para o amostrador em aprendizado por reforço (RL). O repositório é licenciado sob Apache-2.0. Como funciona o TPU Sync? A biblioteca abrange HBM de TPU, DRAM do host fixada, armazenamento compartilhado e DCN multi-host. O armazenamento utiliza backends POSIX experimentais, incluindo SSDs locais e Google Cloud Managed Lustre. O Google lista 4 escolhas de design: DMA direto com preservação de layout: Ele lê buffers PJRT do JAX ou tensores PyTorch no local. Os dados mantêm seu layout nativo de blocos de TPU, sem serialização Python ou buffers de salto. Transferências sobrepostas por computação: Aliases de buffer PJRT brutos permitem que as transferências ocorram durante a computação. O escalonamento de DMA integrado impede que grandes cópias atrasem o lançamento de kernels. Sem compilação XLA no caminho de dados: T
Fonte: MarkTechPost
Publicado em 2026-11-13