Google Has Open-Sourced Its TPU Raiden Inference Library In Apparent Bid To Externalize TPU Stack

Google appears to be taking steps towards making its TPUs more of a direct competition to NVIDIA’s GPUs.

The company has quietly open-sourced TPU Raiden, an inference optimization library that handles one of the most unglamorous but essential jobs in serving large language models: moving KV-cache data between chips during inference. The repository is live on GitHub under the Apache-2.0 license, listed under Google’s own account, and it is the direct TPU equivalent of NVIDIA’s NIXL, a library that has become a fairly important piece of infrastructure in the NVIDIA inference stack.

The release was first flagged by the research and analysis outfit SemiAnalysis, which pointed out that Raiden sits at the same layer of the stack as NIXL, and its arrival signals that Google is willing to externalize more of what has traditionally been treated as an internal, closely guarded part of its TPU software.

What TPU Raiden Actually Does

To understand why this matters, it helps to know what happens inside a chip when it serves a response to a prompt. Generating a reply from a large language model happens in two distinct phases. The first is prefill, where the model reads and processes the entire input prompt at once, a compute-heavy burst of work. The second is decode, where the model generates the reply one token at a time, a phase that is far lighter on raw compute but heavy on memory bandwidth.

Serving these two phases on the same chip is inefficient, because a chip built for the compute-heavy burst of prefill sits underused during the slower, memory-bound decode phase, and vice versa. So most large-scale inference setups today split the two phases across separate pools of hardware, a technique called disaggregated serving. One set of chips handles prefill, another handles decode, and the model runs across both.

The catch is that when prefill finishes processing a prompt, it produces something called a KV-cache, essentially a running memory of everything the model has read so far, stored as key and value tensors. That KV-cache then has to be handed off to whichever chip is doing decode work, and it needs to get there fast, since the whole point of splitting the work in the first place is to eke out speed. That handoff is exactly what TPU Raiden is built for.

Digging into the repository, Raiden includes modules for direct chip-to-chip KV-cache transfer within a single machine, cache transfer across TPU VMs over the network, and offloading cache blocks from TPU memory down to host RAM when they are not immediately needed, which helps avoid the kind of cold restart that happens when a serving process has to reload everything from scratch. There is also a shared memory mode that lets the cache persist in DRAM even if the model server itself gets restarted, useful during routine binary updates on a production fleet. Google’s documentation is candid that the project is still under active development and not yet meant for general production use, so this looks like an early but deliberate release rather than a finished product being dumped onto GitHub.

On the other side of the fence, NVIDIA’s equivalent library, NIXL, has already become a default option in serving frameworks like vLLM and is baked into NVIDIA’s Dynamo inference platform. It supports transfers over RDMA, NVMe, and even object storage like S3, and it works across GPUs, CPUs, and storage tiers through a single API. Raiden covers similar ground, but scoped specifically to TPUs and the way Google’s own DMA engines and interconnects are built, which is expected given it comes from the chip designer itself rather than a third party bolting support onto someone else’s silicon.

Why This Fits Into A Bigger Pattern

Google has never sold TPUs directly. Anyone who wants to use one has had to go through Google Cloud, and for a long time, the software that made TPUs efficient to run at scale stayed largely inside Google’s own walls, tuned for Google’s own workloads like Search, Gemini and internal ranking systems. That has started to shift over the last year or so, with Google striking a deal to supply Anthropic with a million TPUs, and even OpenAI, its most direct rival in the model race, quietly beginning to use Google TPUs for parts of its own operations.

Every one of those customers, current and prospective, needs the software ecosystem around TPUs to be as mature as the one that has grown up around NVIDIA GPUs over the past decade. That ecosystem is not just about raw chip performance, it is about whether the tooling exists to actually run production inference at scale without engineers having to build everything from scratch. Frameworks like vLLM and SGLang have become the default way most companies serve models, and both are deeply wired into NIXL already. If Raiden gets similar levels of integration into those same frameworks, it becomes far easier for a company already running NVIDIA-based disaggregated serving to point the same workload at TPUs with comparatively little engineering overhead.

There’s also a talent and trust angle to open-sourcing something like this. Infrastructure engineers who might otherwise default to building around NVIDIA’s stack, simply because that is what’s documented and battle-tested in the open, now have a public, inspectable version of Google’s approach to the same problem. That lowers the barrier for outside teams and cloud customers to actually adopt TPUs for serious inference workloads, rather than treating them as a NVIDIA GPU alternative that only Google’s own engineers know how to operate well.

None of this changes the fact that Google still does not sell TPUs as standalone hardware, and Raiden being open-source doesn’t mean the underlying chips are available to just anyone. But it does chip away at one of the bigger practical objections to adopting TPUs, which is that a company would be locking itself into an unfamiliar and comparatively undocumented stack. Google appears to be betting that if the tooling looks and feels similar enough to what teams already know from the NVIDIA world, more of them will be willing to make the switch, or at least run a mixed fleet.

NVIDIA has previously brushed off the growing chatter around TPUs, with the company arguing its GPUs remain a generation ahead and are the only platform that runs every AI model everywhere computing happens. Whether that holds up as Google continues to externalize pieces of its stack, one release at a time, is the question the rest of the industry will be watching closely.

Posted in AI