From 1ccfee728e03d906f111069476c43acf363d0755 Mon Sep 17 00:00:00 2001
From: Alex Zenla <alex@edera.dev>
Date: Sat, 10 Oct 2026 14:22:07 +0000
Subject: [PATCH] feat(zone-tpu): add a TPU zone flavor with vfio-pci binding
 hooks

Add the zone-tpu variant, which states the VFIO PCI options the TPU runtime
needs, and ships setup and hotplug hooks that bind any Google processing
accelerator to vfio-pci and allow unsafe interrupts, which Xen PV guests need.
The hooks use the native bind-driver and sysfs-module-parameters fields, so
the addons carry no overlay or shell for them.

Split the nvidia section out of firmware.sh into firmware-nvidia.sh, add
firmware-tpu.sh, and move the hook TOML from inline heredocs into real files
under hooks/.
---
 config.yaml                             |   6 +
 configs/x86_64/zone-tpu.fragment.config |  15 ++
 hack/build/firmware-nvidia.sh           | 214 ++++++++++++++++++++++
 hack/build/firmware-tpu.sh              |   7 +
 hack/build/firmware.sh                  | 225 +-----------------------
 hack/build/refresh-nvidia-versions.py   |   2 +-
 hooks/nvidia-persist.toml               |  11 ++
 hooks/tpu-vfio.toml                     |  19 ++
 8 files changed, 277 insertions(+), 222 deletions(-)
 create mode 100644 configs/x86_64/zone-tpu.fragment.config
 create mode 100755 hack/build/firmware-nvidia.sh
 create mode 100755 hack/build/firmware-tpu.sh
 create mode 100644 hooks/nvidia-persist.toml
 create mode 100644 hooks/tpu-vfio.toml

diff --git a/config.yaml b/config.yaml
index 2ae4460f..f8ba250f 100644
--- a/config.yaml
+++ b/config.yaml
@@ -52,6 +52,12 @@ flavors:
 - name: zone-amdgpu
   constraints:
     lower: '6.1'
+- name: zone-tpu
+  # zone kernel for Google TPU passthrough: VFIO PCI, plus a hook that binds the TPU to vfio-pci
+  architectures:
+  - x86_64
+  constraints:
+    lower: '6.1'
 - name: zone-nvidiagpu
   # we will generate distinct images for each of these
   # local tags. For `zone-nvidiagpu` specifically,
