diff --git a/libhsakmt/src/queues.c b/libhsakmt/src/queues.c
index 5b1a0317bb..5a5f079bd2 100644
--- a/libhsakmt/src/queues.c
+++ b/libhsakmt/src/queues.c
@@ -363,11 +363,16 @@ void *allocate_exec_aligned_memory_gpu(uint32_t size, uint32_t align, uint32_t g
 
 	if (NodeId != 0) {
 		uint32_t nodes_array[1] = {NodeId};
+		HsaMemMapFlags map_flags = {0};
+		HSAKMT_STATUS result;
 
-		if (hsaKmtRegisterMemoryToNodes(mem, size, 1, nodes_array) != HSAKMT_STATUS_SUCCESS) {
+		result = hsaKmtMapMemoryToGPUNodes(mem, size, &gpu_va, map_flags, 1, nodes_array);
+		if (result != HSAKMT_STATUS_SUCCESS) {
 			hsaKmtFreeMemory(mem, size);
 			return NULL;
 		}
+
+		return mem;
 	}
 
 	if (hsaKmtMapMemoryToGPU(mem, size, &gpu_va) != HSAKMT_STATUS_SUCCESS) {
diff --git a/libhsakmt/tests/kfdtest/src/KFDLocalMemoryTest.cpp b/libhsakmt/tests/kfdtest/src/KFDLocalMemoryTest.cpp
index b37528c651..7f00b31d13 100644
--- a/libhsakmt/tests/kfdtest/src/KFDLocalMemoryTest.cpp
+++ b/libhsakmt/tests/kfdtest/src/KFDLocalMemoryTest.cpp
@@ -529,32 +529,26 @@ TEST_F(KFDLocalMemoryTest, MapVramToGPUNodesTest) {
     HsaMemMapFlags mapFlags = {0};
 
     EXPECT_SUCCESS(hsaKmtAllocMemory(nodes[1], PAGE_SIZE, memFlags, &shared_addr));
-    EXPECT_SUCCESS(hsaKmtRegisterMemoryToNodes(shared_addr, PAGE_SIZE, 2, nodes));
     EXPECT_SUCCESS(hsaKmtMapMemoryToGPUNodes(shared_addr, PAGE_SIZE, NULL, mapFlags, 2, nodes));
     EXPECT_SUCCESS(hsaKmtQueryPointerInfo(shared_addr, &info));
-    EXPECT_EQ(info.NRegisteredNodes, 2);
     EXPECT_EQ(info.NMappedNodes, 2);
 
     EXPECT_SUCCESS(hsaKmtMapMemoryToGPUNodes(shared_addr, PAGE_SIZE, NULL, mapFlags, 1, &nodes[0]));
     EXPECT_SUCCESS(hsaKmtQueryPointerInfo(shared_addr, &info));
-    EXPECT_EQ(info.NRegisteredNodes, 2);
     EXPECT_EQ(info.NMappedNodes, 1);
     EXPECT_EQ(info.MappedNodes[0], nodes[0]);
 
     EXPECT_SUCCESS(hsaKmtMapMemoryToGPUNodes(shared_addr, PAGE_SIZE, NULL, mapFlags, 1, &nodes[1]));
     EXPECT_SUCCESS(hsaKmtQueryPointerInfo(shared_addr, &info));
-    EXPECT_EQ(info.NRegisteredNodes, 2);
     EXPECT_EQ(info.NMappedNodes, 1);
     EXPECT_EQ(info.MappedNodes[0], nodes[1]);
 
     EXPECT_SUCCESS(hsaKmtUnmapMemoryToGPU(shared_addr));
     EXPECT_SUCCESS(hsaKmtQueryPointerInfo(shared_addr, &info));
-    EXPECT_EQ(info.NRegisteredNodes, 2);
     EXPECT_EQ(info.NMappedNodes, 0);
 
     EXPECT_SUCCESS(hsaKmtMapMemoryToGPUNodes(shared_addr, PAGE_SIZE, NULL, mapFlags, 1, &nodes[0]));
     EXPECT_SUCCESS(hsaKmtQueryPointerInfo(shared_addr, &info));
-    EXPECT_EQ(info.NRegisteredNodes, 2);
     EXPECT_EQ(info.NMappedNodes, 1);
     EXPECT_EQ(info.MappedNodes[0], nodes[0]);
 
diff --git a/libhsakmt/tests/kfdtest/src/KFDTestUtil.cpp b/libhsakmt/tests/kfdtest/src/KFDTestUtil.cpp
index d3539f97a9..a31630c510 100644
--- a/libhsakmt/tests/kfdtest/src/KFDTestUtil.cpp
+++ b/libhsakmt/tests/kfdtest/src/KFDTestUtil.cpp
@@ -468,14 +468,11 @@ unsigned int HsaMemoryBuffer::Node() const {
 }
 
 int HsaMemoryBuffer::MapMemToNodes(unsigned int *nodes, unsigned int nodes_num) {
+    HsaMemMapFlags mapFlags = {0};
     int ret, bit;
 
-    ret = hsaKmtRegisterMemoryToNodes(m_pBuf, m_Size, nodes_num, nodes);
-    if (ret != 0)
-        return ret;
-    ret = hsaKmtMapMemoryToGPU(m_pBuf, m_Size, NULL);
+    ret = hsaKmtMapMemoryToGPUNodes(m_pBuf, m_Size, NULL, mapFlags, nodes_num, nodes);
     if (ret != 0) {
-        hsaKmtDeregisterMemory(m_pBuf);
         return ret;
     }
 
@@ -494,7 +491,6 @@ int HsaMemoryBuffer::UnmapMemToNodes(unsigned int *nodes, unsigned int nodes_num
     if (ret)
         return ret;
 
-    hsaKmtDeregisterMemory(m_pBuf);
     for (unsigned int i = 0; i < nodes_num; i++) {
         bit = 1 << nodes[i];
         m_MappedNodes &= ~bit;
@@ -528,7 +524,6 @@ void HsaMemoryBuffer::UnmapAllNodes() {
      * TODO: When thunk is updated, use hsaKmtRegisterToNodes. Then nodes will be used
      */
     hsaKmtUnmapMemoryToGPU(m_pBuf);
-    hsaKmtDeregisterMemory(m_pBuf);
 
     m_MappedNodes = 0;
 
@@ -543,7 +538,6 @@ HsaMemoryBuffer::~HsaMemoryBuffer() {
         if (is_dgpu()) {
             if (m_MappedNodes) {
                 hsaKmtUnmapMemoryToGPU(m_pBuf);
-                hsaKmtDeregisterMemory(m_pBuf);
             }
         }
         hsaKmtFreeMemory(m_pBuf, m_Size);