diff --git a/images/tpu/jax/Dockerfile.x86_64 b/images/tpu/jax/Dockerfile.x86_64
index d1e088b1c..647e6c718 100644
--- a/images/tpu/jax/Dockerfile.x86_64
+++ b/images/tpu/jax/Dockerfile.x86_64
@@ -1,4 +1,4 @@
-FROM python:3.8
+FROM python:3.10
 
 RUN pip install --upgrade jax
 RUN pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
diff --git a/pkg/abi/tpu/tpu.go b/pkg/abi/tpu/tpu.go
index 358bcfff7..d42193be0 100644
--- a/pkg/abi/tpu/tpu.go
+++ b/pkg/abi/tpu/tpu.go
@@ -41,6 +41,9 @@ const (
 
 	// TPUV5eDeviceID is the PCI device ID of TPU V5e hardware.
 	TPUV5eDeviceID = 0x0063
+
+	// TPUV5pDeviceID is the PCI device ID of TPU V5p hardware.
+	TPUV5pDeviceID = 0x0062
 )
 
 // TPUV4InterruptsMap maps BAR indices to valid register offsets.
diff --git a/pkg/sentry/devices/tpuproxy/device.go b/pkg/sentry/devices/tpuproxy/device.go
index e6bdf9e8c..807d19847 100644
--- a/pkg/sentry/devices/tpuproxy/device.go
+++ b/pkg/sentry/devices/tpuproxy/device.go
@@ -15,8 +15,8 @@
 package tpuproxy
 
 import (
-	"fmt"
 	"path/filepath"
+	"strconv"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -32,6 +32,7 @@ import (
 const (
 	// VFIO_MINOR is the VFIO minor number from include/linux/miscdevice.h.
 	VFIO_MINOR = 196
+
 	// VFIOPath is the path to a VFIO device, it is usually used to
 	// construct a VFIO container.
 	VFIOPath = "/dev/vfio/vfio"
@@ -46,7 +47,10 @@ const (
 type tpuDevice struct {
 	mu sync.Mutex
 
+	// minor is the device minor number.
 	minor uint32
+	// num is the number of the device in the dev filesystem (e.g /dev/vfio/0).
+	num uint32
 }
 
 // Open implements vfs.Device.Open.
@@ -58,10 +62,10 @@ func (dev *tpuDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, o
 	}
 	dev.mu.Lock()
 	defer dev.mu.Unlock()
-	devName := fmt.Sprintf("vfio/%d", dev.minor)
+	devName := filepath.Join("vfio", strconv.Itoa(int(dev.num)))
 	hostFD, err := devClient.OpenAt(ctx, devName, opts.Flags)
 	if err != nil {
-		ctx.Warningf("vfioDevice: failed to open host %s: %v", devName, err)
+		ctx.Warningf("tpuDevice: failed to open host %s: %v", devName, err)
 		return nil, err
 	}
 	fd := &tpuFD{
@@ -83,12 +87,7 @@ func (dev *tpuDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, o
 }
 
 // device implements vfs.Device for /dev/vfio/vfio.
-type vfioDevice struct {
-	mu sync.Mutex
-
-	// +checklocks:mu
-	devAddrSet DevAddrSet
-}
+type vfioDevice struct{}
 
 // Open implements vfs.Device.Open.
 func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, opts vfs.OpenOptions) (*vfs.FileDescription, error) {
@@ -98,9 +97,7 @@ func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry,
 		return nil, linuxerr.ENOENT
 	}
 
-	dev.mu.Lock()
-	defer dev.mu.Unlock()
-	name := fmt.Sprintf("vfio/%s", filepath.Base(VFIOPath))
+	name := filepath.Join("vfio", filepath.Base(VFIOPath))
 	hostFD, err := client.OpenAt(ctx, name, opts.Flags)
 	if err != nil {
 		ctx.Warningf("failed to open host file %s: %v", name, err)
@@ -125,9 +122,10 @@ func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry,
 }
 
 // RegisterTPUDevice registers devices implemented by this package in vfsObj.
-func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32) error {
+func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32) error {
 	return vfsObj.RegisterDevice(vfs.CharDevice, linux.VFIO_MAJOR, minor, &tpuDevice{
 		minor: minor,
+		num:   deviceNum,
 	}, &vfs.RegisterDeviceOptions{
 		GroupName: tpuDeviceGroupName,
 	})
diff --git a/pkg/sentry/devices/tpuproxy/seccomp_filter.go b/pkg/sentry/devices/tpuproxy/seccomp_filter.go
index f4f515d97..1db929cf1 100644
--- a/pkg/sentry/devices/tpuproxy/seccomp_filter.go
+++ b/pkg/sentry/devices/tpuproxy/seccomp_filter.go
@@ -95,6 +95,10 @@ func Filters() seccomp.SyscallRules {
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(linux.VFIO_IOMMU_MAP_DMA),
 			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(linux.VFIO_IOMMU_UNMAP_DMA),
+			},
 			seccomp.PerArg{
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(linux.VFIO_SET_IOMMU),
diff --git a/pkg/sentry/devices/tpuproxy/vfio.go b/pkg/sentry/devices/tpuproxy/vfio.go
index 2af5c277f..19d3474b9 100644
--- a/pkg/sentry/devices/tpuproxy/vfio.go
+++ b/pkg/sentry/devices/tpuproxy/vfio.go
@@ -16,6 +16,7 @@ package tpuproxy
 
 import (
 	"fmt"
+	"sync"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -45,10 +46,15 @@ type vfioFD struct {
 	device     *vfioDevice
 	queue      waiter.Queue
 	memmapFile vfioFDMemmapFile
+
+	mu sync.Mutex
+	// +checklocks:mu
+	devAddrSet DevAddrSet
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
 func (fd *vfioFD) Release(context.Context) {
+	fd.unpinRange(DevAddrRange{0, ^uint64(0)})
 	fdnotifier.RemoveFD(fd.hostFD)
 	fd.queue.Notify(waiter.EventHUp)
 	unix.Close(int(fd.hostFD))
@@ -198,15 +204,16 @@ func (fd *vfioFD) iommuMapDma(ctx context.Context, t *kernel.Task, arg hostarch.
 	// Unmap the reserved range, which is no longer required.
 	unix.RawSyscall(unix.SYS_MUNMAP, m, uintptr(ar.Length()), 0)
 
-	fd.device.mu.Lock()
-	defer fd.device.mu.Unlock()
+	fd.mu.Lock()
+	defer fd.mu.Unlock()
+	dar := devAddr
 	for _, pr := range prs {
-		rlen := uint64(pr.Source.Length())
-		fd.device.devAddrSet.InsertRange(DevAddrRange{
-			devAddr,
-			devAddr + rlen,
+		r := uint64(pr.Source.Length())
+		fd.devAddrSet.InsertRange(DevAddrRange{
+			dar,
+			dar + r,
 		}, pr)
-		devAddr += rlen
+		dar += r
 	}
 	return n, nil
 }
@@ -221,22 +228,29 @@ func (fd *vfioFD) iommuUnmapDma(ctx context.Context, t *kernel.Task, arg hostarc
 		// gVisor working with TPU.
 		return 0, linuxerr.ENOSYS
 	}
-	n, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_IOMMU_MAP_DMA, &dmaUnmap)
+	n, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_IOMMU_UNMAP_DMA, &dmaUnmap)
 	if err != nil {
 		return 0, nil
 	}
-	fd.device.mu.Lock()
-	defer fd.device.mu.Unlock()
-	s := &fd.device.devAddrSet
+	if _, err := dmaUnmap.CopyOut(t, arg); err != nil {
+		return 0, err
+	}
+
 	r := DevAddrRange{Start: dmaUnmap.IOVa, End: dmaUnmap.IOVa + dmaUnmap.Size}
-	seg := s.LowerBoundSegment(r.Start)
+	fd.unpinRange(r)
+	return n, nil
+}
+
+func (fd *vfioFD) unpinRange(r DevAddrRange) {
+	fd.mu.Lock()
+	defer fd.mu.Unlock()
+	seg := fd.devAddrSet.LowerBoundSegment(r.Start)
 	for seg.Ok() && seg.Start() < r.End {
-		seg = s.Isolate(seg, r)
+		seg = fd.devAddrSet.Isolate(seg, r)
 		mm.Unpin([]mm.PinnedRange{seg.Value()})
-		gap := s.Remove(seg)
+		gap := fd.devAddrSet.Remove(seg)
 		seg = gap.NextSegment()
 	}
-	return n, nil
 }
 
 // VFIO extension.
diff --git a/pkg/sentry/fsimpl/sys/pci.go b/pkg/sentry/fsimpl/sys/pci.go
index 1f14a6ecd..5539f5f88 100644
--- a/pkg/sentry/fsimpl/sys/pci.go
+++ b/pkg/sentry/fsimpl/sys/pci.go
@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"path"
 	regex "regexp"
+	"strings"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -29,17 +30,19 @@ import (
 )
 
 const (
-	pciMainBusDevicePath = "/sys/devices/pci0000:00"
-	accelDevice          = "accel"
-	vfioDevice           = "vfio-dev"
+	accelDevice        = "accel"
+	vfioDevice         = "vfio-dev"
+	sysDevicesMainPath = "/sys/devices"
 )
 
 var (
-	// Matches PCI device addresses in the main domain.
-	pciDeviceRegex = regex.MustCompile(`0000:([a-fA-F0-9]{2}|[a-fA-F0-9]{4}):[a-fA-F0-9]{2}\.[a-fA-F0-9]{1,2}`)
+	// pciBusRegex matches PCI bus addresses.
+	pciBusRegex = regex.MustCompile(`pci0000:[[:xdigit:]]{2}`)
+	// Matches PCI device addresses.
+	pciDeviceRegex = regex.MustCompile(`0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}`)
 	// Matches the directories for the main bus (i.e. pci000:00),
 	// individual devices (e.g. 00:00:04.0), accel (TPU v4), and vfio (TPU v5)
-	sysDevicesDirRegex = regex.MustCompile(`pci0000:00|accel|vfio|(0000:([a-fA-F0-9]{2}|[a-fA-F0-9]{4}):[a-fA-F0-9]{2}\.[a-fA-F0-9]{1,2})`)
+	sysDevicesDirRegex = regex.MustCompile(`pci0000:[[:xdigit:]]{2}|accel|vfio|vfio-dev|(0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2})`)
 	// Files allowlisted for host passthrough. These files are read-only.
 	sysDevicesFiles = map[string]any{
 		"vendor": nil, "device": nil, "subsystem_vendor": nil, "subsystem_device": nil,
@@ -49,27 +52,56 @@ var (
 		"is_device_owned": nil, "device_owner": nil, "framework_version": nil,
 		"user_mem_ranges": nil, "interrupt_counts": nil, "chip_model": nil,
 		"bar_offsets": nil, "bar_sizes": nil, "resource0": nil, "resource1": nil,
-		"resource2": nil, "resource3": nil, "resource4": nil, "resource5": nil,
+		"resource2": nil, "resource3": nil, "resource4": nil, "resource5": nil, "enable": nil,
 	}
 )
 
+// sysDevicesPCIPaths returns the paths of all PCI devices on the host in a
+// /sys/devices directory.
+func sysDevicesPCIPaths(sysDevicesPath string) ([]string, error) {
+	sysDevicesDents, err := hostDirEntries(sysDevicesPath)
+	if err != nil {
+		return nil, err
+	}
+	var pciPaths []string
+	for _, dent := range sysDevicesDents {
+		if pciBusRegex.MatchString(dent) {
+			pciDents, err := hostDirEntries(path.Join(sysDevicesPath, dent))
+			if err != nil {
+				return nil, err
+			}
+			for _, pciDent := range pciDents {
+				pciPaths = append(pciPaths, path.Join(sysDevicesPath, dent, pciDent))
+			}
+		}
+	}
+	return pciPaths, nil
+}
+
+// pciBusFromAddress returns the PCI bus address from a PCI address.
+//
+// Preconditions: pciAddr is a valid PCI address.
+func pciBusFromAddress(pciAddr string) string {
+	return strings.Join(strings.Split(pciAddr, ":")[:2], ":")
+}
+
 // Creates TPU devices' symlinks under /sys/class/. TPU device types that are
 // not present on host will be ignored.
 //
 // TPU v4 symlinks are created at /sys/class/accel/accel#.
 // TPU v5 symlinks go to /sys/class/vfio-dev/vfio#.
-func (fs *filesystem) newDeviceClassDir(ctx context.Context, creds *auth.Credentials, tpuDeviceTypes []string, pciMainBusDevicePath string) (map[string]map[string]kernfs.Inode, error) {
+func (fs *filesystem) newDeviceClassDir(ctx context.Context, creds *auth.Credentials, tpuDeviceTypes []string, sysDevicesPath string) (map[string]map[string]kernfs.Inode, error) {
 	dirs := map[string]map[string]kernfs.Inode{}
 	for _, tpuDeviceType := range tpuDeviceTypes {
 		dirs[tpuDeviceType] = map[string]kernfs.Inode{}
 	}
-	pciDents, err := hostDirEntries(pciMainBusDevicePath)
+	pciPaths, err := sysDevicesPCIPaths(sysDevicesPath)
 	if err != nil {
 		return nil, err
 	}
-	for _, pciDent := range pciDents {
+	for _, pciPath := range pciPaths {
 		for _, tpuDeviceType := range tpuDeviceTypes {
-			subPath := path.Join(pciMainBusDevicePath, pciDent, tpuDeviceType)
+			subPath := path.Join(pciPath, tpuDeviceType)
 			deviceDents, err := hostDirEntries(subPath)
 			if err != nil {
 				// Skips the path that doesn't exist.
@@ -81,7 +113,9 @@ func (fs *filesystem) newDeviceClassDir(ctx context.Context, creds *auth.Credent
 			if numOfDeviceDents := len(deviceDents); numOfDeviceDents != 1 {
 				return nil, fmt.Errorf("exactly one entry is expected at %v while there are %d", subPath, numOfDeviceDents)
 			}
-			dirs[tpuDeviceType][deviceDents[0]] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../devices/pci0000:00/%s/%s/%s", pciDent, tpuDeviceType, deviceDents[0]))
+			pciAddr := path.Base(pciPath)
+			pciBus := pciBusFromAddress(pciAddr)
+			dirs[tpuDeviceType][deviceDents[0]] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../devices/pci%s/%s/%s/%s", pciBus, pciAddr, tpuDeviceType, deviceDents[0]))
 		}
 	}
 	if len(dirs) == 0 {
@@ -91,14 +125,16 @@ func (fs *filesystem) newDeviceClassDir(ctx context.Context, creds *auth.Credent
 }
 
 // Create /sys/bus/pci/devices symlinks.
-func (fs *filesystem) newBusPCIDevicesDir(ctx context.Context, creds *auth.Credentials, pciMainBusDevicePath string) (map[string]kernfs.Inode, error) {
+func (fs *filesystem) newBusPCIDevicesDir(ctx context.Context, creds *auth.Credentials, sysDevicesPath string) (map[string]kernfs.Inode, error) {
 	pciDevicesDir := map[string]kernfs.Inode{}
-	pciDents, err := hostDirEntries(pciMainBusDevicePath)
+	pciPaths, err := sysDevicesPCIPaths(sysDevicesPath)
 	if err != nil {
 		return nil, err
 	}
-	for _, pciDent := range pciDents {
-		pciDevicesDir[pciDent] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../../devices/pci0000:00/%s", pciDent))
+	for _, pciPath := range pciPaths {
+		pciAddr := path.Base(pciPath)
+		pciBus := pciBusFromAddress(pciAddr)
+		pciDevicesDir[pciAddr] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../../devices/pci%s/%s", pciBus, pciAddr))
 	}
 
 	return pciDevicesDir, nil
@@ -106,7 +142,7 @@ func (fs *filesystem) newBusPCIDevicesDir(ctx context.Context, creds *auth.Crede
 
 // Recursively build out sysfs directories according to the allowlisted files,
 // directories, and symlinks defined in this package.
-func (fs *filesystem) mirrorPCIBusDeviceDir(ctx context.Context, creds *auth.Credentials, dir string, iommuGroups map[string]string) (map[string]kernfs.Inode, error) {
+func (fs *filesystem) mirrorSysDevicesDir(ctx context.Context, creds *auth.Credentials, dir string, iommuGroups map[string]string) (map[string]kernfs.Inode, error) {
 	subs := map[string]kernfs.Inode{}
 	dents, err := hostDirEntries(dir)
 	if err != nil {
@@ -123,7 +159,7 @@ func (fs *filesystem) mirrorPCIBusDeviceDir(ctx context.Context, creds *auth.Cre
 			if match := sysDevicesDirRegex.MatchString(dent); !match {
 				continue
 			}
-			contents, err := fs.mirrorPCIBusDeviceDir(ctx, creds, dentPath, iommuGroups)
+			contents, err := fs.mirrorSysDevicesDir(ctx, creds, dentPath, iommuGroups)
 			if err != nil {
 				return nil, err
 			}
diff --git a/pkg/sentry/fsimpl/sys/sys.go b/pkg/sentry/fsimpl/sys/sys.go
index 161e282d4..089e523ef 100644
--- a/pkg/sentry/fsimpl/sys/sys.go
+++ b/pkg/sentry/fsimpl/sys/sys.go
@@ -132,18 +132,20 @@ func (fsType FilesystemType) GetFilesystem(ctx context.Context, vfsObj *vfs.Virt
 		idata := opts.InternalData.(*InternalData)
 		productName = idata.ProductName
 		if idata.EnableTPUProxyPaths {
-			deviceToIommuGroup, err := pciDeviceIOMMUGroups(path.Join(idata.TestSysfsPathPrefix, iommuGroupSysPath))
+			deviceToIOMMUGroup, err := pciDeviceIOMMUGroups(path.Join(idata.TestSysfsPathPrefix, iommuGroupSysPath))
 			if err != nil {
 				return nil, nil, err
 			}
-			pciPath := path.Join(idata.TestSysfsPathPrefix, pciMainBusDevicePath)
-			pciMainBusSub, err := fs.mirrorPCIBusDeviceDir(ctx, creds, pciPath, deviceToIommuGroup)
+			sysDevicesPath := path.Join(idata.TestSysfsPathPrefix, sysDevicesMainPath)
+			sysDevicesSub, err := fs.mirrorSysDevicesDir(ctx, creds, sysDevicesPath, deviceToIOMMUGroup)
 			if err != nil {
 				return nil, nil, err
 			}
-			devicesSub["pci0000:00"] = fs.newDir(ctx, creds, defaultSysDirMode, pciMainBusSub)
+			for dir, sub := range sysDevicesSub {
+				devicesSub[dir] = sub
+			}
 
-			deviceDirs, err := fs.newDeviceClassDir(ctx, creds, []string{accelDevice, vfioDevice}, pciPath)
+			deviceDirs, err := fs.newDeviceClassDir(ctx, creds, []string{accelDevice, vfioDevice}, sysDevicesPath)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -151,7 +153,7 @@ func (fsType FilesystemType) GetFilesystem(ctx context.Context, vfsObj *vfs.Virt
 			for tpuDeviceType, symlinkDir := range deviceDirs {
 				classSub[tpuDeviceType] = fs.newDir(ctx, creds, defaultSysDirMode, symlinkDir)
 			}
-			pciDevicesSub, err := fs.newBusPCIDevicesDir(ctx, creds, pciPath)
+			pciDevicesSub, err := fs.newBusPCIDevicesDir(ctx, creds, sysDevicesPath)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -282,7 +284,8 @@ func (fs *filesystem) mirrorIOMMUGroups(ctx context.Context, creds *auth.Credent
 			subs[dent] = fs.newHostFile(ctx, creds, defaultSysMode, absPath)
 		case unix.S_IFLNK:
 			if pciDeviceRegex.MatchString(dent) {
-				subs[dent] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../../../devices/pci0000:00/%s", dent))
+				pciBus := pciBusFromAddress(dent)
+				subs[dent] = kernfs.NewStaticSymlink(ctx, creds, linux.UNNAMED_MAJOR, fs.devMinor, fs.NextIno(), fmt.Sprintf("../../../../devices/pci%s/%s", pciBus, dent))
 			}
 		}
 	}
diff --git a/pkg/sentry/fsimpl/sys/sys_test.go b/pkg/sentry/fsimpl/sys/sys_test.go
index 5d1bfb85f..4bc8da1cc 100644
--- a/pkg/sentry/fsimpl/sys/sys_test.go
+++ b/pkg/sentry/fsimpl/sys/sys_test.go
@@ -199,8 +199,12 @@ func (dev PCIDeviceInfo) path() string {
 func TestEnableTPUProxyPathsV5(t *testing.T) {
 	// Set up the fs tree that will be mirrored in the sentry.
 	sysfsTestDir := t.TempDir()
-	pciPath := path.Join(sysfsTestDir, "sys", "devices", "pci0000:00")
-	if err := os.MkdirAll(pciPath, 0755); err != nil {
+	pciPath0 := path.Join(sysfsTestDir, "sys", "devices", "pci0000:00")
+	if err := os.MkdirAll(pciPath0, 0755); err != nil {
+		t.Fatalf("Failed to create PCI directory: %v", err)
+	}
+	pciPath1 := path.Join(sysfsTestDir, "sys", "devices", "pci0000:10")
+	if err := os.MkdirAll(pciPath1, 0755); err != nil {
 		t.Fatalf("Failed to create PCI directory: %v", err)
 	}
 	busPath := path.Join(sysfsTestDir, "sys", "bus", "pci", "devices")
@@ -215,16 +219,22 @@ func TestEnableTPUProxyPathsV5(t *testing.T) {
 	devices := []PCIDeviceInfo{
 		PCIDeviceInfo{
 			group:      "0",
-			pciPath:    pciPath,
+			pciPath:    pciPath0,
 			pciAddress: "0000:00:04.0",
 			name:       "vfio0",
 		},
 		PCIDeviceInfo{
 			group:      "1",
-			pciPath:    pciPath,
+			pciPath:    pciPath0,
 			pciAddress: "0000:00:05.0",
 			name:       "vfio1",
 		},
+		PCIDeviceInfo{
+			group:      "2",
+			pciPath:    pciPath1,
+			pciAddress: "0000:10:05.0",
+			name:       "vfio2",
+		},
 	}
 	for _, device := range devices {
 		devicePath := device.path()
@@ -234,20 +244,20 @@ func TestEnableTPUProxyPathsV5(t *testing.T) {
 		if err := os.Symlink(path.Join("..", "..", "..", device.pciAddress), path.Join(devicePath, "device")); err != nil {
 			t.Fatalf("Failed to symlink device directory: %v", err)
 		}
-		if err := os.Symlink(path.Join("..", "..", "..", "devices", "pci0000:00", device.pciAddress), path.Join(busPath, device.pciAddress)); err != nil {
+		if err := os.Symlink(path.Join("..", "..", "..", "devices", path.Base(device.pciPath), device.pciAddress), path.Join(busPath, device.pciAddress)); err != nil {
 			t.Fatalf("Failed to symlink bus directory: %v", err)
 		}
-		if err := os.Symlink(path.Join("..", "..", "devices", "pci0000:00", device.pciAddress, vfioDev, device.name), path.Join(sysClassPath, device.name)); err != nil {
+		if err := os.Symlink(path.Join("..", "..", "devices", path.Base(device.pciPath), device.pciAddress, vfioDev, device.name), path.Join(sysClassPath, device.name)); err != nil {
 			t.Fatalf("Failed to symlink class directory: %v", err)
 		}
 		iommuPath := path.Join(sysfsTestDir, "sys", "kernel", "iommu_groups", device.group, "devices")
 		if err := os.MkdirAll(iommuPath, 0755); err != nil {
 			t.Fatalf("Failed to create iommu_groups directory: %v", err)
 		}
-		if err := os.Symlink(path.Join("..", "..", "..", "..", "devices", "pci0000:00", device.pciAddress), path.Join(iommuPath, device.pciAddress)); err != nil {
+		if err := os.Symlink(path.Join("..", "..", "..", "..", "devices", path.Base(device.pciPath), device.pciAddress), path.Join(iommuPath, device.pciAddress)); err != nil {
 			t.Fatalf("Failed to symlink iommu_group devices directory: %v", err)
 		}
-		if err := os.Symlink(path.Join("..", "..", "..", "kernel", "iommu_groups", device.group), path.Join(pciPath, device.pciAddress, "iommu_group")); err != nil {
+		if err := os.Symlink(path.Join("..", "..", "..", "kernel", "iommu_groups", device.group), path.Join(device.pciPath, device.pciAddress, "iommu_group")); err != nil {
 			t.Fatalf("Failed to symlink iommu_groups directory: %v", err)
 		}
 	}
@@ -256,13 +266,13 @@ func TestEnableTPUProxyPathsV5(t *testing.T) {
 
 	for _, device := range devices {
 		// Validate PCI device symlinks.
-		pop := s.PathOpAtRoot(path.Join("devices", "pci0000:00", device.pciAddress))
+		pop := s.PathOpAtRoot(path.Join("devices", path.Base(device.pciPath), device.pciAddress))
 		s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
 			"iommu_group": linux.DT_LNK,
 			vfioDev:       linux.DT_DIR,
 		})
 		// Validate VFIO device symlinks.
-		pop = s.PathOpAtRoot(path.Join("devices", "pci0000:00", device.pciAddress, vfioDev, device.name))
+		pop = s.PathOpAtRoot(path.Join("devices", path.Base(device.pciPath), device.pciAddress, vfioDev, device.name))
 		s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
 			"device": linux.DT_LNK,
 		})
diff --git a/runsc/boot/vfs.go b/runsc/boot/vfs.go
index 9d073beb3..06f55ed19 100644
--- a/runsc/boot/vfs.go
+++ b/runsc/boot/vfs.go
@@ -74,8 +74,9 @@ const (
 const SelfFilestorePrefix = ".gvisor.filestore."
 
 const (
-	pciPathGlobTPUv4 = "/sys/devices/pci0000:00/*/accel/accel*"
-	pciPathGlobTPUv5 = "/sys/devices/pci0000:00/*/vfio-dev/vfio*"
+	pciPathGlobTPUv4   = "/sys/devices/pci0000:*/*/accel/accel*"
+	pciPathGlobTPUv5   = "/sys/devices/pci0000:*/*/vfio-dev/vfio*"
+	iommuGroupPathGlob = "/sys/kernel/iommu_groups/*/devices/*"
 )
 
 // SelfFilestorePath returns the path at which the self filestore file is
@@ -1381,12 +1382,12 @@ func createDeviceFile(ctx context.Context, creds *auth.Credentials, info *contai
 }
 
 // registerTPUDevice registers a TPU device in vfsObj based on the given device ID.
-func registerTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32, deviceID int64) error {
+func registerTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32, deviceID int64) error {
 	switch deviceID {
 	case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID:
 		return accel.RegisterTPUDevice(vfsObj, minor, deviceID == tpu.TPUV4liteDeviceID)
-	case tpu.TPUV5eDeviceID:
-		return tpuproxy.RegisterTPUDevice(vfsObj, minor)
+	case tpu.TPUV5eDeviceID, tpu.TPUV5pDeviceID:
+		return tpuproxy.RegisterTPUDevice(vfsObj, minor, deviceNum)
 	default:
 		return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
 	}
@@ -1396,8 +1397,8 @@ func registerTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32, deviceID int
 // TPU v4 devices are accessible via /sys/devices/pci0000:00/<pci_address>/accel/accel# on the host.
 // TPU v5 devices are accessible via at /sys/devices/pci0000:00/<pci_address>/vfio-dev/vfio# on the host.
 var pathGlobToPathRegex = map[string]string{
-	pciPathGlobTPUv4: `^/sys/devices/pci0000:00/\d+:\d+:\d+\.\d+/accel/accel(\d+)$`,
-	pciPathGlobTPUv5: `^/sys/devices/pci0000:00/\d+:\d+:\d+\.\d+/vfio-dev/vfio(\d+)$`,
+	pciPathGlobTPUv4: `^/sys/devices/pci0000:[[:xdigit:]]{2}/\d+:\d+:\d+\.\d+/accel/accel(\d+)$`,
+	pciPathGlobTPUv5: `^/sys/devices/pci0000:[[:xdigit:]]{2}/\d+:\d+:\d+\.\d+/vfio-dev/vfio(\d+)$`,
 }
 
 func tpuProxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem) error {
@@ -1429,7 +1430,33 @@ func tpuProxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem)
 			if err != nil {
 				return fmt.Errorf("parsing PCI device ID: %w", err)
 			}
-			if err := registerTPUDevice(vfsObj, uint32(deviceNum), deviceID); err != nil {
+			// VFIO iommu groups correspond to the device minor number. Use these
+			// paths to get the correct minor number for the sentry-internal TPU
+			// device files.
+			var minorNum int
+			switch deviceID {
+			case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID:
+				minorNum = int(deviceNum)
+			case tpu.TPUV5eDeviceID, tpu.TPUV5pDeviceID:
+				groupPaths, err := filepath.Glob(iommuGroupPathGlob)
+				if err != nil {
+					return fmt.Errorf("enumerating IOMMU group files: %w", err)
+				}
+				for _, groupPath := range groupPaths {
+					pci := path.Base(groupPath)
+					if strings.Contains(pciPath, pci) {
+						minor, err := strconv.Atoi(strings.Split(groupPath, "/")[4])
+						if err != nil {
+							return fmt.Errorf("parsing IOMMU group minor number: %w", err)
+						}
+						minorNum = minor
+						break
+					}
+				}
+			default:
+				return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
+			}
+			if err := registerTPUDevice(vfsObj, uint32(minorNum), uint32(deviceNum), deviceID); err != nil {
 				return fmt.Errorf("registering TPU driver: %w", err)
 			}
 		}
diff --git a/runsc/cmd/chroot.go b/runsc/cmd/chroot.go
index 70287b635..8e4bbef53 100644
--- a/runsc/cmd/chroot.go
+++ b/runsc/cmd/chroot.go
@@ -151,24 +151,24 @@ func mountTPUSyslinkInChroot(chroot, dest, relativePath string, validator func(l
 }
 
 func mountTPUDeviceInfoInChroot(chroot, devicePath, sysfsFormat, pciDeviceFormat string) error {
-	deviceNum, valid, err := util.ExtractTpuDeviceMinor(devicePath)
+	deviceMinor, valid, err := util.ExtractTPUDeviceMinor(devicePath)
 	if err != nil {
 		return fmt.Errorf("extracting TPU device minor: %w", err)
 	}
 	if !valid {
 		return nil
 	}
-	// Multiple paths link to the /sys/devices/pci0000:00/<pci_address>
+	// Multiple paths link to the /sys/devices/<pci_bus>/<pci_address>
 	// directory that contains all relevant sysfs accel/vfio device info that we need
 	// bind mounted into the sandbox chroot. We can construct this path by
 	// reading the link below, which points to
-	//   * /sys/devices/pci0000:00/<pci_address>/accel/accel#
-	//   * or /sys/devices/pci0000:00/<pci_address>/vfio-dev/vfio# for VFIO-based TPU
+	//   * /sys/devices/<pci_bus>/<pci_address>/accel/accel#
+	//   * or /sys/devices/<pci_bus>/<pci_address>/vfio-dev/vfio# for VFIO-based TPU
 	// and traversing up 2 directories.
-	// The sysDevicePath itself is a soft link to the deivce directory.
-	sysDevicePath := fmt.Sprintf(sysfsFormat, deviceNum)
+	// The sysDevicePath itself is a soft link to the device directory.
+	sysDevicePath := fmt.Sprintf(sysfsFormat, deviceMinor)
 	sysPCIDeviceDir, err := mountTPUSyslinkInChroot(chroot, sysDevicePath, "../..", func(link string) bool {
-		sysDeviceLinkMatcher := regexp.MustCompile(fmt.Sprintf(pciDeviceFormat, deviceNum))
+		sysDeviceLinkMatcher := regexp.MustCompile(fmt.Sprintf(pciDeviceFormat, deviceMinor))
 		return sysDeviceLinkMatcher.MatchString(link)
 	})
 	if err != nil {
@@ -179,7 +179,8 @@ func mountTPUDeviceInfoInChroot(chroot, devicePath, sysfsFormat, pciDeviceFormat
 	iommuGroupPath := path.Join(sysPCIDeviceDir, "iommu_group")
 	if _, err := os.Stat(iommuGroupPath); err == nil {
 		if _, err := mountTPUSyslinkInChroot(chroot, iommuGroupPath, "", func(link string) bool {
-			return fmt.Sprintf("../../../kernel/iommu_groups/%d", deviceNum) == link
+			iommuGroupPathMatcher := regexp.MustCompile(`../../../kernel/iommu_groups/\d+`)
+			return iommuGroupPathMatcher.MatchString(link)
 		}); err != nil {
 			return err
 		}
@@ -196,11 +197,11 @@ func tpuProxyUpdateChroot(chroot string, spec *specs.Spec, conf *config.Config)
 		"/dev/accel*": "/sys/class/accel/accel%d",
 		"/dev/vfio/*": "/sys/class/vfio-dev/vfio%d"}
 	pathGlobToPciDeviceFormat := map[string]string{
-		"/dev/accel*": `../../devices/pci0000:00/(\d+:\d+:\d+\.\d+)/accel/accel%d`,
-		"/dev/vfio/*": `../../devices/pci0000:00/(\d+:\d+:\d+\.\d+)/vfio-dev/vfio%d`}
+		"/dev/accel*": `../../devices/pci0000:[[:xdigit:]]{2}/(\d+:\d+:\d+\.\d+)/accel/accel%d`,
+		"/dev/vfio/*": `../../devices/pci0000:[[:xdigit:]]{2}/(\d+:\d+:\d+\.\d+)/vfio-dev/vfio%d`}
 	// Bind mount device info directories for all TPU devices on the host.
-	// For v4 TPU, the directory /sys/devices/pci0000:00/<pci_address>/accel/accel# is mounted;
-	// For v5e TPU, the directory /sys/devices/pci0000:00/<pci_address>/vfio-dev/vfio# is mounted.
+	// For v4 TPU, the directory /sys/devices/<pci_bus>/<pci_address>/accel/accel# is mounted;
+	// For v5e TPU, the directory /sys/devices/<pci_bus>/<pci_address>/vfio-dev/vfio# is mounted.
 	for pathGlob, sysfsFormat := range pathGlobToSysfsFormat {
 		paths, err := filepath.Glob(pathGlob)
 		if err != nil {
diff --git a/runsc/cmd/gofer.go b/runsc/cmd/gofer.go
index 729839a67..100c6da1a 100644
--- a/runsc/cmd/gofer.go
+++ b/runsc/cmd/gofer.go
@@ -551,7 +551,7 @@ func shouldExposeVFIODevice(path string) bool {
 //
 // Precondition: tpuproxy is enabled.
 func shouldExposeTpuDevice(path string) bool {
-	_, valid, _ := util.ExtractTpuDeviceMinor(path)
+	_, valid, _ := util.ExtractTPUDeviceMinor(path)
 	return valid || shouldExposeVFIODevice(path)
 }
 
diff --git a/runsc/cmd/util/BUILD b/runsc/cmd/util/BUILD
index 2d55017dd..fd57bd66c 100644
--- a/runsc/cmd/util/BUILD
+++ b/runsc/cmd/util/BUILD
@@ -20,5 +20,6 @@ go_library(
         "//pkg/abi/tpu",
         "//pkg/log",
         "@com_github_google_subcommands//:go_default_library",
+        "@org_golang_x_sys//unix:go_default_library",
     ],
 )
diff --git a/runsc/cmd/util/tpu.go b/runsc/cmd/util/tpu.go
index 251ebbc4e..95e170660 100644
--- a/runsc/cmd/util/tpu.go
+++ b/runsc/cmd/util/tpu.go
@@ -20,7 +20,9 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"syscall"
 
+	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/tpu"
 )
 
@@ -35,12 +37,12 @@ const (
 )
 
 var tpuV4DeviceIDs = map[uint64]any{tpu.TPUV4DeviceID: nil, tpu.TPUV4liteDeviceID: nil}
-var tpuV5DeviceIDs = map[uint64]any{tpu.TPUV5eDeviceID: nil}
+var tpuV5DeviceIDs = map[uint64]any{tpu.TPUV5eDeviceID: nil, tpu.TPUV5pDeviceID: nil}
 
-// ExtractTpuDeviceMinor returns the accelerator device minor number for that
+// ExtractTPUDeviceMinor returns the accelerator device minor number for that
 // the passed device path. If the passed device is not a valid TPU device, then
 // it returns false.
-func ExtractTpuDeviceMinor(path string) (uint32, bool, error) {
+func ExtractTPUDeviceMinor(path string) (uint32, bool, error) {
 	devNum, valid, err := tpuV4DeviceMinor(path)
 	if err != nil {
 		return 0, false, err
@@ -60,10 +62,11 @@ func tpuDeviceMinor(devicePath, devicePathRegex, sysfsFormat string, allowedDevi
 	if matches == nil {
 		return 0, false, nil
 	}
-	minor, err := strconv.ParseUint(matches[1], 10, 32)
-	if err != nil {
-		return 0, false, fmt.Errorf("invalid host device file %q: %w", devicePath, err)
+	var st syscall.Stat_t
+	if err := syscall.Stat(devicePath, &st); err != nil {
+		return 0, false, err
 	}
+	minor := unix.Minor(st.Rdev)
 	vendor, err := readHexInt(fmt.Sprintf(sysfsFormat, minor, vendorFile))
 	if err != nil {
 		return 0, false, err
@@ -78,7 +81,7 @@ func tpuDeviceMinor(devicePath, devicePathRegex, sysfsFormat string, allowedDevi
 	if _, ok := allowedDeviceIDs[deviceID]; !ok {
 		return 0, false, nil
 	}
-	return uint32(minor), true, nil
+	return minor, true, nil
 }
 
 // tpuv4DeviceMinor returns v4 and v4lite TPU device minor number for the given path.