mlx5_db_free(ndev->mvdev.mdev, &vcq->db);
 }
 
+static int read_umem_params(struct mlx5_vdpa_net *ndev)
+{
+       u32 in[MLX5_ST_SZ_DW(query_hca_cap_in)] = {};
+       u16 opmod = (MLX5_CAP_VDPA_EMULATION << 1) | (HCA_CAP_OPMOD_GET_CUR & 0x01);
+       struct mlx5_core_dev *mdev = ndev->mvdev.mdev;
+       int out_size;
+       void *caps;
+       void *out;
+       int err;
+
+       out_size = MLX5_ST_SZ_BYTES(query_hca_cap_out);
+       out = kzalloc(out_size, GFP_KERNEL);
+       if (!out)
+               return -ENOMEM;
+
+       MLX5_SET(query_hca_cap_in, in, opcode, MLX5_CMD_OP_QUERY_HCA_CAP);
+       MLX5_SET(query_hca_cap_in, in, op_mod, opmod);
+       err = mlx5_cmd_exec_inout(mdev, query_hca_cap, in, out);
+       if (err) {
+               mlx5_vdpa_warn(&ndev->mvdev,
+                       "Failed reading vdpa umem capabilities with err %d\n", err);
+               goto out;
+       }
+
+       caps =  MLX5_ADDR_OF(query_hca_cap_out, out, capability);
+
+       ndev->umem_1_buffer_param_a = MLX5_GET(virtio_emulation_cap, caps, umem_1_buffer_param_a);
+       ndev->umem_1_buffer_param_b = MLX5_GET(virtio_emulation_cap, caps, umem_1_buffer_param_b);
+
+       ndev->umem_2_buffer_param_a = MLX5_GET(virtio_emulation_cap, caps, umem_2_buffer_param_a);
+       ndev->umem_2_buffer_param_b = MLX5_GET(virtio_emulation_cap, caps, umem_2_buffer_param_b);
+
+       ndev->umem_3_buffer_param_a = MLX5_GET(virtio_emulation_cap, caps, umem_3_buffer_param_a);
+       ndev->umem_3_buffer_param_b = MLX5_GET(virtio_emulation_cap, caps, umem_3_buffer_param_b);
+
+out:
+       kfree(out);
+       return 0;
+}
+
 static void set_umem_size(struct mlx5_vdpa_net *ndev, struct mlx5_vdpa_virtqueue *mvq, int num,
                          struct mlx5_vdpa_umem **umemp)
 {
-       struct mlx5_core_dev *mdev = ndev->mvdev.mdev;
-       int p_a;
-       int p_b;
+       u32 p_a;
+       u32 p_b;
 
        switch (num) {
        case 1:
-               p_a = MLX5_CAP_DEV_VDPA_EMULATION(mdev, umem_1_buffer_param_a);
-               p_b = MLX5_CAP_DEV_VDPA_EMULATION(mdev, umem_1_buffer_param_b);
+               p_a = ndev->umem_1_buffer_param_a;
+               p_b = ndev->umem_1_buffer_param_b;
                *umemp = &mvq->umem1;
                break;
        case 2:
-               p_a = MLX5_CAP_DEV_VDPA_EMULATION(mdev, umem_2_buffer_param_a);
-               p_b = MLX5_CAP_DEV_VDPA_EMULATION(mdev, umem_2_buffer_param_b);
+               p_a = ndev->umem_2_buffer_param_a;
+               p_b = ndev->umem_2_buffer_param_b;
                *umemp = &mvq->umem2;
                break;
        case 3:
-               p_a = MLX5_CAP_DEV_VDPA_EMULATION(mdev, umem_3_buffer_param_a);
-               p_b = MLX5_CAP_DEV_VDPA_EMULATION(mdev, umem_3_buffer_param_b);
+               p_a = ndev->umem_3_buffer_param_a;
+               p_b = ndev->umem_3_buffer_param_b;
                *umemp = &mvq->umem3;
                break;
        }
+
        (*umemp)->size = p_a * mvq->num_ent + p_b;
 }
 
                goto out;
        }
        mlx5_vdpa_add_debugfs(ndev);
+
+       err = read_umem_params(ndev);
+       if (err)
+               goto err_setup;
+
        err = setup_virtqueues(mvdev);
        if (err) {
                mlx5_vdpa_warn(mvdev, "setup_virtqueues\n");