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 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. IThe open-source release makes the library publicly available for developers to inspect and work with.
