build: split NVIDIA rootfs into base-nvidia image + gpu extension

Carve the monolithic NVIDIA GPU rootfs into a driver-agnostic base-nvidia image (NVRC init + agent + base libs + in-tree modules) and a driver-versioned gpu extension (GPU userspace, configs, firmware, NVIDIA modules) laid out for /run/kata-extensions/gpu. Both halves are produced from the same chiseled tree via a partition-after step, so the monolith build stays byte-identical and continues to ship. The shared stage-one (driver install) is reused across the monolith, base and extension variants; base-nvidia is cached without the driver/ctk versions so one base image can back multiple driver extensions. Signed-off-by: Fabiano Fidêncio <ffidencio@nvidia.com> Assisted-by: Cursor <cursoragent@cursor.com>
2026-07-09 22:24:51 +00:00 · 2026-06-11 07:13:18 +02:00
parent f16b2fff3e
commit 522ca781d7
9 changed files with 406 additions and 39 deletions
--- a/tools/osbuilder/rootfs-builder/nvidia/nvidia_chroot.sh
+++ b/tools/osbuilder/rootfs-builder/nvidia/nvidia_chroot.sh
@@ -74,6 +74,13 @@ install_userspace_components() {

 	apt-mark hold cryptsetup-bin dmsetup libargon2-1     \
 		e2fsprogs libxml2
+
+	# NVRC loads the NVIDIA driver modules from the gpu extension's self-contained
+	# module tree via `modprobe --dirname <extension>`, a kmod feature the base
+	# busybox lacks. Install the real kmod here so the base-nvidia chisel can
+	# pull /usr/bin/kmod (and its libzstd/liblzma deps) from this stage.
+	eval "${APT_INSTALL}" kmod
+	apt-mark hold kmod
 }

 setup_apt_repositories() {
--- a/tools/osbuilder/rootfs-builder/nvidia/nvidia_rootfs.sh
+++ b/tools/osbuilder/rootfs-builder/nvidia/nvidia_rootfs.sh
@@ -37,7 +37,29 @@ else
    die "Unsupported architecture: ${machine_arch}"
 fi

-readonly stage_one="${BUILD_DIR:?}/rootfs-${BUILD_VARIANT:?}-stage-one"
+# The base-nvidia and gpu-extension images are carved out of the very same chiseled
+# tree as the monolith, so they share the (expensive) driver stage-one with the
+# monolithic "nvidia-gpu" build instead of rebuilding it per layout.
+nvidia_stage_one_variant() {
+	case "${BUILD_VARIANT}" in
+		nvidia-gpu-base|nvidia-gpu-extension) echo "nvidia-gpu" ;;
+		*) echo "${BUILD_VARIANT}" ;;
+	esac
+}
+
+readonly stage_one="${BUILD_DIR:?}/rootfs-$(nvidia_stage_one_variant)-stage-one"
+
+# Image layout produced from the chiseled tree:
+#   monolith    - the full GPU image (default; unchanged behaviour)
+#   base        - driver-agnostic base-nvidia (NVRC init + agent + base libs)
+#   gpu-extension   - GPU userspace only, laid out for /run/kata-extensions/gpu
+nvidia_image_layout() {
+	case "${BUILD_VARIANT}" in
+		nvidia-gpu-base) echo "base" ;;
+		nvidia-gpu-extension) echo "gpu-extension" ;;
+		*) echo "monolith" ;;
+	esac
+}

 setup_nvidia-nvrc() {
 	# NVRC is built from a pinned git ref by the Kata "nvrc" static-build
@@ -47,7 +69,11 @@ setup_nvidia-nvrc() {
 	local nvrc_tarball="${BUILD_DIR}/kata-static-nvrc.tar.zst"
 	[[ -e "${nvrc_tarball}" ]] || \
 		die "NVRC tarball not found: ${nvrc_tarball} (build the 'nvrc' target first)"
-	tar --zstd -xvf "${nvrc_tarball}" -C .
+	# The tarball carries a ./bin/ entry; on a usr-merged rootfs /bin is a
+	# symlink to usr/bin, so extract with --keep-directory-symlink to follow it
+	# instead of clobbering the symlink with a real (near-empty) directory,
+	# which would hide /bin/bash and break the driver chroot below.
+	tar --keep-directory-symlink --zstd -xvf "${nvrc_tarball}" -C .
 }

 setup_nvidia_gpu_rootfs_stage_one() {
@@ -238,9 +264,14 @@ chisseled_init() {
 	echo "nvidia: chisseling init"
 	tar --zstd -xvf "${BUILD_DIR}"/kata-static-busybox.tar.zst -C .

+	# Create bin/ and sbin/ explicitly rather than relying on busybox's
+	# `make install` to emit them: busybox only creates a directory when it has
+	# an applet living there, and once the nvidia busybox dropped modprobe (its
+	# only /sbin applet, now provided by kmod) it stopped emitting sbin/, which
+	# broke the `sbin/init` symlink below.
 	mkdir -p dev etc proc run/cdi sys tmp usr var lib/modules lib/firmware \
 		 usr/share/nvidia lib/"${machine_arch}"-linux-gnu lib64        \
-		 usr/bin etc/modprobe.d etc/ssl/certs
+		 bin sbin usr/bin etc/modprobe.d etc/ssl/certs

 	ln -sf ../run var/run
 	ln -sf ../run var/log
@@ -311,7 +342,11 @@ compress_rootfs() {
 	[[ ${machine_arch} == "aarch64" ]] && libdir="lib"
 	[[ ${machine_arch} == "x86_64" ]]  && libdir="lib64"

-	chmod +x "${libdir}"/ld-linux-*
+	# The gpu-extension layout ships no program interpreter (it lives in the
+	# base-nvidia image), so only fix up the loader when it is actually present.
+	if compgen -G "${libdir}/ld-linux-*" > /dev/null; then
+		chmod +x "${libdir}"/ld-linux-*
+	fi
 }

 copy_cdh_runtime_deps() {
@@ -391,10 +426,240 @@ coco_guest_components() {
 	copy_cdh_runtime_deps
 }

+# GPU userspace owned by the gpu extension. Anything not listed here stays in the
+# driver-agnostic base-nvidia image. Paths are relative to the chiseled rootfs.
+readonly nvidia_gpu_extension_bins=(
+	bin/nvidia-smi
+	bin/nvidia-ctk
+	bin/nvidia-cdi-hook
+	bin/nvidia-persistenced
+	bin/nv-hostengine
+	bin/nv-fabricmanager
+	sbin/nvlsm
+)
+
+# GPU shared-library globs (inside the multiarch lib dir, plus libgrpc_mgr in
+# /lib) owned by the gpu extension.
+readonly nvidia_gpu_extension_lib_globs=(
+	'libnv*'
+	'libcuda.so*'
+	'libdcgm.*'
+	'libnvidia-nscq.so*'
+)
+
+# Lay the GPU userspace out for the extension mount (/run/kata-extensions/gpu) and drop
+# everything else. Matches how NVRC consumes the extension: bins from <root>/bin and
+# <root>/sbin, libraries via LD_LIBRARY_PATH=<root>/usr/lib, kernel modules via
+# `modprobe --dirname <root>`, configs from <root>/usr/share/nvidia, and
+# <root>/lib/firmware/nvidia bind-mounted onto /lib/firmware/nvidia. Runs inside
+# the (full, chiseled) ${stage_two}/${ROOTFS_DIR}.
+#
+# The GPU shared libraries go under <root>/usr/lib (not a flat <root>/lib) so
+# that NVRC can run `nvidia-ctk cdi generate --driver-root=<root>`: nvidia-ctk
+# records the in-container mount path as the host path with the driver root
+# stripped, so libraries at <root>/usr/lib land at the canonical /usr/lib inside
+# the container. Apps that scan /usr for the driver (e.g. NVIDIA NIM, which
+# bails with "libnvidia-ml.so.1 not found under /usr") then find it; a flat
+# <root>/lib would strip to /lib and hide the driver from those checks.
+partition_gpu_extension() {
+	echo "nvidia: building gpu extension layout"
+
+	local extension
+	extension="$(mktemp -d "${BUILD_DIR}/.nvidia-gpu-extension.XXXX")"
+	mkdir -p "${extension}/bin" "${extension}/sbin" "${extension}/usr/lib" \
+		 "${extension}/usr/share/nvidia" "${extension}/lib/firmware" "${extension}/lib/modules"
+
+	local f
+	for f in "${nvidia_gpu_extension_bins[@]}"; do
+		[[ -e "${f}" ]] && install -D -m0755 "${f}" "${extension}/${f}"
+	done
+
+	# Collect the GPU shared libraries under <root>/usr/lib so both NVRC's
+	# LD_LIBRARY_PATH and `nvidia-ctk --driver-root=<root>` resolve them and the
+	# container sees them at /usr/lib (see header); libc/loader stay in the base.
+	local md="lib/${machine_arch}-linux-gnu"
+	local g
+	for g in "${nvidia_gpu_extension_lib_globs[@]}"; do
+		find "${md}" -maxdepth 1 -name "${g}" -exec cp -a {} "${extension}/usr/lib/" \;
+	done
+	[[ -e lib/libgrpc_mgr.so ]] && cp -a lib/libgrpc_mgr.so "${extension}/usr/lib/"
+
+	# Materialize the SONAME symlinks (e.g. libcuda.so.1 -> libcuda.so.595.58.03)
+	# inside the lib dir. The extension ships only the versioned files, so without
+	# this `nvidia-ctk cdi generate` has no symlink to replicate into the
+	# container (it reproduces existing links, it does not synthesize SONAMEs) and
+	# the container can't resolve libcuda.so.1 -> it then falls back to the image's
+	# older cuda-compat libcuda and CUDA fails with "driver version is insufficient
+	# for CUDA runtime version". `ldconfig -n` only creates the versioned symlinks
+	# in the given dir (no cache, no chroot), which is all the loader-less extension
+	# needs. The monolith gets the same links via the `chroot . ldconfig` below.
+	ldconfig -n "${extension}/usr/lib"
+
+	# GPU configs (fabricmanager.cfg, nvlsm.conf, ...).
+	[[ -d usr/share/nvidia ]] && cp -a usr/share/nvidia/. "${extension}/usr/share/nvidia/"
+
+	# GPU firmware (GSP, ...); NVRC binds this onto /lib/firmware/nvidia.
+	[[ -d lib/firmware/nvidia ]] && cp -a lib/firmware/nvidia "${extension}/lib/firmware/"
+
+	# Ship a self-contained module tree so `modprobe --dirname <root>` resolves
+	# the NVIDIA modules and their dependencies.
+	cp -a lib/modules/. "${extension}/lib/modules/"
+	if command -v depmod >/dev/null 2>&1; then
+		local kdir kver
+		for kdir in "${extension}"/lib/modules/*/; do
+			[[ -d "${kdir}" ]] || continue
+			kver="$(basename "${kdir}")"
+			depmod -b "${extension}" "${kver}" || true
+		done
+	fi
+
+	# Replace the rootfs with the extension-only content.
+	find . -mindepth 1 -maxdepth 1 -exec rm -rf {} +
+	cp -a "${extension}/." .
+	rm -rf "${extension}"
+}
+
+# Strip the GPU userspace from the chiseled tree, leaving a driver-agnostic
+# base-nvidia: NVRC init + kata-agent + busybox + loader/libc. No kernel
+# modules are shipped (see below). The empty /lib/firmware/nvidia directory is
+# kept as the bind mountpoint NVRC uses for the extension firmware. Runs inside
+# ${ROOTFS_DIR}.
+partition_base() {
+	echo "nvidia: building driver-agnostic base layout"
+
+	local f
+	for f in "${nvidia_gpu_extension_bins[@]}"; do
+		rm -f "${f}"
+	done
+
+	local md="lib/${machine_arch}-linux-gnu"
+	local g
+	for g in "${nvidia_gpu_extension_lib_globs[@]}"; do
+		find "${md}" -maxdepth 1 -name "${g}" -delete
+	done
+	rm -f lib/libgrpc_mgr.so
+
+	# GPU configs live in the extension; keep usr/share/nvidia as an empty stub.
+	rm -rf usr/share/nvidia
+	mkdir -p usr/share/nvidia
+
+	# Keep /lib/firmware/nvidia as an empty mountpoint for NVRC's firmware bind.
+	rm -rf lib/firmware/nvidia
+	mkdir -p lib/firmware/nvidia
+
+	# Ship no kernel modules in the base: the NVIDIA driver modules are
+	# GPU-specific and live in the gpu extension (NVRC loads them via
+	# `modprobe --dirname <extension>`), and the remaining in-tree dependencies
+	# (mlx5, infiniband, ...) are built into the NVIDIA kernel. Keeping
+	# /lib/modules empty is what makes the base driver-agnostic and reusable
+	# across driver versions.
+	rm -rf lib/modules
+	mkdir -p lib/modules
+}
+
+# NVRC opens every cold-plugged extension (the gpu extension always, and the coco extension
+# on confidential guests) as a dm-verity device by exec'ing /usr/sbin/veritysetup
+# before mounting it. The base-nvidia image is the one that boots and runs NVRC,
+# so it must carry veritysetup and its shared-library closure unconditionally -
+# regardless of whether the guest is confidential. This closure is also exactly
+# what cryptsetup links, so the cryptsetup binary shipped in the coco extension
+# (encrypted storage, see build-static-coco-guest-components.sh) resolves its
+# libraries against the base without bundling any of its own. Runs inside
+# ${ROOTFS_DIR}.
+chisseled_veritysetup() {
+	echo "nvidia: chisseling veritysetup"
+
+	local libdir="lib/${machine_arch}-linux-gnu"
+
+	# NVRC execs the absolute path /usr/sbin/veritysetup; cryptsetup-bin is
+	# installed in the (usr-merged) stage-one, so /sbin/veritysetup resolves to
+	# the real binary there.
+	mkdir -p usr/sbin
+	cp -a "${stage_one}/sbin/veritysetup" usr/sbin/.
+
+	# veritysetup -> libcryptsetup runtime closure (same set cryptsetup links).
+	cp -a "${stage_one}/${libdir}"/libcryptsetup.so.12*    "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libpopt.so.0*           "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libuuid.so.1*           "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libblkid.so.1*          "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libdevmapper.so.1.02.1* "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libcrypto.so.3*         "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libargon2.so.1*         "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libjson-c.so.5*         "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libselinux.so.1*        "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libudev.so.1*           "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libpcre2-8.so.0*        "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libcap.so.2*            "${libdir}/."
+}
+
+# NVRC loads each extension's NVIDIA kernel modules from that extension's self-contained
+# module tree via `modprobe --dirname <extension>` (a kmod feature). The base ships
+# busybox, whose modprobe is built without long options and has no --dirname, so
+# the base must carry the real kmod. This keeps modules composable: every
+# module-bearing extension stays independent (its own modules.dep from depmod -b) and
+# nothing has to shadow the read-only /lib/modules. kmod is a single multi-call
+# binary that embeds libkmod (no libkmod2 to ship); modprobe/insmod/... are
+# argv[0] symlinks to it. Runs inside ${ROOTFS_DIR}.
+chisseled_kmod() {
+	echo "nvidia: chisseling kmod"
+
+	local libdir="lib/${machine_arch}-linux-gnu"
+
+	cp -a "${stage_one}/usr/bin/kmod" usr/bin/.
+
+	# kmod picks its applet from argv[0]. Expose the module tools as symlinks;
+	# /sbin/modprobe (NVRC's absolute path) shadows the busybox modprobe applet.
+	# Absolute targets so they resolve regardless of whether /sbin is a real dir
+	# or a usr-merge symlink.
+	local tool
+	for tool in modprobe insmod rmmod depmod lsmod; do
+		ln -sf /usr/bin/kmod "sbin/${tool}"
+	done
+
+	# kmod links libzstd/liblzma (compressed-module support) and libcrypto
+	# (signed-module support) unconditionally; our modules are uncompressed and
+	# unsigned but the NEEDED entries must still resolve. libc is already present
+	# (from chisseled_compute). libcrypto is also pulled by chisseled_veritysetup
+	# on the base, but the monolith runs kmod without veritysetup, so copy it
+	# here to keep this function self-contained (cp -a is idempotent for base).
+	cp -a "${stage_one}/${libdir}"/libzstd.so.1*    "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/liblzma.so.5*    "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libcrypto.so.3*  "${libdir}/."
+}
+
+# CDH's secure_mount and the agent's plain ephemeral-storage handlers format
+# scratch volumes with mke2fs/mkfs.ext4 (and zero/size their backing with dd)
+# before mounting them. This tooling backs *unencrypted* storage too, so unlike
+# cryptsetup (encrypted storage, which is CoCo-only and therefore lives in the
+# coco extension) the base-nvidia image must carry it unconditionally - just like
+# veritysetup. libblkid/libuuid/libc are already provided by
+# chisseled_veritysetup/chisseled_compute; only the e2fsprogs libs are new here.
+# Runs inside ${ROOTFS_DIR}.
+chisseled_storage() {
+	echo "nvidia: chisseling storage tooling (mke2fs/mkfs.ext4/dd)"
+
+	local libdir="lib/${machine_arch}-linux-gnu"
+
+	# e2fsprogs (mke2fs/mkfs.ext4) runtime libs not already shipped by
+	# chisseled_veritysetup.
+	cp -a "${stage_one}/${libdir}"/libext2fs.so.2*  "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libcom_err.so.2* "${libdir}/."
+	cp -a "${stage_one}/${libdir}"/libe2p.so.2*     "${libdir}/."
+
+	# mkfs.ext4 is a symlink to mke2fs in e2fsprogs; cp -a preserves it.
+	mkdir -p sbin etc bin
+	cp -a "${stage_one}/sbin/mke2fs"     sbin/.
+	cp -a "${stage_one}/sbin/mkfs.ext4"  sbin/.
+	cp -a "${stage_one}/etc/mke2fs.conf" etc/.
+	cp -a "${stage_one}/usr/bin/dd"      bin/.
+}
+
 setup_nvidia_gpu_rootfs_stage_two() {
 	readonly stage_two="${ROOTFS_DIR:?}"
 	readonly stack="${NVIDIA_GPU_STACK:?}"
 	readonly type=${1:-""}
+	local layout
+	layout="$(nvidia_image_layout)"

 	# If devkit flag is set, skip chisseling, use stage_one
 	if echo "${stack}" | grep -q '\<devkit\>'; then
@@ -416,6 +681,16 @@ setup_nvidia_gpu_rootfs_stage_two() {

 		pushd "${stage_two}" >> /dev/null

+		# stage-one archives the full base+driver tree with `tar
+		# --remove-files`, emptying ${stage_two} on its first run. When
+		# stage-one is served from cache that emptying never happens, so the
+		# freshly-built distro rootfs is still here. The chisel assembles a
+		# minimal tree purely from ${stage_one}, so wipe any leftover distro
+		# content first to keep the result deterministic regardless of whether
+		# stage-one was cached (otherwise base/monolith layouts bloat with the
+		# full Ubuntu rootfs).
+		find . -mindepth 1 -maxdepth 1 -exec rm -rf {} +
+
 		chisseled_init
 		chisseled_iptables

@@ -439,11 +714,30 @@ setup_nvidia_gpu_rootfs_stage_two() {

 		coco_guest_components
 		chisseled_nvat
+
+		# Carve the freshly chiseled (monolith) tree into the requested layout.
+		# The monolith path is left untouched.
+		case "${layout}" in
+			base) partition_base; chisseled_veritysetup; chisseled_storage; chisseled_kmod ;;
+			gpu-extension) partition_gpu_extension ;;
+			# The monolith boots and runs the whole GPU stack itself, so NVRC's
+			# `/sbin/modprobe nvidia` must resolve here too. busybox no longer
+			# ships a modprobe applet (see chisseled_init), so carry real kmod.
+			monolith) chisseled_kmod ;;
+		esac
 	fi

-	copy_mkfs_ext4_runtime_deps
+	# The mkfs/e2fsprogs storage tooling belongs in the images that actually
+	# boot and run the agent (the monolith and the base-nvidia image, the latter
+	# via chisseled_storage). The gpu extension ships GPU userspace only, and
+	# partition_gpu_extension has already replaced the tree with extension-only
+	# content (no lib/<arch>-linux-gnu), so skip it there.
+	[[ "${layout}" != "gpu-extension" ]] && copy_mkfs_ext4_runtime_deps
 	compress_rootfs
-	chroot . ldconfig
+	# The gpu extension has no loader/ldconfig of its own; its libraries are found
+	# via NVRC's LD_LIBRARY_PATH, so skip the ld.so cache rebuild there. Its
+	# SONAME symlinks are created with `ldconfig -n` in partition_gpu_extension().
+	[[ "${layout}" != "gpu-extension" ]] && chroot . ldconfig

 	popd >> /dev/null
 }
--- a/tools/osbuilder/rootfs-builder/rootfs.sh
+++ b/tools/osbuilder/rootfs-builder/rootfs.sh
@@ -932,7 +932,13 @@ main()
 	init="${ROOTFS_DIR}/sbin/init"
 	setup_rootfs

-	if [[ "${BUILD_VARIANT}" = "nvidia-gpu" ]]; then
+	# The base-nvidia and gpu-extension layouts are carved from the same chiseled
+	# tree as the monolith (which they share the driver stage-one with), so they
+	# drive the non-confidential stage-one/two and only differ in the final
+	# partition step (see nvidia_image_layout in nvidia_rootfs.sh).
+	if [[ "${BUILD_VARIANT}" = "nvidia-gpu" || \
+	      "${BUILD_VARIANT}" = "nvidia-gpu-base" || \
+	      "${BUILD_VARIANT}" = "nvidia-gpu-extension" ]]; then
 		setup_nvidia_gpu_rootfs_stage_one
 		setup_nvidia_gpu_rootfs_stage_two
 		return $?
--- a/tools/packaging/guest-image/build_image.sh
+++ b/tools/packaging/guest-image/build_image.sh
@@ -97,7 +97,11 @@ build_image() {
 		PAUSE_IMAGE_TARBALL="${PAUSE_IMAGE_TARBALL:-}" \
 		GUEST_HOOKS_TARBALL="${GUEST_HOOKS_TARBALL}"

-	if [[ "${image_initrd_suffix}" == "nvidia-gpu"* ]]; then
+	# The driver-versioned NVIDIA artefacts (the monolith, the confidential
+	# monolith and the gpu extension) carry the driver version in their filename so
+	# multiple driver builds can coexist.  The base-nvidia image is
+	# driver-agnostic, so it must not get the driver-version suffix.
+	if [[ "${image_initrd_suffix}" == "nvidia-gpu"* && "${image_initrd_suffix}" != "nvidia-gpu-base" ]]; then
 		nvidia_driver_version=$(get_from_kata_deps .externals.nvidia.driver.version)
 		artifact_name=${artifact_name/.image/"-${nvidia_driver_version}".image}
 	fi
--- a/tools/packaging/kata-deploy/local-build/kata-deploy-binaries.sh
+++ b/tools/packaging/kata-deploy/local-build/kata-deploy-binaries.sh
@@ -194,7 +194,7 @@ cleanup_and_fail_shim_v2_specifics() {
 	local extra_tarballs="${3:-}"
 	local tarball_dir="${repo_root_dir}/tools/packaging/kata-deploy/local-build/build"

-	for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential; do
+	for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-extension; do
 		local root_hash_file="${tarball_dir}/${component}-root_hash_${variant}.txt"
 		[[ -f "${root_hash_file}" ]] && rm -f "${root_hash_file}"
 	done
@@ -224,7 +224,7 @@ install_cached_shim_v2_tarball_get_root_hash() {
 	local tarball_dir="${repo_root_dir}/tools/packaging/kata-deploy/local-build/build"
 	local root_hash_basedir="./opt/kata/share/kata-containers/"

-	for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential; do
+	for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-extension; do
 		# The measured base image ships as kata-static-rootfs-image.tar.zst
 		# (no variant suffix), but carries its root hash under the "base" label.
 		local image_conf_tarball
@@ -252,7 +252,7 @@ install_cached_shim_v2_tarball_compare_root_hashes() {
 	local found_any=""
 	local tarball_dir="${repo_root_dir}/tools/packaging/kata-deploy/local-build/build"

-	for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential; do
+	for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-extension; do
 		local image_root_hash="${tarball_dir}/root_hash_${variant}.txt"
 		local cached_root_hash="${component}-root_hash_${variant}.txt"

@@ -492,14 +492,26 @@ install_image() {
 		latest_artefact+="-$(get_latest_pause_image_artefact_and_builder_image_version)"
 	fi

-	if [[ "${variant}" == "nvidia-gpu" ]]; then
-		# If we bump the kernel we need to rebuild the image
+	if [[ "${variant}" == "nvidia-gpu" || "${variant}" == "nvidia-gpu-extension" ]]; then
+		# If we bump the kernel we need to rebuild the image.  The gpu extension
+		# carries the driver userspace carved out of the same chiseled tree,
+		# so it is driver-versioned just like the monolith.
 		latest_artefact+="-$(get_latest_kernel_nvidia_artefact_and_builder_image_version)"
 		latest_artefact+="-$(get_latest_nvidia_driver_version)"
 		latest_artefact+="-$(get_latest_nvidia_ctk_version)"
 		latest_artefact+="-$(get_latest_nvidia_nvrc_version)"
 	fi

+	if [[ "${variant}" == "nvidia-gpu-base" ]]; then
+		# The base-nvidia image strips all driver userspace and resets the
+		# kernel modules to in-tree only, so it is driver-agnostic: it depends
+		# on the NVIDIA kernel build and NVRC (its init), but not on the driver
+		# or container-toolkit versions.  That lets a single base image back
+		# multiple driver-specific gpu extensions.
+		latest_artefact+="-$(get_latest_kernel_nvidia_artefact_and_builder_image_version)"
+		latest_artefact+="-$(get_latest_nvidia_nvrc_version)"
+	fi
+
 	# The base guest image (empty variant) is built as a measured rootfs so
 	# that confidential configurations can dm-verity-protect it; non-confidential
 	# configurations simply boot its data partition unverified.  Reflect the
@@ -910,6 +922,48 @@ install_image_nvidia_gpu_confidential() {
 	install_image "nvidia-gpu-confidential"
 }

+# Install the driver-agnostic base-nvidia image.
+#
+# This is the NVRC-init half of the chiseled NVIDIA tree: NVRC, the kata-agent,
+# busybox, the base libc/loader and the in-tree kernel modules.  The driver
+# userspace is stripped out and shipped separately in the gpu extension, so this
+# image is reused across driver versions.  The driver still has to be installed
+# to build the shared stage-one (the GPU files are carved out afterwards), so we
+# keep the same NVIDIA_GPU_STACK as the monolith.
+install_image_nvidia_gpu_base() {
+	export AGENT_POLICY
+	export MEASURED_ROOTFS="yes"
+	export FS_TYPE="erofs"
+	export SKIP_DAX_HEADER="yes"
+	local version
+	version=$(get_latest_nvidia_driver_version)
+	EXTRA_PKGS="apt curl ${EXTRA_PKGS}"
+	NVIDIA_GPU_STACK=${NVIDIA_GPU_STACK:-"driver=${version},compute,dcgm,nvswitch"}
+	install_image "nvidia-gpu-base"
+}
+
+# Install the gpu extension image.
+#
+# This is the driver half of the chiseled NVIDIA tree, laid out for
+# /run/kata-extensions/gpu: the GPU binaries, libraries, configs, firmware and the
+# NVIDIA kernel modules.  It is an erofs+verity image (MEASURED_ROOTFS) and is
+# driver-versioned, so multiple driver extensions can coexist against a single
+# base-nvidia image.
+install_image_nvidia_gpu_extension() {
+	export AGENT_POLICY
+	export MEASURED_ROOTFS="yes"
+	export FS_TYPE="erofs"
+	export SKIP_DAX_HEADER="yes"
+	# The gpu extension is GPU-userspace-only content mounted into base-nvidia; it
+	# ships no /sbin/init, so skip the rootfs init/agent sanity check.
+	export SKIP_ROOTFS_CHECK="yes"
+	local version
+	version=$(get_latest_nvidia_driver_version)
+	EXTRA_PKGS="apt curl ${EXTRA_PKGS}"
+	NVIDIA_GPU_STACK=${NVIDIA_GPU_STACK:-"driver=${version},compute,dcgm,nvswitch"}
+	install_image "nvidia-gpu-extension"
+}
+
 install_se_image() {
 	# shellcheck disable=SC2154
 	info "Create IBM SE image configured with AA_KBC=${AA_KBC}"
@@ -1261,7 +1315,7 @@ install_nydus() {
 # Shared helper: extract measured-rootfs root hashes from confidential image tarballs.
 # These are needed by the Rust runtime (runtime-rs) at build time for dm-verity.
 _collect_root_hashes() {
-	for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential; do
+	for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-extension; do
 		# The measured base image ships as kata-static-rootfs-image.tar.zst
 		# (no variant suffix), but carries its root hash under the "base" label.
 		local tarball_glob="kata-static-rootfs-image-${variant}.tar.zst"
@@ -1744,6 +1798,10 @@ handle_build() {

 	rootfs-image-nvidia-gpu-confidential) install_image_nvidia_gpu_confidential ;;

+	rootfs-image-nvidia-gpu-base) install_image_nvidia_gpu_base ;;
+
+	rootfs-image-nvidia-gpu-extension) install_image_nvidia_gpu_extension ;;
+
 	rootfs-cca-confidential-image) install_image_confidential ;;

 	rootfs-cca-confidential-initrd) install_initrd_confidential ;;
@@ -1803,7 +1861,7 @@ handle_build() {
 			;;
 		shim-v2-go|shim-v2-rust)
 			if [[ "${MEASURED_ROOTFS}" == "yes" ]]; then
-				for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential; do
+				for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-extension; do
 					[[ -f "${workdir}/root_hash_${variant}.txt" ]] && mv "${workdir}/root_hash_${variant}.txt" "${workdir}/${build_target}-root_hash_${variant}.txt"
 				done
 			fi
@@ -1869,7 +1927,7 @@ handle_build() {
 		shim-v2-go|shim-v2-rust)
 			if [[ "${MEASURED_ROOTFS}" == "yes" ]]; then
 				local found_any=""
-				for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential; do
+				for variant in base confidential coco-extension nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-extension; do
 					# The variants could be built independently we need to check if
 					# they exist and then push them to the registry
 					[[ -f "${workdir}/${build_target}-root_hash_${variant}.txt" ]] && files_to_push+=("${build_target}-root_hash_${variant}.txt") && found_any="yes"
--- a/tools/packaging/kernel/configs/fragments/gpu/nvidia.arm64.conf.in
+++ b/tools/packaging/kernel/configs/fragments/gpu/nvidia.arm64.conf.in
@@ -27,11 +27,3 @@ CONFIG_ARM_SMMU_V3_SVA=y
 CONFIG_CRYPTO_ECC=y
 CONFIG_CRYPTO_ECDH=y
 CONFIG_CRYPTO_ECDSA=y
-
-# HGX/DGX platform
-CONFIG_INFINIBAND=m
-CONFIG_INFINIBAND_USER_MAD=m
-
-CONFIG_NET_VENDOR_MELLANOX=y
-CONFIG_MLX5_CORE=m
-CONFIG_MLX5_INFINIBAND=m
--- a/tools/packaging/kernel/configs/fragments/gpu/nvidia.x86_64.conf.in
+++ b/tools/packaging/kernel/configs/fragments/gpu/nvidia.x86_64.conf.in
@@ -26,11 +26,3 @@ CONFIG_CRYPTO_ECDSA=y

 # Dependency of _CRYPTO_
 CONFIG_MODULE_SIG=y
-
-# HGX/DGX platform
-CONFIG_INFINIBAND=m
-CONFIG_INFINIBAND_USER_MAD=m
-
-CONFIG_NET_VENDOR_MELLANOX=y
-CONFIG_MLX5_CORE=m
-CONFIG_MLX5_INFINIBAND=m
--- a/tools/packaging/static-build/busybox/busybox.nvidia.conf
+++ b/tools/packaging/static-build/busybox/busybox.nvidia.conf
@@ -564,14 +564,18 @@ CONFIG_FEATURE_DEFAULT_PASSWD_ALGO=""
 # CONFIG_LSMOD is not set
 # CONFIG_FEATURE_LSMOD_PRETTY_2_6_OUTPUT is not set
 # CONFIG_MODINFO is not set
-CONFIG_MODPROBE=y
-CONFIG_FEATURE_MODPROBE_BLACKLIST=y
+# CONFIG_MODPROBE is not set
+# CONFIG_FEATURE_MODPROBE_BLACKLIST is not set
 # CONFIG_RMMOD is not set

 #
 # Options common to multiple modutils
 #
-CONFIG_FEATURE_CMDLINE_MODULE_OPTIONS=y
+# The base-nvidia image ships the real kmod (modprobe/insmod/...) so that NVRC
+# can load each extension's modules with `modprobe --dirname <extension>`; busybox's
+# modprobe lacks long options (no --dirname), so all busybox modutils are
+# disabled here to avoid a dead, conflicting applet shadowing kmod.
+# CONFIG_FEATURE_CMDLINE_MODULE_OPTIONS is not set
 # CONFIG_FEATURE_MODPROBE_SMALL_CHECK_ALREADY_LOADED is not set
 # CONFIG_FEATURE_2_4_MODULES is not set
 # CONFIG_FEATURE_INSMOD_VERSION_CHECKING is not set
@@ -581,10 +585,8 @@ CONFIG_FEATURE_CMDLINE_MODULE_OPTIONS=y
 # CONFIG_FEATURE_INSMOD_LOAD_MAP_FULL is not set
 # CONFIG_FEATURE_CHECK_TAINTED_MODULE is not set
 # CONFIG_FEATURE_INSMOD_TRY_MMAP is not set
-CONFIG_FEATURE_MODUTILS_ALIAS=y
-CONFIG_FEATURE_MODUTILS_SYMBOLS=y
-CONFIG_DEFAULT_MODULES_DIR="/lib/modules"
-CONFIG_DEFAULT_DEPMOD_FILE="modules.dep"
+# CONFIG_FEATURE_MODUTILS_ALIAS is not set
+# CONFIG_FEATURE_MODUTILS_SYMBOLS is not set

 #
 # Linux System Utilities
--- a/versions.yaml
+++ b/versions.yaml
@@ -137,6 +137,12 @@ assets:
        nvidia-gpu-confidential:
          name: "ubuntu"
          version: "noble"  # 24.04 LTS
+        nvidia-gpu-base:
+          name: "ubuntu"
+          version: "noble"  # 24.04 LTS
+        nvidia-gpu-extension:
+          name: "ubuntu"
+          version: "noble"  # 24.04 LTS
      ppc64le:
        name: "ubuntu"
        version: "noble"  # 24.04 LTS
@@ -161,6 +167,12 @@ assets:
        nvidia-gpu-confidential:
          name: "ubuntu"
          version: "noble"  # 24.04 LTS
+        nvidia-gpu-base:
+          name: "ubuntu"
+          version: "noble"  # 24.04 LTS
+        nvidia-gpu-extension:
+          name: "ubuntu"
+          version: "noble"  # 24.04 LTS

  initrd:
    description: |