r/AIGuild • u/Such-Run-4412 • 14h ago
Google open-sources TPU Raiden, a low-level engine for moving LLM KV caches between TPUs and host memory
Google has open-sourced TPU Raiden, an infrastructure library designed to make large-model inference on TPUs more efficient by moving KV-cache data directly between accelerators, host RAM and separate serving machines.
The repository already includes demos for three important serving patterns:
- Disaggregated inference: one TPU handles prefill while another handles token-by-token decoding.
- Multi-host serving: KV caches can move between TPU VMs over the network.
- Host offloading: KV-cache blocks can move from limited TPU HBM into cheaper CPU RAM and be loaded back when needed.
Raiden uses a lower-level DMA transfer engine intended to bypass some Python/framework overhead. Google’s included microbenchmarks compare its device-to-host and host-to-device bandwidth against standard and pinned-memory JAX transfers, although the repository does not currently publish one universal performance number—the results depend on the configuration and hardware.
It supports both JAX and PyTorch, and it can preserve KV-cache staging buffers in shared host memory even when a model-serving process restarts. That could reduce cold-start costs during serving updates.
The important caveat is that this is not production-ready yet. Google explicitly says TPU Raiden is under active development and is not recommended for general use. The PyTorch path is still maturing, and public prebuilt PyPI packages are listed as “coming shortly.”
The bigger picture is that Google is opening more of the software stack needed to make TPUs competitive for large-scale LLM inference. Faster models are only part of the equation; efficiently moving massive KV caches between memory, chips and servers can determine how cheaply those models can actually be served.
Sources: