Google’s TPU Sync cut TTFT by up to 90% in its own tests

TPU Sync, Google’s new open-source C++ library for moving the KV cache between TPU memory, host DRAM and remote nodes, cut time to first token by up to 90% in Google’s own benchmarks. Google announced it alongside its eighth-generation TPUs on November 9. The KV cache is the attention state a model keeps for tokens it has already processed. If your team serves large models on Cloud TPU, there are two limits for now: the library is validated on v7x and v6e, and it ships only as pre-release nightlies. Google’s README warns that the APIs are evolving rapidly and recommends pinning a tagged release or the verified commit in lkg.version instead of tracking main.

Installing it: nightlies, pins and two traps

The JAX wheel is the simpler path. pip install --pre tpu-sync-jax is self-contained and pulls the pinned jax, jaxlib and libtpu stack along with it. You need --pre because only .dev nightlies exist until the first stable release is cut. To stay on a build that passed all end-to-end functional and performance tests, check out the hash Google publishes in the repository root with git checkout $(cat lkg.version).

PyTorch takes more work. The tpu-sync-torch wheel does not pull torch or torch_tpu, and neither is on a public index. You install a matching torch_tpu and its exact torch pin first (the README’s example is torch==2.11.0), because the extension is ABI-locked to the build it was compiled against. The torch wheel’s release process is also not yet gated on end-to-end performance tests, since the torch path “is still maturing”.

Two README warnings concern the supply chain. The project used to be called TPU Raiden, and its old package names, tpu-raiden-jax and tpu-raiden-torch, are retired and not controlled by the project on PyPI. Google says not to install them. It also says not to combine a PyPI install with --extra-index-url pointing at a private registry. pip resolves across every configured index and picks the highest version from any of them, which leaves the install open to dependency-confusion attacks.

A source build needs Bazel 8.6.0 and Python 3.12, plus patchelf for the PyTorch extension. Each build targets one of seven listed JAX versions, from 0.10.0 through 0.12.0, with 0.11.2 as the default. You choose the version with RAIDEN_JAX_VERSION, and anything outside the list is rejected. If you switch versions, rerun build.sh in the environment you will actually use. An extension compiled against one jaxlib and loaded next to another doesn’t fail at import. Instead it crashes at run time, “usually far from the cause”. Although the library arrived with the eighth-generation announcement, it is validated and optimized only for Cloud TPU v7x and v6e.

What Google measured

The News Stack laid out Google’s benchmark results. On the 235-billion-parameter Qwen3, a KV-cache transfer that used to take several seconds took about 230 milliseconds. That cut average time to first token (TTFT) by 60% when relatively few requests were running at once. On a larger 480-billion-parameter model, the TTFT reduction reached 90%. In prefix-caching tests with 64,000-token contexts, TTFT fell by as much as 90%, throughput rose by more than 50% and P99 inter-token latency dropped 30%.

The News Stack also notes that results will vary with the workload, with how much cached data can be reused and with the network bandwidth available. It puts combined host-to-device and device-to-host throughput at 223 GB/s per Ironwood TPU host, close to the theoretical 256 GB/s limit of PCIe Gen5. The README gives a more conservative figure: about 200 GB/s of bidirectional PCIe bandwidth with production-scale KV-cache block sizes.

SemiAnalysis’s analysis suggests that external TPU serving previously lacked the infrastructure to separate prefill (processing the prompt) from decode (generating tokens) effectively. SemiAnalysis’s InferenceX data also shows Nvidia’s GB200 and GB300 NVL72 systems ahead in parts of the disaggregated-serving performance curve. A direct comparison with a fully tuned TPU v7 setup has yet to be published.

How it moves the cache

TPU Sync runs underneath Python orchestrators. It reads JAX arrays and PyTorch tensors through their native PjRtBuffer hardware descriptors, so transfers skip the intermediate CPU copy. The transfers, networking and memory allocation run on C++ background thread pools outside the Python GIL, which lets the data moves overlap with prefill and decode on the TPU.

When HBM fills up, least-recently-used cache blocks are evicted to host RAM, which works as a second-tier prefix cache. Those blocks come back through page-locked, asynchronous DMA when a prompt needs them. The host-side cache and its eviction metadata live in POSIX shared memory under /dev/shm. So if an inference container restarts, upgrades or crashes, the offloaded cache stays in host DRAM, and the server reattaches to it on reboot rather than running prefill again from a cold start.

TPU Sync shipped together with TorchTPU, which supports vLLM and SGLang. The News Stack says teams already using those engines can run disaggregated workloads on TPUs without rebuilding their serving stack. For reinforcement learning, the library moves weights and activations between trainers and samplers without extra copies. Google asks teams wiring it into a custom serving or RL stack to open a GitHub issue and discuss compatibility and workload requirements with the maintainers.