diff --git a/drivers/accel/ivpu/ivpu_mmu_context.c b/drivers/accel/ivpu/ivpu_mmu_context.c
index 1d2e554e2c4a071d6b6974051252762c9ee2e790..ce94f4029127138618b9bb548f3ded0f9a603178 100644
--- a/drivers/accel/ivpu/ivpu_mmu_context.c
+++ b/drivers/accel/ivpu/ivpu_mmu_context.c
@@ -11,6 +11,7 @@
 #include "ivpu_mmu.h"
 #include "ivpu_mmu_context.h"
 
+#define IVPU_MMU_VPU_ADDRESS_MASK        GENMASK(47, 12)
 #define IVPU_MMU_PGD_INDEX_MASK          GENMASK(47, 39)
 #define IVPU_MMU_PUD_INDEX_MASK          GENMASK(38, 30)
 #define IVPU_MMU_PMD_INDEX_MASK          GENMASK(29, 21)
@@ -328,12 +329,8 @@ ivpu_mmu_context_map_sgt(struct ivpu_device *vdev, struct ivpu_mmu_context *ctx,
 
 	if (!IS_ALIGNED(vpu_addr, IVPU_MMU_PAGE_SIZE))
 		return -EINVAL;
-	/*
-	 * VPU is only 32 bit, but DMA engine is 38 bit
-	 * Ranges < 2 GB are reserved for VPU internal registers
-	 * Limit range to 8 GB
-	 */
-	if (vpu_addr < SZ_2G || vpu_addr > SZ_8G)
+
+	if (vpu_addr & ~IVPU_MMU_VPU_ADDRESS_MASK)
 		return -EINVAL;
 
 	prot = IVPU_MMU_ENTRY_MAPPED;