diff --git a/drivers/vdpa/mlx5/mlx5_vdpa.c b/drivers/vdpa/mlx5/mlx5_vdpa.c
index 5ab7c525c2..9c9a552ba0 100644
--- a/drivers/vdpa/mlx5/mlx5_vdpa.c
+++ b/drivers/vdpa/mlx5/mlx5_vdpa.c
@@ -11,12 +11,11 @@
 #include <rte_malloc.h>
 #include <rte_log.h>
 #include <rte_errno.h>
-#include <rte_pci.h>
 #include <rte_string_fns.h>
+#include <rte_bus_pci.h>
 
 #include <mlx5_glue.h>
 #include <mlx5_common.h>
-#include <mlx5_common_pci.h>
 #include <mlx5_devx_cmds.h>
 #include <mlx5_prm.h>
 #include <mlx5_nl.h>
@@ -552,34 +551,13 @@ mlx5_vdpa_sys_roce_disable(const char *addr)
 }
 
 static int
-mlx5_vdpa_roce_disable(struct rte_pci_addr *addr, struct ibv_device **ibv)
+mlx5_vdpa_roce_disable(struct rte_device *dev)
 {
-	char addr_name[64] = {0};
-
-	rte_pci_device_name(addr, addr_name, sizeof(addr_name));
 	/* Firstly try to disable ROCE by Netlink and fallback to sysfs. */
-	if (mlx5_vdpa_nl_roce_disable(addr_name) == 0 ||
-	    mlx5_vdpa_sys_roce_disable(addr_name) == 0) {
-		/*
-		 * Succeed to disable ROCE, wait for the IB device to appear
-		 * again after reload.
-		 */
-		int r;
-		struct ibv_device *ibv_new;
-
-		for (r = MLX5_VDPA_MAX_RETRIES; r; r--) {
-			ibv_new = mlx5_os_get_ibv_device(addr);
-			if (ibv_new) {
-				*ibv = ibv_new;
-				return 0;
-			}
-			usleep(MLX5_VDPA_USEC);
-		}
-		DRV_LOG(ERR, "Cannot much device %s after ROCE disable, "
-			"retries exceed %d", addr_name, MLX5_VDPA_MAX_RETRIES);
-		rte_errno = EAGAIN;
-	}
-	return -rte_errno;
+	if (mlx5_vdpa_nl_roce_disable(dev->name) != 0 &&
+	    mlx5_vdpa_sys_roce_disable(dev->name) != 0)
+		return -rte_errno;
+	return 0;
 }
 
 static int
@@ -647,44 +625,33 @@ mlx5_vdpa_config_get(struct rte_devargs *devargs, struct mlx5_vdpa_priv *priv)
 	DRV_LOG(DEBUG, "no traffic max is %u.", priv->no_traffic_max);
 }
 
-/**
- * DPDK callback to register a mlx5 PCI device.
- *
- * This function spawns vdpa device out of a given PCI device.
- *
- * @param[in] pci_drv
- *   PCI driver structure (mlx5_vpda_driver).
- * @param[in] pci_dev
- *   PCI device information.
- *
- * @return
- *   0 on success, 1 to skip this driver, a negative errno value otherwise
- *   and rte_errno is set.
- */
 static int
