NEWS

Google open-sources TPU Sync code to synchronize cache and weights across TPUs

Google open-sources TPU Sync code to synchronize cache and weights across TPUs
Image: Redação iMasters
News
Reading
4 min
In short

The TPU Sync library, now open-sourced under Apache 2.0, moves KV cache blocks and model weights between TPUs and promises to cut time to first token by up to 90% in long-context workloads.

On November 12, 2026, Google Cloud open-sourced the code for TPU Sync, a native C++ library that moves KV cache blocks and model weights between TPU HBM memory, host DRAM, shared storage, and the data center network. The announcement came alongside the new TPU 8tx and TPU 8i Max systems, but the tool's focus isn't the silicon: it's solving the data bottleneck that leaves expensive accelerators idle while waiting for cache or weights to arrive, according to a report from InfoQ.

The project targets two specific scenarios. The first is long-context inference with disaggregated prefill and decode (disaggregated serving), where the KV cache needs to travel between different nodes in the cluster. The second is the reinforcement learning (RL) loop, where updated weights need to leave the trainer and quickly reach the workers that generate responses (rollout samplers).

The numbers Google released

For prefill and decode disaggregation, Google reports that the median KV cache transfer latency for the Qwen3 235B model dropped from several seconds to around 230 ms. That translated into average reductions of 60% in TTFT (time to first token, the time until the first response token) for that model, and 90% for a 480B model under low concurrency. For prefix caching in workloads with a 64,000-token window, the company reports drops of up to 90% in average TTFT.

These latency gains don't come from bandwidth alone. The pseudonymous developer @lgoodbookd did the math publicly on the new TPU 8i Max:

8i Max doubles HBM bandwidth to ~17.2TB/s, yet one 64K session on a 235B MoE still eats ~12.6GB of KV cache. That's why open-sourcing TPU Sync (host-DRAM KV tiering, ~223GB/s per v7x host) may matter as much as the new package. Bandwidth helps; capacity is still the wall.

@lgoodbookd, developer

In other words: the problem isn't just moving data faster between chips, it's having somewhere to offload what doesn't fit in HBM without losing speed. That's where the KV cache tiering to host DRAM that TPU Sync implements comes in.

What developers gain in practice

The library ships with Python wheels for JAX and PyTorch, exposing three sets of APIs: KV cache manager, KV storage (KV store), and weight synchronizer. In practice, it already connects to projects that much of the community running inference at scale knows well:

  • Acts as a KV connector for vLLM-TorchTPU, through vLLM's own KVConnector interface.
  • Works as a disaggregation and hierarchical cache backend for SGLang-TorchTPU.
  • Integrates with llm-d 1.0 running on GKE for prefix-cache-aware routing.
  • Has an experimental weight synchronization module for RL, aimed at Pathways and multi-controller JAX.

The GitHub repository absorbed the code from the previous project, called TPU Raiden, and is validated for TPU v7x and v6e. Anyone testing it needs to know that the PyTorch wheel requires a specific torch_tpu build that isn't available on a public index, and that only development builds are published on PyPI: the project itself recommends pinning to the last commit known to be stable, since the APIs still change quickly.

How this compares to what already exists on GPU

Those who work with GPU inference already deal with variations of the same problem, just with different tools. NVIDIA Dynamo and its NIXL transfer library tackle prefill and decode disaggregation in that ecosystem. LMCache offloads KV cache to CPU memory, disk, or remote storage. Mooncake is the KVCache-centric serving architecture behind Kimi, from Moonshot AI. And vLLM itself has its own native disaggregated prefill connectors.

The difference between TPU Sync and this group, according to InfoQ, lies in two points: it's built on top of PJRT and TPU-specific tile layouts, and it covers weight synchronization for RL in the same package, something the GPU tools mentioned handle separately. For those already running a hybrid GPU and TPU stack, this means deciding whether it's worth keeping two data transfer layers or migrating TPU workloads to the new library.

The take from those already testing it

Part of the technical community's initial reaction was less about the new hardware and more about what open-sourcing the code actually solves. Developer Anand Butani summed it up this way:

TPU sync cutting kv transfer latency 90% matters more to me than the chip names, long-context serving mostly dies on moving cache between nodes.

Anand Butani, developer

This kind of comment points to who TPU Sync is really for: teams that already feel firsthand the cost of keeping long context alive in production, whether in chatbots with extensive memory or in RAG pipelines that reuse large prefixes across calls.

What still needs to mature

Not everything is ready for production. The repository itself warns that the APIs change quickly, that only development builds are on PyPI, and that the disk storage layer and the RL weight synchronization module remain experimental. For teams evaluating adoption now, this means treating TPU Sync as a lab piece, not a production dependency, until a stable release is cut.

The code is under the Apache 2.0 license, and the repository includes architecture documentation, C++ and Python APIs, and usage examples for both single-host and multi-host disaggregated serving, plus examples of offloading KV to the host. For those running large LLM inference on TPU on GCP, this is the first public sign that Google is treating data movement between accelerators as a problem separate from the hardware, and is willing to open this layer for the community to test and contribute to.

Translated from the Brazilian Portuguese original · Read the original