diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_gmc.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_gmc.c
index 3a71028e41930..4e0a2d86a7268 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_gmc.c
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_gmc.c
@@ -1244,14 +1244,14 @@ void amdgpu_gmc_sysfs_fini(struct amdgpu_device *adev)
 
 int amdgpu_gmc_get_nps_memranges(struct amdgpu_device *adev,
 				 struct amdgpu_mem_partition_info *mem_ranges,
-				 int exp_ranges)
+				 uint8_t *exp_ranges)
 {
 	struct amdgpu_gmc_memrange *ranges;
 	int range_cnt, ret, i, j;
 	uint32_t nps_type;
 	bool refresh;
 
-	if (!mem_ranges)
+	if (!mem_ranges || !exp_ranges)
 		return -EINVAL;
 
 	refresh = (adev->init_lvl->level != AMDGPU_INIT_LEVEL_MINIMAL_XGMI) &&
@@ -1265,16 +1265,16 @@ int amdgpu_gmc_get_nps_memranges(struct amdgpu_device *adev,
 	/* TODO: For now, expect ranges and partition count to be the same.
 	 * Adjust if there are holes expected in any NPS domain.
 	 */
-	if (range_cnt != exp_ranges) {
+	if (*exp_ranges && (range_cnt != *exp_ranges)) {
 		dev_warn(
 			adev->dev,
 			"NPS config mismatch - expected ranges: %d discovery - nps mode: %d, nps ranges: %d",
-			exp_ranges, nps_type, range_cnt);
+			*exp_ranges, nps_type, range_cnt);
 		ret = -EINVAL;
 		goto err;
 	}
 
-	for (i = 0; i < exp_ranges; ++i) {
+	for (i = 0; i < range_cnt; ++i) {
 		if (ranges[i].base_address >= ranges[i].limit_address) {
 			dev_warn(
 				adev->dev,
@@ -1315,6 +1315,8 @@ int amdgpu_gmc_get_nps_memranges(struct amdgpu_device *adev,
 			ranges[i].limit_address - ranges[i].base_address + 1;
 	}
 
+	if (!*exp_ranges)
+		*exp_ranges = range_cnt;
 err:
 	kfree(ranges);
 
diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_gmc.h b/drivers/gpu/drm/amd/amdgpu/amdgpu_gmc.h
index c66b6dbe9ac65..459a30fe239f7 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_gmc.h
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_gmc.h
@@ -466,7 +466,7 @@ void amdgpu_gmc_sysfs_fini(struct amdgpu_device *adev);
 
 int amdgpu_gmc_get_nps_memranges(struct amdgpu_device *adev,
 				 struct amdgpu_mem_partition_info *mem_ranges,
-				 int exp_ranges);
+				 uint8_t *exp_ranges);
 
 int amdgpu_gmc_request_memory_partition(struct amdgpu_device *adev,
 					int nps_mode);
diff --git a/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c b/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
index bfbfbcc9b89b0..010db0e586500 100644
--- a/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
+++ b/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
@@ -1386,11 +1386,30 @@ gmc_v9_0_get_memory_partition(struct amdgpu_device *adev, u32 *supp_modes)
 	return mode;
 }
 
+static enum amdgpu_memory_partition
+gmc_v9_0_query_vf_memory_partition(struct amdgpu_device *adev)
+{
+	switch (adev->gmc.num_mem_partitions) {
+	case 0:
+		return UNKNOWN_MEMORY_PARTITION_MODE;
+	case 1:
+		return AMDGPU_NPS1_PARTITION_MODE;
+	case 2:
+		return AMDGPU_NPS2_PARTITION_MODE;
+	case 4:
+		return AMDGPU_NPS4_PARTITION_MODE;
+	default:
+		return AMDGPU_NPS1_PARTITION_MODE;
+	}
+
+	return AMDGPU_NPS1_PARTITION_MODE;
+}
+
 static enum amdgpu_memory_partition
 gmc_v9_0_query_memory_partition(struct amdgpu_device *adev)
 {
 	if (amdgpu_sriov_vf(adev))
-		return AMDGPU_NPS1_PARTITION_MODE;
+		return gmc_v9_0_query_vf_memory_partition(adev);
 
 	return gmc_v9_0_get_memory_partition(adev, NULL);
 }
@@ -1935,6 +1954,8 @@ gmc_v9_0_init_sw_mem_ranges(struct amdgpu_device *adev,
 
 	switch (mode) {
 	case UNKNOWN_MEMORY_PARTITION_MODE:
+		adev->gmc.num_mem_partitions = 0;
+		break;
 	case AMDGPU_NPS1_PARTITION_MODE:
 		adev->gmc.num_mem_partitions = 1;
 		break;
@@ -1954,7 +1975,7 @@ gmc_v9_0_init_sw_mem_ranges(struct amdgpu_device *adev,
 
 	/* Use NPS range info, if populated */
 	r = amdgpu_gmc_get_nps_memranges(adev, mem_ranges,
-					 adev->gmc.num_mem_partitions);
+					 &adev->gmc.num_mem_partitions);
 	if (!r) {
 		l = 0;
 		for (i = 1; i < adev->gmc.num_mem_partitions; ++i) {
@@ -1964,6 +1985,11 @@ gmc_v9_0_init_sw_mem_ranges(struct amdgpu_device *adev,
 		}
 
 	} else {
+		if (!adev->gmc.num_mem_partitions) {
+			dev_err(adev->dev,
+				"Not able to detect NPS mode, fall back to NPS1");
+			adev->gmc.num_mem_partitions = 1;
+		}
 		/* Fallback to sw based calculation */
 		size = (adev->gmc.real_vram_size + SZ_16M) >> AMDGPU_GPU_PAGE_SHIFT;
 		size /= adev->gmc.num_mem_partitions;