# Introducing TPU Sync: Fast KV cache and weight transfer for inference and RL on TPU

> Source: <https://cloud.google.com/blog/topics/ai-infrastructure/introducing-tpu-sync-fast-kv-cache-and-weight-transfer-on-tpu/>
> Published: 2026-11-09 13:00:00+00:00

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<sup>1</sup> below on the testing configuration. 

TPU Sync serves as the open, composable memory substrate across our PyTorch-on-TPU and Kubernetes serving ecosystem, built in close collaboration with key open-source maintainers and infrastructure partners:

**Inferact (vLLM-TorchTPU)**: Co-engineered with Inferact, maintainers of [vLLM-TorchTPU](https://github.com/vllm-project/vllm-torchtpu), TPU Sync is integrated directly via vLLM's native KVConnector interface. It powers Direct-DMA host DRAM offloading and multi-host P/D disaggregation across frontier open-source models.

**RadixArk (SGLang-TorchTPU)**: In partnership with RadixArk, maintainers of [SGLang-TorchTPU](https://github.com/sgl-project/sglang-torchtpu), TPU Sync serves as both the high-speed P/D disaggregation transfer backend and the native I/O backend for SGLang’s hierarchical cache.

**llm-d 1.0 integration**: [llm-d](https://github.com/llm-d/llm-d) 1.0 pairs with TPU Sync on GKE to deliver global prefix-cache-affinity routing, per-rank decode session pinning, and experimental two-stage capacity spill to Google Cloud Managed Lustre.

**Validated by frontier AI labs and enterprises**: Frontier AI organizations are already adopting our PyTorch-on-TPU serving stack and evaluating TPU Sync multi-tier architectures for production deployments — including Meta, alongside major enterprises like Wayfair, and multi-silicon partners.

[TPU Sync](https://github.com/google/tpu-sync) is available today as open source with Python wheels for both PyTorch (`tpu-sync-torch`) and JAX (` tpu-sync-jax`), and is integrated with [vLLM-TorchTPU](https://github.com/vllm-project/vllm-torchtpu), [SGLang-TorchTPU](https://github.com/sgl-project/sglang-torchtpu), and [llm-d](https://github.com/llm-d/llm-d) on GKE. Visit the [GitHub repository](https://github.com/google/tpu-sync) to explore the architecture docs, C++ and Python APIs, and deployment guides.

We are committed to working with the open-source community to advance the model-state layer on TPU — supporting new frontier architectures, expanded storage tiering, and longer-horizon agentic and RL workloads. For more up-to-date performance results, please see our ongoing efforts on public benchmarks ([InferenceX](https://inferencex.semianalysis.com/)).

*1. Performance results are based on internal testing conducted in July 2026 on TPU v7x (2x2x1 topology), configured with PCIe Gen5 (x16) per TPU chip and one or two 200 Gbps NICs enabled per host, running tpu-inference v0.24.0 (using JAX 0.10.2 and libtpu 0.0.42.1). Metrics presented represent average and median performance improvements observed under controlled laboratory conditions. Actual real-world performance, cache hit rates, and latency reductions will vary based on system configuration, network bandwidth, concurrent traffic, and specific user workload patterns.*