diff --git a/configs/x86_64/zone-tpu.fragment.config b/configs/x86_64/zone-tpu.fragment.config
new file mode 100644
index 00000000..7c282d44
--- /dev/null
+++ b/configs/x86_64/zone-tpu.fragment.config
@@ -0,0 +1,15 @@
+#
+# Edera kernel config snippet
+# - Generated from delta config: configs/x86_64/zone-tpu.fragment.config
+# - Against base config: configs/x86_64/zone.config
+#
+# The TPU runtime opens the chip through /dev/vfio, so the VFIO PCI stack is
+# stated here explicitly instead of relying on the base zone config.
+CONFIG_PCI_MSI=y
+CONFIG_VFIO=y
+CONFIG_VFIO_GROUP=y
+CONFIG_VFIO_CONTAINER=y
+CONFIG_VFIO_IOMMU_TYPE1=y
+CONFIG_VFIO_PCI_CORE=y
+CONFIG_VFIO_PCI_INTX=y
+CONFIG_VFIO_PCI=y
diff --git a/hack/build/firmware-nvidia.sh b/hack/build/firmware-nvidia.sh
new file mode 100755
index 00000000..53bc9e08
--- /dev/null
+++ b/hack/build/firmware-nvidia.sh
@@ -0,0 +1,214 @@
+#!/bin/bash
+set -e
+
+# Sourced from firmware.sh; expects its FIRMWARE_OUTPUT_PATH and WORKLOAD_OVERLAY_PATH.
+
+# For nvidia kernel, firmware is distributed via their out-of-tree .run userspace
+# package, which we need to fetch and extract.
+# TODO(BML) fix arm64 builds
+if [ "${KERNEL_FLAVOR}" = "zone-nvidiagpu" ] && [ "${TARGET_ARCH_STANDARD}" = "x86_64" ]; then
+	# Check that we already set NV_VERSION when building the out-of-tree kmod,
+	# naturally, we must fetch the matching runtime package.
+	if [ -z "${NV_VERSION}" ]; then
+		echo "ERROR: flavor is zone-nvidiagpu but no NV_VERSION is set"
+		exit 1
+	fi
+
+	OLDDIR=$PWD
+
+	# TBH it's probably the same firmware regardless
+	if [ "${TARGET_ARCH_STANDARD}" = "aarch64" ]; then
+		NV_RUN_FILE="NVIDIA-Linux-aarch64-${NV_VERSION}.run"
+		NV_RUN_URL="https://us.download.nvidia.com/XFree86/aarch64/${NV_VERSION}/${NV_RUN_FILE}"
+	elif [ "${TARGET_ARCH_STANDARD}" = "x86_64" ]; then
+		NV_RUN_FILE="NVIDIA-Linux-x86_64-${NV_VERSION}.run"
+		NV_RUN_URL="https://us.download.nvidia.com/XFree86/Linux-x86_64/${NV_VERSION}/${NV_RUN_FILE}"
+	fi
+
+	NV_EXTRACT_PATH="$(mktemp -d)/extracted-${NV_VERSION}"
+	mkdir -p "$NV_EXTRACT_PATH"
+
+	echo "Downloading NVIDIA runtime package for driver ${NV_VERSION} from: $NV_RUN_URL"
+	curl --retry 5 --retry-delay 2 --retry-max-time 30 --retry-all-errors -L -o "$NV_EXTRACT_PATH/$NV_RUN_FILE" "$NV_RUN_URL"
+	chmod +x "$NV_EXTRACT_PATH/$NV_RUN_FILE"
+	"$NV_EXTRACT_PATH/$NV_RUN_FILE" -x --target "$NV_EXTRACT_PATH/out"
+	# Compress firmwares on-disk
+	# As of 6.x kernels the kconfig explicitly says you must use crc32 or none, not the default crc64.
+	xz -C crc32 "$NV_EXTRACT_PATH"/out/firmware/*
+	ls -lah "$NV_EXTRACT_PATH/out/firmware/"
+	NV_FW_PATH="${FIRMWARE_OUTPUT_PATH}/nvidia/$NV_VERSION"
+	mkdir -p "${NV_FW_PATH}"
+	cp "$NV_EXTRACT_PATH"/out/firmware/*.xz "${NV_FW_PATH}"
+
+	#
+	# Workload overlay
+	#
+	# Install various binaries necessary to make compute workloads run properly
+	# under containers.
+	#
+
+	create_links() {
+		local base_path="$1"
+		# create soname links
+		find "$base_path" -type f -name '*.so*' ! -path '*xorg/*' -print0 | while read -r -d $'\0' _lib; do
+			_soname=$(dirname "${_lib}")/$(readelf -d "${_lib}" | grep -Po 'SONAME.*: \[\K[^]]*' || true)
+			_base=$(echo "${_soname}" | sed -r 's/(.*)\.so.*/\1.so/')
+			[[ -e "${_soname}" ]] || ln -s "$(basename "${_lib}")" "${_soname}"
+			[[ -e "${_base}" ]] || ln -s "$(basename "${_soname}")" "${_base}"
+		done
+	}
+
+	multiarch_symlink_mirror() {
+		local base_path="$1"
+		local triplet="$2"
+		local src="${base_path}/usr/lib"
+		local dst="${base_path}/usr/lib/${triplet}"
+		pushd "$src" >/dev/null
+
+		# 1) Recreate directory structure (excluding the triplet subtree to avoid recursion)
+		# Use find from within src to get clean relative paths like "./gbm"
+
+		# Create directories
+		find . \
+			-mindepth 1 \
+			-path "./${triplet}/*" -prune -o \
+			-path "./${triplet}" -prune -o \
+			-type d -print0 |
+			while IFS= read -r -d '' rel_dir; do
+				# Strip leading "./"
+				rel_dir="${rel_dir#./}"
+				mkdir -p "${dst}/${rel_dir}"
+			done
+
+		# Helper: decide whether to mirror this path
+		should_mirror() {
+			local p="$1"
+			if [[ -f "$p" ]]; then
+				# regular file -> yes
+				return 0
+			elif [[ -L "$p" ]]; then
+				# symlink -> only if it ultimately targets a regular file
+				# resolve; readlink -f fails (non-zero) on broken loops/broken links
+				local tgt
+				if ! tgt="$(readlink -f -- "$p" 2>/dev/null)"; then
+					return 1
+				fi
+				[[ -f "$tgt" ]] && return 0 || return 1
+			else
+				# sockets, fifos, device nodes, etc. -> no
+				return 1
+			fi
+		}
+
+		# 2) Mirror regular files and symlinks-to-regular-files
+		find . \
+			-path "./${triplet}/*" -prune -o \
+			\( -type f -o -type l \) \
+			-print0 |
+			while IFS= read -r -d '' rel_item; do
+				rel_item="${rel_item#./}"
+				# Filter: only files and symlinks that resolve to files
+				if ! should_mirror "$rel_item"; then
+					continue
+				fi
+
+				local link_path="${dst}/${rel_item}"
+				local link_dir
+				link_dir="$(dirname "$link_path")"
+				mkdir -p "$link_dir"
+
+				# Always point the mirror symlink to the *original path in /usr/lib*,
+				# preserving whether the original is a file or a symlink.
+				local target="../${rel_item}"
+
+				if [[ -L "$link_path" ]]; then
+					ln -snf "$target" "$link_path"
+				elif [[ -e "$link_path" ]]; then
+					echo "WARN: exists and not a symlink, skipping: $link_path" >&2
+				else
+					ln -s "$target" "$link_path"
+				fi
+			done
+
+		popd >/dev/null
+	}
+
+	# GTK (used by nvidia-settings, which we don't install)
+	rm -f "$NV_EXTRACT_PATH/out"/libnvidia-gtk*
+
+	# Wayland/GBM
+	mkdir -p "$WORKLOAD_OVERLAY_PATH/usr/lib/gbm"
+	ln -s "../libnvidia-allocator.so.${NV_VERSION}" "$WORKLOAD_OVERLAY_PATH/usr/lib/gbm/nvidia-drm_gbm.so"
+
+	# DRI driver
+	install -Dm755 "$NV_EXTRACT_PATH/out/"nvidia_drv.so "$WORKLOAD_OVERLAY_PATH/usr/lib/xorg/modules/drivers/nvidia_drv.so"
+
+	# GLX extensions
+	install -Dm755 "$NV_EXTRACT_PATH/out/libglxserver_nvidia.so.${NV_VERSION}" "$WORKLOAD_OVERLAY_PATH/usr/lib/nvidia/xorg/libglxserver_nvidia.so.${NV_VERSION}"
+	ln -s "libglxserver_nvidia.so.${NV_VERSION}" "$WORKLOAD_OVERLAY_PATH/usr/lib/nvidia/xorg/libglxserver_nvidia.so.1"
+	ln -s "libglxserver_nvidia.so.${NV_VERSION}" "$WORKLOAD_OVERLAY_PATH/usr/lib/nvidia/xorg/libglxserver_nvidia.so"
+
+	# Remove already-installed libs, so we don't accidentally include them in
+	# filters below
+	rm -f "$NV_EXTRACT_PATH/out"/nvidia-drv.so
+	rm -f "$NV_EXTRACT_PATH/out"/libglxserver_*
+
+	# Remove unnecessary libraries
+	rm -f "$NV_EXTRACT_PATH/out"/libnvidia-wayland*
+
+	for LIBRARY in "$NV_EXTRACT_PATH/out/"lib*.so*; do
+		BN="$(basename "$LIBRARY")"
+		install -Dm755 "$LIBRARY" "$WORKLOAD_OVERLAY_PATH/usr/lib/$BN"
+	done
+
+	for BINARY in "$NV_EXTRACT_PATH/out/"nvidia-{cuda-mps-control,cuda-mps-server,debugdump,pcc,smi}; do
+		BN="$(basename "$BINARY")"
+		install -Dm755 "$BINARY" "$WORKLOAD_OVERLAY_PATH/usr/bin/$BN"
+	done
+
+	for MAN1 in "$NV_EXTRACT_PATH/out/"nvidia-{cuda-mps-control,smi}.1.gz; do
+		BN="$(basename "$MAN1")"
+		install -Dm644 "$MAN1" "$WORKLOAD_OVERLAY_PATH/usr/share/man/man1/$BN"
+	done
+
+	# Install ICD loaders for OpenGL, Vulkan, and Vulkan SC
+	install -Dm644 "$NV_EXTRACT_PATH/out/"10_nvidia.json "$WORKLOAD_OVERLAY_PATH/usr/share/glvnd/egl_vendor.d/10_nvidia.json"
+	install -Dm644 "$NV_EXTRACT_PATH/out/"nvidia.icd "$WORKLOAD_OVERLAY_PATH/etc/OpenCL/vendors/nvidia.icd"
+	install -Dm644 "$NV_EXTRACT_PATH/out/"nvidia_icd.json "$WORKLOAD_OVERLAY_PATH/usr/share/vulkan/icd.d/nvidia_icd.json"
+	install -Dm644 "$NV_EXTRACT_PATH/out/"nvidia_layers.json "$WORKLOAD_OVERLAY_PATH/usr/share/vulkan/implicit_layer.d/nvidia_layers.json"
+	install -Dm644 "$NV_EXTRACT_PATH/out/"nvidia_icd_vksc.json "$WORKLOAD_OVERLAY_PATH/usr/share/vulkansc/icd.d/nvidia_icd_vksc.json"
+
+	create_links "$WORKLOAD_OVERLAY_PATH"
+
+	# For Debian-like distributions
+	multiarch_symlink_mirror "$WORKLOAD_OVERLAY_PATH" x86_64-linux-gnu
+
+	#
+	# Create NVIDIA persistence mode hook, ensures the driver is loaded and
+	# initialized, and always ready to run a workload
+	#
+	NVIDIA_BOOTSTRAP_OVERLAY_PATH="$ADDONS_OUTPUT_PATH/overlays/nvidia-bootstrap"
+
+	mkdir -p "$NVIDIA_BOOTSTRAP_OVERLAY_PATH"
+
+	# shellcheck disable=SC2043 # one entry today; loop mirrors the sibling install blocks
+	for BINARY in "$NV_EXTRACT_PATH/out/"nvidia-smi; do
+		BN="$(basename "$BINARY")"
+		install -Dm755 "$BINARY" "$NVIDIA_BOOTSTRAP_OVERLAY_PATH/usr/bin/$BN"
+	done
+
+	# HACK: Just using the builder's glibc binaries so we can run nvidia-smi
+	# without a full workload image
+	for LIBRARY in "$NV_EXTRACT_PATH/out/"libnvidia-ml.so* /lib64/ld-linux-x86-64.so.2 /lib/x86_64-linux-gnu/lib{pthread,m,dl,c,rt}.so.*; do
+		BN="$(basename "$LIBRARY")"
+		install -Dm755 "$LIBRARY" "$NVIDIA_BOOTSTRAP_OVERLAY_PATH/usr/lib/$BN"
+	done
+	ln -s usr/lib "$NVIDIA_BOOTSTRAP_OVERLAY_PATH/lib64"
+
+	create_links "$NVIDIA_BOOTSTRAP_OVERLAY_PATH"
+	multiarch_symlink_mirror "$NVIDIA_BOOTSTRAP_OVERLAY_PATH" x86_64-linux-gnu
+
+	install -Dm644 "${KERNEL_DIR}/hooks/nvidia-persist.toml" "${ADDONS_OUTPUT_PATH}/hooks/nvidia-persist.toml"
+
+	cd "$OLDDIR"
+fi
diff --git a/hack/build/firmware-tpu.sh b/hack/build/firmware-tpu.sh
new file mode 100755
index 00000000..1e4cd84b
--- /dev/null
+++ b/hack/build/firmware-tpu.sh
@@ -0,0 +1,7 @@
+#!/bin/bash
+set -e
+
+# Sourced from firmware.sh; the TPU needs no firmware, only the hook that binds it to vfio-pci.
+if [ "${KERNEL_FLAVOR}" = "zone-tpu" ]; then
+	install -Dm644 "${KERNEL_DIR}/hooks/tpu-vfio.toml" "${ADDONS_OUTPUT_PATH}/hooks/tpu-vfio.toml"
+fi
diff --git a/hack/build/firmware.sh b/hack/build/firmware.sh
index 7ef44c5d..4bb49a4c 100644
--- a/hack/build/firmware.sh
+++ b/hack/build/firmware.sh
@@ -57,225 +57,8 @@ if [ -n "${FIRMWARE_URL}" ]; then
 	fi
 fi
 
