diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_device.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_device.c
index ae31744dac28a8df6550abef270d98403e5a9103..93061a439dbc1a25a12e70c229513dc00438e54a 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_device.c
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_device.c
@@ -1043,10 +1043,10 @@ static bool amdgpu_check_pot_argument(int arg)
 static void amdgpu_get_block_size(struct amdgpu_device *adev)
 {
 	/* from AI, asic starts to support multiple level VMPT */
-	if (adev->family >= AMDGPU_FAMILY_AI) {
+	if (adev->asic_type >= CHIP_VEGA10) {
 		if (amdgpu_vm_block_size != 9)
-			dev_warn(adev->dev, "Multi-VMPT limits block size to"
-				 "one page!\n");
+			dev_warn(adev->dev,
+				 "Multi-VMPT limits block size to one page!\n");
 		amdgpu_vm_block_size = 9;
 		return;
 	}
diff --git a/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c b/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
index 1e4734daa6ed79e79ea590d2e94401395d2395b6..df69aae99df48cb95c0a16a3e26486268f0a307a 100644
--- a/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
+++ b/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
@@ -511,12 +511,6 @@ static int gmc_v9_0_vm_init(struct amdgpu_device *adev)
 	 * amdkfd will use VMIDs 8-15
 	 */
 	adev->vm_manager.num_ids = AMDGPU_NUM_OF_VMIDS;
-	/* Because of four level VMPTs, vm size at least is 256GB.
-	256TB is OK as well */
-	if (amdgpu_vm_size < 256) {
-		DRM_WARN("vm size at least is 256GB!\n");
-		amdgpu_vm_size = 256;
-	}
 	adev->vm_manager.num_level = 3;
 	amdgpu_vm_manager_init(adev);
 
@@ -563,11 +557,14 @@ static int gmc_v9_0_sw_init(void *handle)
 	if (r)
 		return r;
 
-	/* Adjust VM size here.
-	 * Currently default to 64GB ((16 << 20) 4k pages).
-	 * Max GPUVM size is 48 bits.
+	/* Because of four level VMPTs, vm size is at least 512GB.
+	 * The maximum size is 256TB (48bit).
 	 */
-	adev->vm_manager.max_pfn = amdgpu_vm_size << 18;
+	if (amdgpu_vm_size < 512) {
+		DRM_WARN("VM size is at least 512GB!\n");
+		amdgpu_vm_size = 512;
+	}
+	adev->vm_manager.max_pfn = (uint64_t)amdgpu_vm_size << 18;
 
 	/* Set the internal MC address mask
 	 * This is the max address of the GPU's