Introducing TPU Sync: Fast KV cache and weight transfer for inference and RL on TPU
Google introduces TPU Sync, a new open-source C++ library designed to accelerate AI inference and reinforcement learning on TPU clusters by optimizing model-state transfers across memory hierarchies and networks.
TPU Sync addresses the growing challenge of managing live model state—key-value caches for inference and parameter updates for reinforcement learning—by treating TPU HBM, host DRAM, shared storage, and the data center network as a unified memory hierarchy. The open-source C++ library enables direct, asynchronous DMA transfers of TPU-tiled tensors, allowing computation to overlap with data movement. This design reduces idle time for TPU accelerators by decoupling compute from network backpressure and enabling multi-tier KV caching across host memory and storage.
For inference workloads, TPU Sync transforms host DRAM, local SSDs, and Google Cloud Managed Lustre into an automated, multi-tier KV cache, eliminating redundant prefill computations in prefix-caching and disaggregated serving scenarios. In reinforcement learning, it provides low-level DMA and DCN primitives for streaming updated model weights from trainers to samplers, including automatic resharding across parallelism layouts. The library integrates with existing frameworks like PyTorch (via TorchTPU) and JAX, supporting tools such as vLLM, SGLang, llm-d, Pathways, and McJAX.
The system achieves over 90% reductions in KV transfer latency and Time to First Token for large open-weight MoE models (235B–480B parameters) in prefix-caching and disaggregated serving, compared to systems without TPU Sync. It operates by preserving tensor layouts during transfers, avoiding Python serialization or layout conversions, and enabling compute-overlapped DMA operations. TPU Sync’s architecture supports high-bandwidth PCIe DMA and DCN streaming, reaching up to 223 GB/s per host and ~192 Gbps on a single NIC.
TPU Sync introduces a modular orchestration stack for distributed cache management, including a global KV store for prompt prefix hashes, a TPU Sync controller for transfer planning, and a KV cache manager for local execution. The system handles dynamic lookups, LRU eviction, tier demotion, and asynchronous DMA transfers across HBM, host DRAM, and experimental storage backends. It aims to improve throughput and latency in long-context AI agents, multi-turn chats, enterprise prefix caching, disaggregated serving, and RL rollouts.