-# For nvidia kernel, firmware is distributed via their out-of-tree .run userspace
-# package, which we need to fetch and extract.
-# TODO(BML) fix arm64 builds
-if [ "${KERNEL_FLAVOR}" = "zone-nvidiagpu" ] && [ "${TARGET_ARCH_STANDARD}" = "x86_64" ]; then
-	# Check that we already set NV_VERSION when building the out-of-tree kmod,
-	# naturally, we must fetch the matching runtime package.
-	if [ -z "${NV_VERSION}" ]; then
-		echo "ERROR: flavor is zone-nvidiagpu but no NV_VERSION is set"
-		exit 1
-	fi
-
-	OLDDIR=$PWD
-
-	# TBH it's probably the same firmware regardless
-	if [ "${TARGET_ARCH_STANDARD}" = "aarch64" ]; then
-		NV_RUN_FILE="NVIDIA-Linux-aarch64-${NV_VERSION}.run"
-		NV_RUN_URL="https://us.download.nvidia.com/XFree86/aarch64/${NV_VERSION}/${NV_RUN_FILE}"
-	elif [ "${TARGET_ARCH_STANDARD}" = "x86_64" ]; then
-		NV_RUN_FILE="NVIDIA-Linux-x86_64-${NV_VERSION}.run"
-		NV_RUN_URL="https://us.download.nvidia.com/XFree86/Linux-x86_64/${NV_VERSION}/${NV_RUN_FILE}"
-	fi
-
-	NV_EXTRACT_PATH="$(mktemp -d)/extracted-${NV_VERSION}"
-	mkdir -p "$NV_EXTRACT_PATH"
-
-	echo "Downloading NVIDIA runtime package for driver ${NV_VERSION} from: $NV_RUN_URL"
-	curl --retry 5 --retry-delay 2 --retry-max-time 30 --retry-all-errors -L -o "$NV_EXTRACT_PATH/$NV_RUN_FILE" "$NV_RUN_URL"
-	chmod +x "$NV_EXTRACT_PATH/$NV_RUN_FILE"
-	"$NV_EXTRACT_PATH/$NV_RUN_FILE" -x --target "$NV_EXTRACT_PATH/out"
-	# Compress firmwares on-disk
-	# As of 6.x kernels the kconfig explicitly says you must use crc32 or none, not the default crc64.
-	xz -C crc32 "$NV_EXTRACT_PATH"/out/firmware/*
-	ls -lah "$NV_EXTRACT_PATH/out/firmware/"
-	NV_FW_PATH="${FIRMWARE_OUTPUT_PATH}/nvidia/$NV_VERSION"
-	mkdir -p "${NV_FW_PATH}"
-	cp "$NV_EXTRACT_PATH"/out/firmware/*.xz "${NV_FW_PATH}"
-
-	#
-	# Workload overlay
-	#
-	# Install various binaries necessary to make compute workloads run properly
-	# under containers.
-	#
-
-	create_links() {
-		local base_path="$1"
-		# create soname links
-		find "$base_path" -type f -name '*.so*' ! -path '*xorg/*' -print0 | while read -r -d $'\0' _lib; do
-			_soname=$(dirname "${_lib}")/$(readelf -d "${_lib}" | grep -Po 'SONAME.*: \[\K[^]]*' || true)
-			_base=$(echo "${_soname}" | sed -r 's/(.*)\.so.*/\1.so/')
-			[[ -e "${_soname}" ]] || ln -s "$(basename "${_lib}")" "${_soname}"
-			[[ -e "${_base}" ]] || ln -s "$(basename "${_soname}")" "${_base}"
-		done
-	}
-
-	multiarch_symlink_mirror() {
-		local base_path="$1"
-		local triplet="$2"
-		local src="${base_path}/usr/lib"
-		local dst="${base_path}/usr/lib/${triplet}"
-		pushd "$src" >/dev/null
-
-		# 1) Recreate directory structure (excluding the triplet subtree to avoid recursion)
-		# Use find from within src to get clean relative paths like "./gbm"
-
-		# Create directories
-		find . \
-			-mindepth 1 \
-			-path "./${triplet}/*" -prune -o \
-			-path "./${triplet}" -prune -o \
-			-type d -print0 |
-			while IFS= read -r -d '' rel_dir; do
-				# Strip leading "./"
-				rel_dir="${rel_dir#./}"
-				mkdir -p "${dst}/${rel_dir}"
-			done
-
-		# Helper: decide whether to mirror this path
-		should_mirror() {
-			local p="$1"
-			if [[ -f "$p" ]]; then
-				# regular file -> yes
-				return 0
-			elif [[ -L "$p" ]]; then
-				# symlink -> only if it ultimately targets a regular file
-				# resolve; readlink -f fails (non-zero) on broken loops/broken links
-				local tgt
-				if ! tgt="$(readlink -f -- "$p" 2>/dev/null)"; then
-					return 1
-				fi
-				[[ -f "$tgt" ]] && return 0 || return 1
-			else
-				# sockets, fifos, device nodes, etc. -> no
-				return 1
-			fi
-		}
-
-		# 2) Mirror regular files and symlinks-to-regular-files
-		find . \
-			-path "./${triplet}/*" -prune -o \
-			\( -type f -o -type l \) \
-			-print0 |
-			while IFS= read -r -d '' rel_item; do
-				rel_item="${rel_item#./}"
-				# Filter: only files and symlinks that resolve to files
-				if ! should_mirror "$rel_item"; then
-					continue
-				fi
-
-				local link_path="${dst}/${rel_item}"
-				local link_dir
-				link_dir="$(dirname "$link_path")"
-				mkdir -p "$link_dir"
-
-				# Always point the mirror symlink to the *original path in /usr/lib*,
-				# preserving whether the original is a file or a symlink.
-				local target="../${rel_item}"
-
-				if [[ -L "$link_path" ]]; then
-					ln -snf "$target" "$link_path"
-				elif [[ -e "$link_path" ]]; then
-					echo "WARN: exists and not a symlink, skipping: $link_path" >&2
-				else
-					ln -s "$target" "$link_path"
-				fi
-			done
-
-		popd >/dev/null
-	}
+# shellcheck source-path=SCRIPTDIR source=firmware-nvidia.sh
+. "${KERNEL_DIR}/hack/build/firmware-nvidia.sh"
 
