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 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) 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.
Understanding TPU Sync: The kitchen analogy
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).
Key use cases for TPU Sync
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.
Breaking through the memory wall with TPU Sync
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.
Higher-level architecture
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.
