diff --git a/drivers/gpu/drm/amd/amdgpu/soc15.c b/drivers/gpu/drm/amd/amdgpu/soc15.c
index 0986ea59cca29385063fc960cd57c7e111952498..333edbf564df373ec80f6c6bf287491bb9f0fe25 100644
--- a/drivers/gpu/drm/amd/amdgpu/soc15.c
+++ b/drivers/gpu/drm/amd/amdgpu/soc15.c
@@ -58,6 +58,7 @@
 #include "uvd_v7_0.h"
 #include "vce_v4_0.h"
 #include "amdgpu_powerplay.h"
+#include "mxgpu_ai.h"
 
 MODULE_FIRMWARE("amdgpu/vega10_smc.bin");
 
@@ -471,6 +472,9 @@ int soc15_set_ip_blocks(struct amdgpu_device *adev)
 {
 	nbio_v6_1_detect_hw_virt(adev);
 
+	if (amdgpu_sriov_vf(adev))
+		adev->virt.ops = &xgpu_ai_virt_ops;
+
 	switch (adev->asic_type) {
 	case CHIP_VEGA10:
 		amdgpu_ip_block_add(adev, &vega10_common_ip_block);