-	# GTK (used by nvidia-settings, which we don't install)
-	rm -f "$NV_EXTRACT_PATH/out"/libnvidia-gtk*
-
-	# Wayland/GBM
-	mkdir -p "$WORKLOAD_OVERLAY_PATH/usr/lib/gbm"
-	ln -s "../libnvidia-allocator.so.${NV_VERSION}" "$WORKLOAD_OVERLAY_PATH/usr/lib/gbm/nvidia-drm_gbm.so"
-
-	# DRI driver
-	install -Dm755 "$NV_EXTRACT_PATH/out/"nvidia_drv.so "$WORKLOAD_OVERLAY_PATH/usr/lib/xorg/modules/drivers/nvidia_drv.so"
-
-	# GLX extensions
-	install -Dm755 "$NV_EXTRACT_PATH/out/libglxserver_nvidia.so.${NV_VERSION}" "$WORKLOAD_OVERLAY_PATH/usr/lib/nvidia/xorg/libglxserver_nvidia.so.${NV_VERSION}"
-	ln -s "libglxserver_nvidia.so.${NV_VERSION}" "$WORKLOAD_OVERLAY_PATH/usr/lib/nvidia/xorg/libglxserver_nvidia.so.1"
-	ln -s "libglxserver_nvidia.so.${NV_VERSION}" "$WORKLOAD_OVERLAY_PATH/usr/lib/nvidia/xorg/libglxserver_nvidia.so"
-
-	# Remove already-installed libs, so we don't accidentally include them in
-	# filters below
-	rm -f "$NV_EXTRACT_PATH/out"/nvidia-drv.so
-	rm -f "$NV_EXTRACT_PATH/out"/libglxserver_*
-
-	# Remove unnecessary libraries
-	rm -f "$NV_EXTRACT_PATH/out"/libnvidia-wayland*
-
-	for LIBRARY in "$NV_EXTRACT_PATH/out/"lib*.so*; do
-		BN="$(basename "$LIBRARY")"
-		install -Dm755 "$LIBRARY" "$WORKLOAD_OVERLAY_PATH/usr/lib/$BN"
-	done
-
-	for BINARY in "$NV_EXTRACT_PATH/out/"nvidia-{cuda-mps-control,cuda-mps-server,debugdump,pcc,smi}; do
-		BN="$(basename "$BINARY")"
-		install -Dm755 "$BINARY" "$WORKLOAD_OVERLAY_PATH/usr/bin/$BN"
-	done
-
-	for MAN1 in "$NV_EXTRACT_PATH/out/"nvidia-{cuda-mps-control,smi}.1.gz; do
-		BN="$(basename "$MAN1")"
-		install -Dm644 "$MAN1" "$WORKLOAD_OVERLAY_PATH/usr/share/man/man1/$BN"
-	done
-
-	# Install ICD loaders for OpenGL, Vulkan, and Vulkan SC
-	install -Dm644 "$NV_EXTRACT_PATH/out/"10_nvidia.json "$WORKLOAD_OVERLAY_PATH/usr/share/glvnd/egl_vendor.d/10_nvidia.json"
-	install -Dm644 "$NV_EXTRACT_PATH/out/"nvidia.icd "$WORKLOAD_OVERLAY_PATH/etc/OpenCL/vendors/nvidia.icd"
-	install -Dm644 "$NV_EXTRACT_PATH/out/"nvidia_icd.json "$WORKLOAD_OVERLAY_PATH/usr/share/vulkan/icd.d/nvidia_icd.json"
-	install -Dm644 "$NV_EXTRACT_PATH/out/"nvidia_layers.json "$WORKLOAD_OVERLAY_PATH/usr/share/vulkan/implicit_layer.d/nvidia_layers.json"
-	install -Dm644 "$NV_EXTRACT_PATH/out/"nvidia_icd_vksc.json "$WORKLOAD_OVERLAY_PATH/usr/share/vulkansc/icd.d/nvidia_icd_vksc.json"
-
-	create_links "$WORKLOAD_OVERLAY_PATH"
-
-	# For Debian-like distributions
-	multiarch_symlink_mirror "$WORKLOAD_OVERLAY_PATH" x86_64-linux-gnu
-
-	#
-	# Create NVIDIA persistence mode hook, ensures the driver is loaded and
-	# initialized, and always ready to run a workload
-	#
-	NVIDIA_BOOTSTRAP_OVERLAY_PATH="$ADDONS_OUTPUT_PATH/overlays/nvidia-bootstrap"
-
-	mkdir -p "$NVIDIA_BOOTSTRAP_OVERLAY_PATH"
-
-	# shellcheck disable=SC2043 # one entry today; loop mirrors the sibling install blocks
-	for BINARY in "$NV_EXTRACT_PATH/out/"nvidia-smi; do
-		BN="$(basename "$BINARY")"
-		install -Dm755 "$BINARY" "$NVIDIA_BOOTSTRAP_OVERLAY_PATH/usr/bin/$BN"
-	done
-
-	# HACK: Just using the builder's glibc binaries so we can run nvidia-smi
-	# without a full workload image
-	for LIBRARY in "$NV_EXTRACT_PATH/out/"libnvidia-ml.so* /lib64/ld-linux-x86-64.so.2 /lib/x86_64-linux-gnu/lib{pthread,m,dl,c,rt}.so.*; do
-		BN="$(basename "$LIBRARY")"
-		install -Dm755 "$LIBRARY" "$NVIDIA_BOOTSTRAP_OVERLAY_PATH/usr/lib/$BN"
-	done
-	ln -s usr/lib "$NVIDIA_BOOTSTRAP_OVERLAY_PATH/lib64"
-
-	create_links "$NVIDIA_BOOTSTRAP_OVERLAY_PATH"
-	multiarch_symlink_mirror "$NVIDIA_BOOTSTRAP_OVERLAY_PATH" x86_64-linux-gnu
-
-	mkdir -p "$ADDONS_OUTPUT_PATH/hooks"
-	cat >"$ADDONS_OUTPUT_PATH/hooks/nvidia-persist.toml" <<-EOF
-		[[hooks.setup]]
-		overlay = "nvidia-bootstrap"
-		modules = ["nvidia", "nvidia_drm", "nvidia_uvm"]
-		execute = ["/usr/bin/nvidia-smi", "-pm", "1"]
-		ignore-failure = true
-
-		[[hooks.hotplug]]
-		overlay = "nvidia-bootstrap"
-		modules = ["nvidia", "nvidia_drm", "nvidia_uvm"]
-		execute = ["/usr/bin/nvidia-smi", "-pm", "1"]
-		ignore-failure = true
-	EOF
-
-	cd "$OLDDIR"
-fi
+# shellcheck source-path=SCRIPTDIR source=firmware-tpu.sh
+. "${KERNEL_DIR}/hack/build/firmware-tpu.sh"
diff --git a/hack/build/refresh-nvidia-versions.py b/hack/build/refresh-nvidia-versions.py
index ce5ea5bf..aa228b47 100755
--- a/hack/build/refresh-nvidia-versions.py
+++ b/hack/build/refresh-nvidia-versions.py
@@ -54,7 +54,7 @@ def fetch_latest_versions() -> dict[str, str]:
             # Version is a mangled form (595.1040 for 595.104.02); the page
             # itself prefers DisplayVersion.
             version = info["DisplayVersion"]
