Google Cloud launches TPU Sync, cutting KV transfer latency by over 90%

[Google Cloud] Google Cloud introduced TPU Sync, an open-source native C++ library for moving model state across TPU clusters. According to Google Cloud, the library treats accelerator HBM, host DRAM, shared storage and the data center network as a single co-designed memory hierarchy, using asynchronous DMA transfers to overlap data movement with active computation. For inference, Google Cloud said TPU Sync turns host DRAM, local SSDs and Google Cloud Managed Lustre into an automated multi-tier KV cache; for reinforcement learning, it exposes low-level DMA and DCN primitives that higher-level frameworks can use to stream updated weights from trainers to samplers. The company said TPU Sync cuts KV transfer latency by more than 90% and reduces Time to First Token by an average of 90% for heavyweight open-weight MoE models of 235B to 480B parameters in prefix-caching and disaggregated-serving workloads, compared with running without it. TPU Sync integrates with PyTorch via TorchTPU and with JAX, and serves as the TPU KV cache connector for vLLM, SGLang and llm-d disaggregated serving, according to Google Cloud.