TPU Sync Explained: KV Cache Offloading, Setup and Performance Limits

Google introduced TPU Sync on November 9 as an open-source library for moving model state across TPU clusters. It targets a specific serving problem: expensive accelerators waiting while large conversation caches move between memory tiers or between the workers that process prompts and generate responses.

For a team already serving long-context models on TPUs, the question is whether context movement or repeated prompt processing is consuming enough time to justify another infrastructure component. A small, mostly uncached workload may have a different bottleneck entirely.

What TPU Sync changes in an inference server

A model's key-value cache stores attention state from tokens it has already processed. Keeping that state available avoids recomputing the entire conversation on every turn, but long prompts and concurrent sessions can fill accelerator high-bandwidth memory, or HBM.

The TPU Sync repository describes a native C++ layer that moves this state between TPU HBM, host RAM and other nodes. It accesses device buffers directly and performs transfers on background threads, avoiding Python serialization in the data path. Host RAM becomes an additional cache tier rather than merely a staging area for loading a model.

There are two related but distinct use cases. Prefix caching reuses context that requests share, such as a long system instruction or a repeated document. Disaggregated serving separates prefill, which processes the incoming prompt, from decode, which generates output tokens. The latter still needs an efficient way to get the processed context to the decode worker.

More cache capacity only helps when the application can reuse the state it preserves. A workload with unique prompts and little shared context will not get the same benefit as a fleet repeatedly processing the same documents.

Which serving stacks can use it

Google's PyTorch on TPU documentation connects TPU Sync to vLLM-TorchTPU and SGLang-TorchTPU for host-cache offloading and transfers between prefill and decode workers. It also describes llm-d as the fleet-level layer for routing requests using cache state and coordinating distributed inference.

That separation is useful when planning a rollout. TPU Sync moves state; the serving engine schedules model execution; a fleet router decides which replica should receive a request. Fast movement cannot compensate for routing that repeatedly sends a conversation away from the worker holding its context.

The project's currently validated hardware list is TPU v7x and v6e. Do not infer support for every newly announced TPU system from the timing of the launch. Likewise, its experimental reinforcement-learning and storage-tiering features should be evaluated separately from the core host-memory cache path.

Installation starts with matching versions

The repository warns that APIs are evolving quickly and recommends a tagged release or the verified revision recorded in lkg.version rather than following main. Its current installation instructions say that only development wheels are available until the first stable release, so the package commands include the prerelease flag:

pip install --pre tpu-sync-jax

or, for the PyTorch variant:

pip install --pre tpu-sync-torch

These are alternative framework paths, not a recommendation to install both into a production environment. Check the version-specific requirements before running either. The JAX wheel brings its pinned JAX, jaxlib and libtpu stack; the PyTorch extension requires a matching torch and torch_tpu pairing, documented in the project's release instructions.

For a source build, the repository validates Python 3.12 and specifies Bazel 8.6.0. Its PyTorch build also requires patchelf. These requirements matter because the extension shares a process with the framework's native libraries. A mismatched XLA or C++ ABI can produce a runtime crash even when compilation succeeded.

Use a dedicated environment and record the exact package versions or source revision. If an import fails after an upgrade, compare that pairing first instead of assuming the model checkpoint is damaged. This article follows the installation documentation; it does not report an independently executed TPU deployment.

Read the 90% claim with its workload attached

Google's launch post reports large reductions in transfer overhead and time to first token. Its detailed results distinguish workload sizes: controlled low-concurrency tests report an average TTFT reduction of 60% for a 235B model and 90% for a 480B model. The test note identifies internal July 2026 testing on TPU v7x with specific network and software configurations.

Those results are evidence of what Google measured, not a promise that every endpoint becomes 90% faster. TTFT, transfer latency, inter-token latency and overall throughput measure different parts of the request. Improving one does not establish the same percentage improvement in the others.

A useful evaluation should therefore compare the same model, precision, prompt distribution, output lengths and traffic level with and without the relevant cache path. Include a cold-cache run, repeated-prefix traffic and a mixed workload with low reuse. Report tail latency as well as averages so a faster common case does not conceal worse slow requests.

A practical decision test before rollout

Consider an assistant that repeatedly answers questions about the same large technical repository. If each worker discards the repository's processed context under memory pressure, requests repeatedly pay the prefill cost. A larger host-backed prefix cache could help, provided the router and cache configuration make that state reusable.

Now change the workload to unrelated documents arriving once each. There may be little prefix reuse, so the useful experiment shifts toward transfer costs and concurrency rather than cache-hit rate. Treating both workloads as one benchmark would hide the reason the system improved or regressed.

Also test recovery. The repository describes host shared-memory cache persistence across inference-server restarts. That is not durable storage that survives losing the host. Include server restart, host loss and cache invalidation in the operational plan, and make sure a stale cache cannot silently become the basis of a correctness claim.

TPU Sync is worth evaluating when state movement is a measured constraint. The next infrastructure decision should follow the shape of the workload, not the largest percentage in the announcement.

I’m Paras. I work at the intersection of AI, engineering and distribution.

I help technical founders explain infrastructure products and performance trade-offs in stories developers can use.

Book a 30-minute call.