An open source inference library handles KV-cache transfers across TPU hardware, providing components for chip-to-chip, network-based, and memory-level data movement.
Google has open-sourced TPU Raiden, an inference optimisation library designed to manage the movement of key-value (KV)-cache data during large language model inference. The library is available under the Apache 2.0 licence and is intended to address data-transfer requirements when inference workloads are distributed across different TPU resources.
Large language model inference happens in two main phases: prefill, where the model processes the entire input prompt, and decode, where it generates the response one token at a time. Disaggregated serving runs these phases on separate hardware pools. After the prefill phase, the resulting KV-cache is transferred to the hardware handling the decode phase. TPU Raiden is designed to manage this transfer.
The library supports KV-cache transfers between chips within a single machine, across TPU virtual machines over a network, and from TPU memory to host RAM. It also includes a shared-memory mode that allows KV-cache data to persist in DRAM when the model server is restarted.
The project is still under active development and is not yet intended for general production use. The open source release makes the library publicly available for developers to inspect and work with.
