TPU Sync Explained: KV Cache, Memory Tiers and Transfers
Trace TPU Sync across HBM, host DRAM and remote nodes; understand its reported transfer gains, experimental storage and RL paths, and how it differs from Agent Compute Engine.
BY JONAS WEBER · AIINFRAINTERVIEWS EDITORIAL · UPDATED NOVEMBER 9, 2026 · 5 MIN READ
PRACTICE THIS:GPU and accelerator questions ·Inference and serving questions ·Calculate KV-cache memory
TPU Sync moves model state between TPU memory, host memory and remote workers. Its main inference uses are retaining reusable key/value (KV) cache outside scarce high-bandwidth memory and transferring that state from a prefill worker to a decode worker. The public project is google/tpu-sync, with integrations described for vLLM-TorchTPU, SGLang-TorchTPU and llm-d.
Google's launch article reports large gains in selected tests. Those results need their workload and baseline attached. This November 10, 2026 review explains the mechanism and its limits; it does not reproduce a TPU benchmark.
What state moves, and why?
The KV cache stores attention keys and values for processed tokens. A later generation step reuses them. When two requests share a compatible prefix, saved cache blocks can also avoid repeating work on that prefix. The free KV-cache concept explains this token-level lifecycle.
Device HBM is fast but limited. Host DRAM is larger, but a block must cross the device-host connection before it can be used on the accelerator again. Remote memory adds a network transfer. A storage tier adds another access path. Each tier trades capacity against retrieval cost.
Moving inactive reusable blocks out of HBM can preserve more prefixes. It does not make host memory equivalent to local HBM, and it does not remove the need to bring required data back before dependent computation can proceed.
Step 1 of 4: Prefill TPU to Prefill host DRAM
- Prefill TPUPrefill host DRAM
A completed block moves out through device-to-host DMA.
- Prefill hostDecode host
The data center network carries the block to the selected worker.
- Decode host DRAMDecode TPU
Host-to-device DMA restores the block in the required device layout.
- Transfer completionDecode scheduler
Dependent work can use the complete state; unrelated work may overlap earlier transfers.
The slowest path still has a byte budget
Suppose an illustrative 10 GiB payload crosses a 200 Gbit/s link. Ignoring protocol overhead and contention:
Payload = 10 × 2^30 = 10,737,418,240 bytes
Link ceiling = 200 × 10^9 / 8 = 25,000,000,000 bytes/s
Network-only floor = payload / rate ≈ 0.4295 seconds
Actual transfer time includes effective bandwidth, queueing and any non-overlapped staging work. A pipelined transfer can overlap stages, so adding each stage's full duration may overestimate the critical path. Dividing by the sum of all link rates would underestimate it. Measure the complete interval from state-ready at the producer to state-usable at the consumer.
Google describes direct buffer access and layout-preserving transfers that avoid extra conversion and copies. “Zero-copy” in that context does not mean zero bytes traverse PCIe or the network. The TPU Sync repository also recommends a tagged or known-good revision because its APIs are evolving.
Chunking changes the transfer critical path
The three-stage path can operate as a pipeline if the implementation and buffers allow different chunks to advance independently. Suppose a 10 GiB payload is divided into four equal chunks. Assume each chunk takes 50 ms for device-to-host transfer, 110 ms across the network and 40 ms for host-to-device transfer. These hypothetical rates are effective stage rates, not TPU Sync benchmark results.
| Chunk | Device to host | Network | Host to device |
|---|---|---|---|
| 1 | 0–50 ms | 50–160 ms | 160–200 ms |
| 2 | 50–100 ms | 160–270 ms | 270–310 ms |
| 3 | 100–150 ms | 270–380 ms | 380–420 ms |
| 4 | 150–200 ms | 380–490 ms | 490–530 ms |
Serializing all three stages for all four chunks takes 4 × (50 + 110 + 40) = 800 ms. The ideal pipeline takes the first chunk's 200 ms plus three more 110 ms network intervals, totaling 530 ms. The network stage sets the sustained pace. The first chunk still has to traverse every stage, and the last chunk still has to drain into device memory.
This schedule assumes the stages do not compete for another shared bottleneck and that the receiver can accept the next completed chunk. Limited staging buffers, host-memory bandwidth or contention can create gaps. The 440 ms of network work remains above the earlier 429.5 ms link-only floor, so the example does not create bandwidth by overlapping transfers.
A usable chunk is not necessarily a usable request. Decode can begin only when the state its next computation depends on is complete and correctly laid out. If cancellation arrives while a transfer is running, the integration must also prevent that completion from writing into a buffer reassigned to another request. Check ownership and completion handling in the exact serving integration; asynchronous transfer alone does not establish those guarantees.
When reloading beats recomputation
Imagine a compatible prefix that takes 800 ms to recompute. If lookup, transfer and installation take 250 ms under the current load, reloading saves 550 ms on that path. If network contention increases the reload to 900 ms, recomputing is faster. A cache policy needs a fallback for that change.
The matching key must include the information that determines the state: model and adapter revision, token sequence, positional treatment and cache representation. Reusing a block from a different model can produce a fluent but wrong result. A transfer completing successfully does not prove semantic compatibility.
The Premium KV Cache Offloading and Tiered KV Storage reference develops the admission and eviction decision. For the system-level interview application, use inference-serving questions.
Read the 90 percent claim with its denominator
The launch reports transfer-overhead reductions exceeding 90% and different first-token improvements for selected workloads. Its test note identifies internal testing on , a 2×2×1 topology, 200 Gbit/s NIC configurations and a JAX-based software stack. It does not establish the same result on TPU 8 or on every TorchTPU configuration.
For a simple example, suppose a request takes 500 ms before the first token, of which 100 ms is transfer. Cutting transfer by 90% saves 90 ms and leaves 410 ms total. The end-to-end improvement is 18%. A 90% transfer reduction and a 90% request-latency reduction are different claims.
Compare cache-hit and cache-miss traffic separately, then combine them using the expected traffic mix. Also record token gaps for requests already decoding: a faster cache warmup is not a win if its transfers cause unacceptable stalls elsewhere. The latency-metrics reference names those measurement boundaries.
Storage and RL are separate maturity questions
The launch labels secondary-storage tiering and reinforcement-learning support experimental. In RL, updated model weights travel from trainers to sampling workers; that requires version coordination in addition to bytes arriving. A sampler must not combine tensors from different updates and treat them as one coherent policy.
Likewise, a warm cache in host shared memory can survive some process restarts, but host failure can destroy it. Treat cached prefixes as reconstructible state unless the deployment explicitly provides stronger durability. The transfer system's lifecycle must agree with the serving engine's ownership and cancellation rules.
Agent Compute Engine manages another kind of state
An agent can pause while waiting for a tool, human input or a remote system. Its execution environment may contain processes and files, while the model server holds model weights and attention state. These are different resources with different recovery needs.
Agent Compute Engine addresses isolated agent execution and suspension/resumption. TPU Sync addresses model-state movement. Suspending an agent environment does not by itself guarantee that a separate model server releases its KV cache, and evicting a cache block does not checkpoint the agent's filesystem.
For a platform decision, trace both lifecycles and specify what must survive each pause. Pair this article with TPU Agent Suite for development workflows and TorchTPU migration for the execution backend. Keep the existing serving path if transfer overhead, compatibility or contention erases the benefit under representative load.
Turn it into offers. Work the real questions and concepts this maps to:
FAQ
TPU Sync is an open-source library for TPU model-state movement and KV-cache management. It connects device memory, host memory and remote workers through asynchronous transfer paths used by serving integrations.
No. Google reports improvements for specific prefix-cache and disaggregated-serving tests. Transfer time and end-to-end first-token time have different denominators; benefits depend on workload, cache reuse and contention.
The reviewed launch explicitly labels secondary-storage tiering and RL support experimental. Check the exact release and integration rather than treating every tier in an architecture diagram as equally mature.
No. TPU Sync manages model tensors and attention state. Agent Compute Engine manages agent execution environments and their suspend/resume lifecycle. An application may need both kinds of state, with separate recovery contracts.
Discussion (4)
A cache hit rate should say whether it counts requests, blocks or tokens. A high request hit rate can still reuse very little of the long prompt that dominates latency.
The transfer size must be the bytes actually crossing the slow link, after sharding. Aggregate model cache size can overstate one link, while replicated copies can undercut the expected saving.
And multiple replicas can share that link. A solo bandwidth test does not allocate the same rate to each replica during a burst.
Host shared memory surviving a process restart is useful. It still needs a cache-generation check before reuse, and it does not survive loss of the host.