Google has open-sourced TPU Sync, an Apache-2.0 C++ library with Python bindings for PyTorch and JAX that moves KV caches and model weights across a unified memory hierarchy spanning TPU HBM, pinned host DRAM, POSIX storage (including Cloud Managed Lustre), and the data center network (DCN) on Cloud TPU v7x (Ironwood) and v6e. The library uses layout-preserving direct DMA on PJRT buffers and PyTorch tensors, avoiding serialization and bounce buffers, while overlapping transfers with compute via DMA pacing and executing resharding outside the XLA graph. On a TPU v7x (2x2x1) testbed with PCIe Gen5 and 200 Gbps NICs, TPU Sync achieved up to 223 GB/s per host on PCIe D2H/H2D and ~350 Gbps aggregate DCN throughput across two NICs. In disaggregated prefill/decode serving benchmarks using Qwen3 235B and 480B MoE models, median KV transfer latency dropped from several seconds to ~230 ms (over 90% reduction), average time-to-first-token (TTFT) fell 60% for 235B and 90% for 480B at low concurrency, and token throughput rose 30–70% during traffic spikes. For 64K-context prefix caching, effective cache capacity grew an order of magnitude, TTFT dropped up to 90%, throughput increased over 50%, and P99 inter-token latency fell 30%. TPU Sync integrates with vLLM-TorchTPU, SGLang-TorchTPU, and llm-d on GKE; RL weight broadcast and storage tiering remain experimental, and all figures are self-reported from July 2026 internal tests.

Read original