diff --git a/drivers/misc/habanalabs/common/mmu/mmu.c b/drivers/misc/habanalabs/common/mmu/mmu.c
index 792d25b79ea68b5c82b2365f43854f02a613bf8f..aa96917f62e57bc78c3462dcbf0466173009cbc9 100644
--- a/drivers/misc/habanalabs/common/mmu/mmu.c
+++ b/drivers/misc/habanalabs/common/mmu/mmu.c
@@ -501,23 +501,25 @@ static void hl_mmu_pa_page_with_offset(struct hl_ctx *ctx, u64 virt_addr,
 
 	if ((hops->range_type == HL_VA_RANGE_TYPE_DRAM) &&
 			!is_power_of_2(prop->dram_page_size)) {
-		unsigned long dram_page_size = prop->dram_page_size;
-		u64 page_offset_mask;
-		u64 phys_addr_mask;
-		u32 bit;
+		u64 dram_page_size, dram_base, abs_phys_addr, abs_virt_addr,
+			page_id, page_start;
+		u32 page_off;
 
 		/*
-		 * find last set bit in page_size to cover all bits of page
-		 * offset. note that 1 has to be added to bit index.
-		 * note that the internal ulong variable is used to avoid
-		 * alignment issue.
+		 * Bit arithmetics cannot be used for non power of two page
+		 * sizes. In addition, since bit arithmetics is not used,
+		 * we cannot ignore dram base. All that shall be considerd.
 		 */
-		bit = find_last_bit(&dram_page_size,
-					sizeof(dram_page_size) * BITS_PER_BYTE) + 1;
-		page_offset_mask = (BIT_ULL(bit) - 1);
-		phys_addr_mask = ~page_offset_mask;
-		*phys_addr = (tmp_phys_addr & phys_addr_mask) |
-				(virt_addr & page_offset_mask);
+
+		dram_page_size = prop->dram_page_size;
+		dram_base = prop->dram_base_address;
+		abs_phys_addr = tmp_phys_addr - dram_base;
+		abs_virt_addr = virt_addr - dram_base;
+		page_id = DIV_ROUND_DOWN_ULL(abs_phys_addr, dram_page_size);
+		page_start = page_id * dram_page_size;
+		div_u64_rem(abs_virt_addr, dram_page_size, &page_off);
+
+		*phys_addr = page_start + page_off + dram_base;
 	} else {
 		/*
 		 * find the correct hop shift field in hl_mmu_properties