Running Google Cloud TPUs Inside Edera Zones

Running AI workloads from different customers or teams creates an infrastructure problem. You need to isolate those workloads, but making a separate TPU VM or Kubernetes cluster the isolation boundary for every workload also means more infrastructure to provision, secure, upgrade, and operate.

We wanted another option.

We've now got Google Cloud TPU workloads running inside Edera Protect zones. Each zone gets its own Linux kernel and IOMMU context for its assigned TPU chips. Multiple zones can run on the same TPU host, and Kubernetes pods can request TPU devices through our DRA driver without requiring changes to the existing libtpu userspace stack.

For platform teams, this means they can isolate workloads without making a dedicated TPU VM or cluster the default isolation unit. That could simplify how multi-tenant AI infrastructure is operated, particularly as the number of workloads and customers grows. 

There was a significant engineering problem to solve, though. Google's TPU VMs don't expose VT-x or AMD-V, so conventional hardware-assisted virtualization wasn't available to us. We used Xen paravirtualization, but making the TPU work inside an unprivileged guest required changes across Xen, Linux, and the Edera runtime. We also had to implement per-zone DMA isolation, work around Google's emulated IOMMU behavior, and make mapping millions of memory pages practical.

Here's what it took to get there, including the isolation limitations that still exist.

Why TPU Workload Isolation Matters for Multi-Tenant AI Infrastructure

If you're running AI workloads from different customers or teams, you probably don't want them all sharing a Linux kernel. You can separate them across TPU VMs or clusters, but that also means more infrastructure to provision and operate. On the Google TPU VMs we tested, you can't just reach for a hypervisor that needs hardware virtualization extensions, because the guest doesn't get those extensions.

We can now run those workloads in separate Edera zones on the same TPU host. Each zone gets its own kernel and its own IOMMU context for the devices assigned to it. Kubernetes pods can claim TPU devices through our DRA driver, and the existing libtpu userspace stack runs without changes. For a platform team, that makes TPU workload isolation possible without making a dedicated TPU VM or Kubernetes cluster the isolation unit for every workload. We're still assigning whole chips, not splitting one TPU chip between tenants, and there are device-isolation limitations I'll get to below.

What Per-Zone TPU Isolation Changes for Platform Teams

If you're building a multi-tenant AI platform, isolation affects how you design and operate your infrastructure.

Giving every workload its own VM or cluster is one way to separate workloads, but it also increases the number of environments your platform team needs to manage. With Edera zones, workloads can have separate kernels while using TPU chips assigned from the same host. You still need to manage device allocation and capacity, but the isolation boundary no longer has to be an entire VM or cluster.

We also wanted this to work with the tooling platform teams already use. Device allocation goes through Kubernetes DRA, using a driver we built that is independent of vendor DRA drivers, and the existing libtpu stack works without application changes. Teams don't need to adopt a new TPU programming model just to use the isolation.

There are limits. We're not splitting individual TPU chips across tenants, and Google's emulated IOMMU doesn't support interrupt remapping. The per-zone IOMMU contexts isolate DMA, but they don't provide equivalent interrupt isolation. Those constraints matter when evaluating the design against your security requirements.

The result is another option for running isolated accelerator workloads. Whether it reduces your operating costs will depend on how your infrastructure is organized. What we can demonstrate today is the isolation architecture and the engineering needed to make it work.

Why Xen PV Booted but the TPU Would Not Attach

The TPU VMs don't expose hardware virtualization to the guest. Xen PV doesn't need VT-x or AMD-V, so running a zone wasn't the problem. The trouble started when we tried to give that zone a TPU.

TPUs don't use a conventional kernel driver for the device operations. The TPU runtime, libtpu, binds each TPU to vfio-pci and drives it from userspace. It maps PCI BARs, programs DMA, and handles interrupts. VFIO expects an IOMMU, and a Xen PV guest didn't have one.

Google emulates an IOMMU on the TPU host and presents it as an AMD IOMMU. It handles DMA translation, but it doesn't provide interrupt remapping. That means vfio-pci needs allow_unsafe_interrupts, and while DMA is confined by the IOMMU, interrupts don't have the same isolation. We need to be clear about that limitation.

Rewriting Xen’s IOMMU Driver for Google’s Emulated AMD IOMMU

Xen's AMD IOMMU driver worked on physical AMD machines. Google's emulated AMD IOMMU was less forgiving. It implements the behavior Linux's driver expects, and some things Xen does differently from Linux didn't work there at all.

A few of the failures we had to fix

•  Xen programmed the 64-bit device-table and command-buffer base registers with two 32-bit writes. Linux uses a single 8-byte write. The emulator only decoded the latter, so Xen never got the command buffer started.

•  Xen assigned domain ID 0 to its first IOMMU domain. Linux doesn't use that ID, and the emulator rejected the first invalidation command referencing it.

•  Xen cleared the whole control register at startup, including a bit the emulator had set itself. Once that bit was gone, the emulator started rejecting commands.

