diff --git a/drivers/gpu/msm/adreno_a6xx_gmu.c b/drivers/gpu/msm/adreno_a6xx_gmu.c
index 8a4ea1752b24..07b1e1d24f45 100644
--- a/drivers/gpu/msm/adreno_a6xx_gmu.c
+++ b/drivers/gpu/msm/adreno_a6xx_gmu.c
@@ -593,7 +593,7 @@ static int find_vma_block(struct a6xx_gmu_device *gmu, u32 addr, u32 size)
 {
 	int i;
 
-	for (i = 0; i < GMU_MEM_TYPE_MAX; i++) {
+	for (i = 0; i < gmu->num_vmas; i++) {
 		struct gmu_vma_entry *vma = &gmu->vma[i];
 
 		if ((addr >= vma->start) &&
@@ -2685,10 +2685,13 @@ int a6xx_gmu_probe(struct kgsl_device *device,
 	if (ret)
 		goto error;
 
-	if (adreno_is_a650_family(adreno_dev))
+	if (adreno_is_a650_family(adreno_dev)) {
 		gmu->vma = a6xx_gmu_vma;
-	else
+		gmu->num_vmas = ARRAY_SIZE(a6xx_gmu_vma);
+	} else {
 		gmu->vma = a6xx_gmu_vma_legacy;
+		gmu->num_vmas = ARRAY_SIZE(a6xx_gmu_vma_legacy);
+	}
 
 	/* Map and reserve GMU CSRs registers */
 	ret = a6xx_gmu_reg_probe(adreno_dev);
diff --git a/drivers/gpu/msm/adreno_a6xx_gmu.h b/drivers/gpu/msm/adreno_a6xx_gmu.h
index d39597683ec9..b671771c46b4 100644
--- a/drivers/gpu/msm/adreno_a6xx_gmu.h
+++ b/drivers/gpu/msm/adreno_a6xx_gmu.h
@@ -187,6 +187,8 @@ struct a6xx_gmu_device {
 	/** @global_entries: To keep track of number of gmu buffers */
 	u32 global_entries;
 	struct gmu_vma_entry *vma;
+	/** @num_vmas: Number of entries in the @vma array */
+	u32 num_vmas;
 	unsigned int log_wptr_retention;
 	/** @cm3_fault: whether gmu received a cm3 fault interrupt */
 	atomic_t cm3_fault;