diff --git a/hipamd/src/hip_vm.cpp b/hipamd/src/hip_vm.cpp
index 31eb381a76..2c81a42f1b 100644
--- a/hipamd/src/hip_vm.cpp
+++ b/hipamd/src/hip_vm.cpp
@@ -202,8 +202,8 @@ hipError_t hipMemMap(void* ptr, size_t size, size_t offset, hipMemGenericAllocat
   cmd->awaitCompletion();
   cmd->release();
 
-  amd::Memory* va = amd::MemObjMap::FindMemObj(ptr);
-  va->getUserData().deviceId = ga->GetProperties().location.id;
+  // update the internal svm address to ptr
+  ga->asAmdMemory().setSvmPtr(ptr);
 
   HIP_RETURN(hipSuccess);
 }
@@ -235,6 +235,18 @@ hipError_t hipMemRetainAllocationHandle(hipMemGenericAllocationHandle_t* handle,
 
   if (handle == nullptr || addr == nullptr) HIP_RETURN(hipErrorInvalidValue);
 
+  amd::Memory* mem = amd::MemObjMap::FindMemObj(addr);
+
+  if (mem == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+
+  *handle = reinterpret_cast<hipMemGenericAllocationHandle_t>(mem->getUserData().data);
+
+  if (*handle == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+
   HIP_RETURN(hipSuccess);
 }
 
@@ -265,6 +277,10 @@ hipError_t hipMemUnmap(void* ptr, size_t size) {
   cmd->awaitCompletion();
   cmd->release();
 
+  // restore the original va of the generic allocation
+  hip::GenericAllocation* ga = reinterpret_cast<hip::GenericAllocation*>(va->getUserData().data);
+  va->setSvmPtr(ga->genericAddress());
+
   HIP_RETURN(hipSuccess);
 }
 
diff --git a/hipamd/src/hip_vm.hpp b/hipamd/src/hip_vm.hpp
index a6b8f7ed7d..d87768ac46 100644
--- a/hipamd/src/hip_vm.hpp
+++ b/hipamd/src/hip_vm.hpp
@@ -38,7 +38,8 @@ public:
 
   const hipMemAllocationProp& GetProperties() const { return properties_; }
   hipMemGenericAllocationHandle_t asMemGenericAllocationHandle() { return reinterpret_cast<hipMemGenericAllocationHandle_t>(this); }
-  amd::Memory& asAmdMemory() { return *amd::MemObjMap::FindMemObj(ptr_); }
+  amd::Memory& asAmdMemory() { return *amd::MemObjMap::FindMemObj(genericAddress()); }
+  void* genericAddress() const { return ptr_; }
 };
 };