•  Xen built five-level IOMMU page tables. Linux used three levels, and the emulator rejected commands for the deeper entries.

One rejected command could stop the IOMMU command processor. Xen had no recovery path, so everything waiting on later invalidations got stuck too. We added command-buffer recovery and bounded the completion waits instead of spinning forever. The GCE-specific behavior, including the page-table-depth limit, is behind options so it doesn't change what Xen does on other hosts.

Bringing a PV-IOMMU Into an Unprivileged Xen Zone

Fixing Xen's host-side driver wasn't enough. vfio-pci inside the zone still needed to see an IOMMU it could use. We started with Teddy Astie's Xen IOMMU subsystem redesign and PV-IOMMU interface RFC. The design supports multiple IOMMU contexts per domain and hypercalls for managing them, but the original driver work focused on dom0. We extended it to work in an unprivileged domU and adapted it for GCE.

Xen remains in charge of the physical IOMMU. Linux inside the zone sees a PV-IOMMU device and manages mappings through hypercalls. Xen puts each zone's mappings in its own IOMMU context, covering only the devices assigned to that zone.

A chip's DMA goes through its assigned zone's IOMMU context, so it can reach only the memory that zone has mapped for it. Once this worked, vfio-pci could operate as it does with a physical IOMMU, and libtpu needed no changes. We also had to fix PV PCI passthrough behavior across Xen, Linux, and our runtime to make the full path work.

Replacing swiotlb Bounce Buffers With Machine-Contiguous DMA

Memory that's contiguous from the PV guest's point of view isn't necessarily contiguous in host machine memory. Linux normally handles that through swiotlb, copying DMA data through bounce buffers the device can access. Doing that for accelerator traffic would be a terrible tradeoff.

We added a way for the zone kernel and Xen to allocate machine-contiguous memory regions together. The kernel uses those regions to program the IOMMU, and the TPU can read and write the zone's memory directly instead of copying data through bounce buffers.

Batching Millions of 4K IOMMU Mappings With Xen Multicalls

By this point, the TPU actually worked, but some tests still timed out. PV guests map memory in 4K pages. libtpu maps a lot of memory for DMA, and we were making IOMMU mapping calls one page at a time. With millions of pages, the overhead was enormous.

We batched the memory operations behind vfio-pci using Xen multicalls, so one hypercall could carry many operations. That fixed the mapping bottleneck.

Zone teardown was slow too. Xen spent about five seconds unmapping a zone's memory, during which dom0 userspace couldn't run. Batching unmaps cut that stall by about 30x. We can now tear down a zone with roughly 90 GB mapped, more than 20 million 4K pages, with only a brief dom0 pause.

A Vendor-Independent Kubernetes DRA Driver for TPU Zones

Getting the TPU working inside a zone solved the low-level problem, but platform teams also need a way to request those devices through Kubernetes. Vendor DRA drivers are designed to expose a device to a container that shares the host's kernel. In an Edera zone, the device has to be assigned to the zone itself, bound to vfio-pci inside the zone's kernel, and mapped into that zone's IOMMU context before the workload can use it.

We built a DRA driver that handles that path directly and doesn't depend on any vendor's DRA driver. A pod claims TPU devices through a standard DRA resource claim, and our driver assigns the matching chips to the pod's zone. Because the driver works at the level of PCI passthrough and per-zone IOMMU contexts instead of a specific vendor's device model, we aren't waiting on each accelerator vendor to support isolated zones in their own drivers, and the same allocation model can extend to other accelerators.

TPU Zone Isolation Status, Merged Changes, and Known Limits

The Xen, Linux kernel, and Edera runtime changes are merged. We have a TPU-specific zone kernel variant that binds the chips to vfio-pci and handles hotplug. Our Kubernetes DRA driver, which works independently of vendor DRA drivers, lets pods claim the devices. The remaining integration work is for the installer to deploy that driver automatically.

There are three limitations worth stating.

• Google's emulated IOMMU doesn't support interrupt remapping. DMA is isolated through the per-zone IOMMU contexts, but TPU interrupts don't get equivalent isolation.
• A zone receives whole TPU chips. A single chip can't be shared between zones with this implementation.
• DMA mappings still use 4K pages, so IOMMU mapping work grows with how much memory a workload maps. Batching makes that manageable but doesn't remove the cost, and we plan to add support for larger page sizes so that mapping work no longer scales with 4K page counts.

Beyond TPUs: Reusable PV-IOMMU and PCI Passthrough for Accelerators

The bulk of the changes are in PV-IOMMU support, PCI passthrough, and how we handle DMA. None of those paths depend on the device being a TPU. We've already used the same work to run NVIDIA GPUs inside Xen PV zones.

For customers building AI infrastructure, we now have a way to give workloads separate kernels and per-zone device DMA mappings on TPU hosts that don't expose hardware virtualization extensions. For us, we've extended what Xen PV guests can do with accelerators, not just added support for one more device. That's a lot more useful than a TPU demo that happens to boot.