Back Home

GitHub Repo

Google Open-Sources TPU Sync, Using Tiered KV Caching and Asynchronous Transfers to Reduce Inference Wait Times

On November 9, Google released TPU Sync, which integrates TPU memory, host DRAM, and cross-node data transfers. The system supports inference caching and reinforcement learning weight synchronization, but some features remain experimental, and performance figures apply to specific test conditions.

Alesker Abdullayev MWA1 · CC BY-SA 3.0 · Image source
zh-Hant

Google Cloud announced the open-source release of TPU Sync on November 9. It provides a low-level mechanism for moving model state for long-context inference and reinforcement learning. It handles the movement of KV cache and model weights across memory tiers and nodes, and provides PyTorch and JAX packages that integrate with the vLLM, SGLang, and llm-d ecosystems. Official announcement

The technical core preserves native TPU tensor layouts and moves data between HBM, host memory, and remote nodes using PJRT buffer aliasing and asynchronous DMA. Native C++ background threads handle transfers and memory operations, reducing Python serialization and Global Interpreter Lock involvement so data movement can overlap with prefill and decode computation. Host DRAM serves as a second-tier prefix cache: when HBM space is insufficient, it holds evicted blocks, which can later be loaded back on demand. Repository architecture overview

This is directly relevant to services that repeatedly use the same system prompt or accumulate multi-turn conversations: retaining computed prefixes can reduce the need to rerun prefill. In deployments that separate prefill and decode, the node that processes the prompt can also stream KV blocks to the generation node. Google reports that, in a low-concurrency test, average time to first token for Qwen3 235B fell by 60%. This is an internal result for a specific workload and should not be assumed to apply to all services. Test results and conditions

Two deployment caveats are worth noting. First, shared-storage tiering and the reinforcement learning transfer module are still experimental. Second, the repository states that the API is evolving rapidly and recommends pinning a release tag or the verified commit specified by lkg.version. The currently listed hardware validation covers TPU v7x and v6e. Engineering teams should evaluate whether tail latency and throughput both improve under their own cache hit rates, host memory budgets, and concurrent traffic. Version and hardware information

Sources

  1. Introducing TPU Sync: Fast KV cache and weight transfer for inference and RL on TPU
  2. google/tpu-sync:架構、硬體支援與部署指引