Introducing TPU Sync: Fast KV cache and weight transfer for inference and RL on TPU Google released TPU Sync, an open-source native C++ library that moves TPU-tiled tensors directly across HBM, host DRAM, shared storage and the data center network, cutting KV cache transfer latency by over 90% and reducing Time to First Token by an average of 90% for 235B–480B-parameter open-weight MoE models in prefix-caching and disaggregated-serving workloads versus running without it. TPU Sync integrates with PyTorch via TorchTPU and with JAX, serving as the KV cache connector for vLLM, SGLang and llm-d disaggregated serving while providing low-level DMA and DCN primitives for Pathways and multi-controller JAX (McJAX) RL pipelines. The library addresses the shift from stateless single-turn prompts to million-token agents and continuous RL, where orchestrating live KV and parameter state, not raw compute, has become the primary infrastructure bottleneck. As frontier AI evolves from stateless single-turn prompts to million-token agents, disaggregated serving, and continuous reinforcement learning RL , the primary bottleneck in AI infrastructure has shifted from raw compute to orchestrating live model state — both the contextual state held in massive key-value KV caches during inference and the parameter state updated and broadcast from trainers to rollout samplers during RL. Whenever this state is trapped inside a single chip's memory or stalled on network transfers, AI accelerators sit idle. We designed TPU Sync https://github.com/google/tpu-sync from first principles to serve as the model-state layer across TPU clusters — built on the architectural insight that accelerator High Bandwidth Memory HBM , host DRAM, shared storage, and the data center network DCN should operate as a single, co-designed memory hierarchy. An open-source, native C++ library, TPU Sync moves TPU-tiled tensors directly across this hierarchy with asynchronous DMA transfers that overlap data movement with active computation. For inference, it decouples TPU compute from network backpressure and turns host DRAM, local SSDs, and Google Cloud Managed Lustre into an automated, multi-tier KV cache; and for RL, it exposes low-level DMA and DCN primitives that higher-level RL frameworks can use to stream updated weights from trainers to samplers. By accessing device buffers in place rather than copying or locking them, TPU Sync consistently slashes KV transfer latency by over 90% and reduces Time to First Token TTFT by an average of 90% for heavyweight open-weight MoE models 235B–480B parameters in prefix-caching and disaggregated-serving workloads, compared to running without TPU Sync. Across this shared model-state layer, TPU Sync integrates directly into existing infrastructure stacks across PyTorch via TorchTPU https://github.com/google-pytorch/torch tpu and JAX — serving as the TPU KV cache connector for vLLM, SGLang, and llm-d disaggregated serving, while providing low-level transfer primitives for Pathways and multi-controller JAX McJAX RL pipelines. To understand TPU Sync’s impact, imagine a high-end restaurant kitchen. Think of the TPU as an incredibly fast executive chef working on a tiny countertop the chip's HBM . When preparing a massive feast, the chef needs lots of ingredients the KV cache , which quickly clutter the counter. The standard kitchen without TPU Sync : When the small counter fills up, the chef has to toss out prepped ingredients and chop them again from scratch the next time an order needs them — stalling the kitchen on repetitive prep work prefill recomputation . The upgraded kitchen with TPU Sync : TPU Sync adds a spacious side pantry pinned host RAM , a shared walk-in cold storage room accessible to every station in the kitchen shared cluster storage , and a coordinated team of runners. The runners slide full prep trays directly between the counter, the pantry, and peer stations without repacking them direct DMA in native TPU layout . Crucially, the runners swap trays in and out while the chef keeps cooking asynchronous, non-blocking transfers . TPU Sync improves system throughput and latency in workloads requiring deep context, high concurrency, multi-turn reasoning and RL rollouts: Long-context AI agents and multi-turn chats : When interacting with an AI agent over an extended 256K to 1M+ token session or analyzing massive documents, the model generates a continuously growing KV cache. Instead of allowing this state to exhaust HBM capacity, TPU Sync automatically tiers older KV blocks out of TPU HBM and into large pools of host RAM with experimental tiering to shared secondary storage . Enterprise applications with heavy system prompts prefix caching : AI serving platforms receive thousands of queries that share common instructions or documents. Rather than recomputing prefixes that the model has already processed, TPU Sync extends the prefix cache into distributed host memory so concurrent sessions can reuse precomputed KV blocks and bypass redundant prefill computation. Disaggregated serving architecture prefill/decode : In modern enterprise deployments, workloads are isolated: Compute-bound nodes process the initial massive prompt prefill , while memory-bound nodes draft the response token-by-token decode . TPU Sync bridges this hardware gap by acting as a high-bandwidth Host-to-Host H2H streaming pipeline. It routes the processed context directly from the prefill nodes to the decode nodes, consistently delivering an over 90% reduction in data transfer overhead. Reinforcement learning pipelines experimental : Validated on internal RL workloads and included in the repository as an experimental module, TPU Sync provides low-level DMA and DCN transfer primitives that higher-level RL frameworks can use to synchronize model weights between trainers and inference samplers — including automatic resharding across different trainer and sampler parallelism layouts — to reduce the time rollout workers spend waiting on weight updates. At its core, TPU Sync establishes a high-performance data transfer path between TPU HBM and host RAM, as well as across TPU hosts via the DCN. It achieves this through four key architectural advantages: Layout-preserving direct DMA with in-place buffer aliasing : TPU Sync accesses the underlying JAX PJRT buffers or PyTorch tensors in place and moves slices directly between TPU HBM and page-locked DmaMap -registered host memory in a single hardware DMA. By keeping tensors in their native TPU physical tile layout across HBM, host DRAM, storage, and the wire, data lands ready for immediate Matrix Multiply Unit MXU execution with no Python serialization, no layout conversion, and no user-space bounce buffers. Compute-overlapped transfers : Rather than taking an exclusive lock or making an on-device copy of static KV cache and weight buffers, TPU Sync creates raw PJRT buffer aliases PJRT RawBuffer without holding PJRT usage locks so block-level transfers can run concurrently with active in-place computation, with built-in DMA pacing to keep large transfers from delaying kernel launches. No XLA compilation in the data path : Because TPU Sync executes transfers and tile-aware strided resharding in native C++ outside the XLA graph, moving KV blocks or resharding weights across different trainer and sampler topologies does not trigger runtime compilation — keeping the data path fast and framework-agnostic. High-bandwidth PCIe DMA and DCN streaming : Using the raw transfer library’s direct PJRT raw-buffer DMA pipeline, TPU Sync reaches up to 223 GB/s of aggregate D2H and H2D throughput per Ironwood TPU v7x host out of the 256 GB/s PCIe Gen5 hardware ceiling, and streams across the DCN at up to ~192 Gbps on a single 200 Gbps NIC or ~350 Gbps across dual 200 Gbps NICs with NUMA-aware multi-NIC striping . TPU Sync exposes a low-level data transfer API, making it easier for developers to integrate high-speed TPU data transfer into higher-level frameworks. Built on top of its core transport and resharding primitives, TPU Sync provides a modular orchestration stack for distributed cache management: Global KV store: A global registry that maps prompt prefix hashes to physical cache blocks, handles dynamic lookups, and coordinates LRU eviction and tier demotion. TPU Sync controller: The orchestration layer that generates transfer plans, coordinates data movement across nodes, and handles lease management and recovery. KV cache manager: The local execution engine that tracks block states, manages HBM, host DRAM including /dev/shm warm-boot persistence alongside experimental POSIX storage backends, and executes asynchronous DMA transfers. Let us take a closer look at how TPU Sync helps with some of the use cases we listed in the earlier section. By treating host pinned memory as a substantial, high-speed Tier 2 cache, TPU Sync reduces the fragmentation and memory limits of standard HBM setups. In state-of-the-art serving topologies, prefill compute-bound and decode memory-bound phases are isolated on separate node pools. To prevent latency bubbles, TPU Sync bridges this gap with a highly optimized, three-step pipeline: Direct-DMA offloading: Once a prefill node processes an initial prompt, TPU Sync immediately moves the resulting KV blocks into pinned host RAM via its asynchronous D2H DMA path. Host-to-host streaming : The TPU Sync controller then orchestrates a high-bandwidth Host-to-Host H2H transfer, streaming the context directly to the assigned decode node. Overlapped ingestion : As blocks arrive in the decode node's host RAM, TPU Sync triggers an asynchronous H2D transfer directly into TPU HBM. This happens concurrently with active decoding computation, substantially minimizing TTFT. For highly concurrent systems sharing large system prompts or running multi-turn agentic loops, TPU HBM is not large enough to hold multiple active cache blocks. TPU Sync extends the prefix cache out of HBM Tier 1 into substantial, distributed host RAM pools Tier 2 . Because TPU Sync transfers are asynchronous, it overlaps the PCIe DMA loading of the next required prefix from host RAM with the compute of the current request — minimizing TPU idle time between requests. For workloads that exceed even host DRAM capacity, TPU Sync also includes an experimental POSIX storage backend Tier 3+ that asynchronously spills older KV blocks to local SSDs or shared Google Cloud Managed Lustre storage — pairing with llm-d prefix-affinity routing on GKE so replicas can recall cached prefixes across tiers without recomputing prompts. P/D disaggregated inference result highlights: Over 90% KV transfer overhead cut : For heavyweight models like Qwen3 235B, median KV transfer latency dropped from several seconds down to a flat ~230 ms. Substantial TTFT reduction : Demonstrated an average of 60% reduction in TTFT for heavyweight models Qwen3 235B under controlled, low-concurrency testing benchmarks, and an average of 90% for 480B under low concurrency. Substantial high-traffic scalability : During sustained traffic spikes, TPU Sync achieves an average of 30% to 70% increase in token output throughput for 235B, depending on traffic distribution. Prefix caching result highlights: Expands KV cache capacity beyond HBM : When heavyweight MoE model weights consume most of device HBM, on-chip memory can hold only a few long-context prefixes. Offloading into pinned host DRAM expands effective cache capacity by an order of magnitude. Sustains high cache-hit rates under concurrency : Pooling distributed host RAM keeps shared system prompts and multi-turn agent contexts warm across thousands of concurrent sessions instead of thrashing HBM. Significantly reduces prefill recompute : Directly bypasses heavy compute layers for repeated instruction blocks and system prompts, reducing average TTFT by up to 90% on 64K-context workloads while delivering over 50% higher throughput and a 30% reduction in P99 inter-token latency ITL . Please see the note