diff --git a/drivers/vfio/vfio.c b/drivers/vfio/vfio.c
index 3b2e87df2e427..a156cbbad52dd 100644
--- a/drivers/vfio/vfio.c
+++ b/drivers/vfio/vfio.c
@@ -84,6 +84,7 @@ struct vfio_group {
 	unsigned int			dev_counter;
 	struct kvm			*kvm;
 	struct blocking_notifier_head	notifier;
+	void				*iommufd;
 };
 
 #ifdef CONFIG_VFIO_NOIOMMU
@@ -1079,6 +1080,12 @@ static void __vfio_group_unset_container(struct vfio_group *group)
 	struct vfio_container *container = group->container;
 	struct vfio_iommu_driver *driver;
 
+	if (group->iommufd) {
+		vfio_group_unset_iommufd(group->iommufd, &group->device_list);
+		group->iommufd = NULL;
+		return;
+	}
+
 	down_write(&container->group_lock);
 
 	driver = container->iommu_driver;
@@ -1150,6 +1157,13 @@ static int vfio_group_set_container(struct vfio_group *group, int container_fd)
 	if (group->type == VFIO_NO_IOMMU && !capable(CAP_SYS_RAWIO))
 		return -EPERM;
 
+	group->iommufd = vfio_group_set_iommufd(container_fd,
+						&group->device_list);
+	if (group->iommufd) {
+		atomic_inc(&group->container_users);
+		return ret;
+	}
+
 	f = fdget(container_fd);
 	if (!f.file)
 		return -EBADF;
@@ -1261,7 +1275,8 @@ static int vfio_group_get_device_fd(struct vfio_group *group, char *buf)
 	int ret = 0;
 
 	if (0 == atomic_read(&group->container_users) ||
-	    !group->container->iommu_driver)
+	    (group->container && !group->container->iommu_driver) ||
+	    (!group->container && !group->iommufd))
 		return -EINVAL;
 
 	if (group->type == VFIO_NO_IOMMU && !capable(CAP_SYS_RAWIO))