-        except (AssertionError, KeyError, IndexError, TypeError):
+        except AssertionError, KeyError, IndexError, TypeError:
             raise RuntimeError(
                 "Unexpected NVIDIA lookup response for %r: %s" % (label, data)
             )
diff --git a/hooks/nvidia-persist.toml b/hooks/nvidia-persist.toml
new file mode 100644
index 00000000..1708846d
--- /dev/null
+++ b/hooks/nvidia-persist.toml
@@ -0,0 +1,11 @@
+[[hooks.setup]]
+overlay = "nvidia-bootstrap"
+modules = ["nvidia", "nvidia_drm", "nvidia_uvm"]
+execute = ["/usr/bin/nvidia-smi", "-pm", "1"]
+ignore-failure = true
+
+[[hooks.hotplug]]
+overlay = "nvidia-bootstrap"
+modules = ["nvidia", "nvidia_drm", "nvidia_uvm"]
+execute = ["/usr/bin/nvidia-smi", "-pm", "1"]
+ignore-failure = true
diff --git a/hooks/tpu-vfio.toml b/hooks/tpu-vfio.toml
new file mode 100644
index 00000000..068e2073
--- /dev/null
+++ b/hooks/tpu-vfio.toml
@@ -0,0 +1,19 @@
+[[hooks.setup]]
+sysfs-module-parameters = [
+  { module = "vfio_iommu_type1", parameter = "allow_unsafe_interrupts", value = "Y" },
+]
+
+[[hooks.setup.bind-driver]]
+driver = "vfio-pci"
+pci.vendor = "0x1ae0"
+pci.class = "0x1200"
+
+[[hooks.hotplug]]
+sysfs-module-parameters = [
+  { module = "vfio_iommu_type1", parameter = "allow_unsafe_interrupts", value = "Y" },
+]
+
+[[hooks.hotplug.bind-driver]]
+driver = "vfio-pci"
+pci.vendor = "0x1ae0"
+pci.class = "0x1200"