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) &&
        /* 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,
                        ranges[i].limit_address - ranges[i].base_address + 1;
        }
 
+       if (!*exp_ranges)
+               *exp_ranges = range_cnt;
 err:
        kfree(ranges);
 
 
        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);
 }
 
        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;
 
        /* 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) {
                }
 
        } 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;