diff --git a/hipamd/src/hip_vm.cpp b/hipamd/src/hip_vm.cpp
index e11e46d96d..fc2533d8a3 100644
--- a/hipamd/src/hip_vm.cpp
+++ b/hipamd/src/hip_vm.cpp
@@ -366,6 +366,12 @@ hipError_t hipMemUnmap(void* ptr, size_t size) {
   cmd->enqueue();
   cmd->awaitCompletion();
   cmd->release();
+  vaddr_sub_obj->release();
+
+  // restore the original pa of the generic allocation
+  hip::GenericAllocation* ga
+    = reinterpret_cast<hip::GenericAllocation*>(phys_mem_obj->getUserData().data);
+  ga->release();
 
   HIP_RETURN(hipSuccess);
 }
diff --git a/rocclr/device/rocm/rocvirtual.cpp b/rocclr/device/rocm/rocvirtual.cpp
index 6881654509..491b72e802 100644
--- a/rocclr/device/rocm/rocvirtual.cpp
+++ b/rocclr/device/rocm/rocvirtual.cpp
@@ -2924,29 +2924,21 @@ void VirtualGPU::submitVirtualMap(amd::VirtualMapCommand& vcmd) {
     dispatchBarrierPacket(kBarrierPacketHeader, false);
     Barriers().WaitCurrent();
 
-    size_t total_unmapped_buffers_size = 0;
-    auto sub_buffers = vaddr_base_obj->subBuffers();
-    for (auto buffer : sub_buffers) {
-      if (total_unmapped_buffers_size + buffer->getSize() <= vcmd.size()) {
-        // Unmap the object, since the physical addr is set.
-        if ((hsa_status = hsa_amd_vmem_unmap(buffer->getSvmPtr(), buffer->getSize())) ==
-            HSA_STATUS_SUCCESS) {
-          buffer->getContext().devices()[0]->DestroyVirtualBuffer(buffer);
-          amd::MemObjMap::RemoveMemObj(buffer->getSvmPtr());
-          if (buffer->getUserData().phys_mem_obj != nullptr) {
-            auto& phys_mem_user_data = buffer->getUserData().phys_mem_obj->getUserData();
-            phys_mem_user_data.vaddr_mem_obj = nullptr;
-            if (phys_mem_user_data.data != nullptr) {
-              reinterpret_cast<amd::RuntimeObject*>(phys_mem_user_data.data)->release();
-            }
-            buffer->getUserData().phys_mem_obj = nullptr;
-            total_unmapped_buffers_size += buffer->getSize();
-          }
-          buffer->release();
-        } else {
-          LogError("HSA Command: hsa_amd_vmem_unmap failed");
-        }
+    amd::Memory* vaddr_sub_obj = amd::MemObjMap::FindMemObj(vcmd.ptr());
+    assert(vaddr_sub_obj != nullptr);
+
+    // Unmap the object, since the physical addr is set.
+    if ((hsa_status = hsa_amd_vmem_unmap(vaddr_sub_obj->getSvmPtr(), vcmd.size()))
+                        == HSA_STATUS_SUCCESS) {
+      // assert the va is mapped and needs to be removed
+      vaddr_sub_obj->getContext().devices()[0]->DestroyVirtualBuffer(vaddr_sub_obj);
+      amd::MemObjMap::RemoveMemObj(vcmd.ptr());
+      if (vaddr_sub_obj->getUserData().phys_mem_obj != nullptr) {
+        vaddr_sub_obj->getUserData().phys_mem_obj->getUserData().vaddr_mem_obj = nullptr;
+        vaddr_sub_obj->getUserData().phys_mem_obj = nullptr;
       }
+    } else {
+      LogError("HSA Command: hsa_amd_vmem_unmap failed");
     }
   }