-mlx5_vdpa_pci_probe(struct rte_pci_driver *pci_drv __rte_unused,
-		    struct rte_pci_device *pci_dev __rte_unused)
+mlx5_vdpa_dev_probe(struct rte_device *dev)
 {
 	struct ibv_device *ibv;
 	struct mlx5_vdpa_priv *priv = NULL;
 	struct ibv_context *ctx = NULL;
 	struct mlx5_hca_attr attr;
+	int retry;
 	int ret;
 
-	ibv = mlx5_os_get_ibv_device(&pci_dev->addr);
-	if (!ibv) {
-		DRV_LOG(ERR, "No matching IB device for PCI slot "
-			PCI_PRI_FMT ".", pci_dev->addr.domain,
-			pci_dev->addr.bus, pci_dev->addr.devid,
-			pci_dev->addr.function);
+	if (mlx5_vdpa_roce_disable(dev) != 0) {
+		DRV_LOG(WARNING, "Failed to disable ROCE for \"%s\".",
+			dev->name);
 		return -rte_errno;
-	} else {
-		DRV_LOG(INFO, "PCI information matches for device \"%s\".",
-			ibv->name);
 	}
-	if (mlx5_vdpa_roce_disable(&pci_dev->addr, &ibv) != 0) {
-		DRV_LOG(WARNING, "Failed to disable ROCE for \"%s\".",
-			ibv->name);
+	/* Wait for the IB device to appear again after reload. */
+	for (retry = MLX5_VDPA_MAX_RETRIES; retry > 0; --retry) {
+		ibv = mlx5_os_get_ibv_dev(dev);
+		if (ibv != NULL)
+			break;
+		usleep(MLX5_VDPA_USEC);
+	}
+	if (ibv == NULL) {
+		DRV_LOG(ERR, "Cannot get IB device after disabling RoCE for "
+				"\"%s\", retries exceed %d.",
+				dev->name, MLX5_VDPA_MAX_RETRIES);
+		rte_errno = EAGAIN;
 		return -rte_errno;
 	}
 	ctx = mlx5_glue->dv_open_device(ibv);
@@ -722,20 +689,18 @@ mlx5_vdpa_pci_probe(struct rte_pci_driver *pci_drv __rte_unused,
 	if (attr.num_lag_ports == 0)
 		priv->num_lag_ports = 1;
 	priv->ctx = ctx;
-	priv->pci_dev = pci_dev;
 	priv->var = mlx5_glue->dv_alloc_var(ctx, 0);
 	if (!priv->var) {
 		DRV_LOG(ERR, "Failed to allocate VAR %u.", errno);
 		goto error;
 	}
-	priv->vdev = rte_vdpa_register_device(&pci_dev->device,
-			&mlx5_vdpa_ops);
+	priv->vdev = rte_vdpa_register_device(dev, &mlx5_vdpa_ops);
 	if (priv->vdev == NULL) {
 		DRV_LOG(ERR, "Failed to register vDPA device.");
 		rte_errno = rte_errno ? rte_errno : EINVAL;
 		goto error;
 	}
-	mlx5_vdpa_config_get(pci_dev->device.devargs, priv);
+	mlx5_vdpa_config_get(dev->devargs, priv);
 	SLIST_INIT(&priv->mr_list);
 	pthread_mutex_init(&priv->vq_config_lock, NULL);
 	pthread_mutex_lock(&priv_list_lock);
@@ -754,26 +719,15 @@ mlx5_vdpa_pci_probe(struct rte_pci_driver *pci_drv __rte_unused,
 	return -rte_errno;
 }
 
-/**
- * DPDK callback to remove a PCI device.
- *
- * This function removes all vDPA devices belong to a given PCI device.
- *
- * @param[in] pci_dev
- *   Pointer to the PCI device.
- *
- * @return
- *   0 on success, the function cannot fail.
- */
 static int
-mlx5_vdpa_pci_remove(struct rte_pci_device *pci_dev)
+mlx5_vdpa_dev_remove(struct rte_device *dev)
 {
 	struct mlx5_vdpa_priv *priv = NULL;
 	int found = 0;
 
 	pthread_mutex_lock(&priv_list_lock);
 	TAILQ_FOREACH(priv, &priv_list, next) {
-		if (!rte_pci_addr_cmp(&priv->pci_dev->addr, &pci_dev->addr)) {
+		if (priv->vdev->device == dev) {
 			found = 1;
 			break;
 		}
@@ -831,17 +785,12 @@ static const struct rte_pci_id mlx5_vdpa_pci_id_map[] = {
 	}
 };
 
-static struct mlx5_pci_driver mlx5_vdpa_driver = {
-	.driver_class = MLX5_CLASS_VDPA,
-	.pci_driver = {
-		.driver = {
-			.name = RTE_STR(MLX5_VDPA_DRIVER_NAME),
-		},
-		.id_table = mlx5_vdpa_pci_id_map,
-		.probe = mlx5_vdpa_pci_probe,
-		.remove = mlx5_vdpa_pci_remove,
-		.drv_flags = 0,
-	},
+static struct mlx5_class_driver mlx5_vdpa_driver = {
+	.drv_class = MLX5_CLASS_VDPA,
+	.name = RTE_STR(MLX5_VDPA_DRIVER_NAME),
+	.id_table = mlx5_vdpa_pci_id_map,
+	.probe = mlx5_vdpa_dev_probe,
+	.remove = mlx5_vdpa_dev_remove,
 };
 
 RTE_LOG_REGISTER_DEFAULT(mlx5_vdpa_logtype, NOTICE)
@@ -853,7 +802,7 @@ RTE_INIT(rte_mlx5_vdpa_init)
 {
 	mlx5_common_init();
 	if (mlx5_glue)
-		mlx5_pci_driver_register(&mlx5_vdpa_driver);
+		mlx5_class_driver_register(&mlx5_vdpa_driver);
 }
 
 RTE_PMD_EXPORT_NAME(MLX5_VDPA_DRIVER_NAME, __COUNTER__);
diff --git a/drivers/vdpa/mlx5/mlx5_vdpa.h b/drivers/vdpa/mlx5/mlx5_vdpa.h
index 722c72b65e..2a04e36607 100644
--- a/drivers/vdpa/mlx5/mlx5_vdpa.h
+++ b/drivers/vdpa/mlx5/mlx5_vdpa.h
@@ -133,7 +133,6 @@ struct mlx5_vdpa_priv {
 	struct rte_vdpa_device *vdev; /* vDPA device. */
 	int vid; /* vhost device id. */
 	struct ibv_context *ctx; /* Device context. */
-	struct rte_pci_device *pci_dev;
 	struct mlx5_hca_vdpa_attr caps;
 	uint32_t pdn; /* Protection Domain number. */
 	struct ibv_pd *pd;
