diff --git a/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp b/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
index 861c10d69b..48506c60e9 100644
--- a/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
+++ b/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
@@ -399,9 +399,14 @@ hsa_status_t KfdDriver::ExportDMABuf(void *mem, size_t size, int *dmabuf_fd,
                                      size_t *offset) {
   int dmabuf_fd_res = -1;
   size_t offset_res = 0;
-  if (HSAKMT_CALL(hsaKmtExportDMABufHandle(mem, size, &dmabuf_fd_res, &offset_res)) !=
-      HSAKMT_STATUS_SUCCESS)
+  HSAKMT_STATUS status =
+      HSAKMT_CALL(hsaKmtExportDMABufHandle(mem, size, &dmabuf_fd_res, &offset_res));
+  if (status != HSAKMT_STATUS_SUCCESS) {
+    if (status == HSAKMT_STATUS_INVALID_PARAMETER) {
+      return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+    }
     return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
 
   *dmabuf_fd = dmabuf_fd_res;
   *offset = offset_res;
diff --git a/runtime/hsa-runtime/core/runtime/runtime.cpp b/runtime/hsa-runtime/core/runtime/runtime.cpp
index 84469ba5ea..ff52495d45 100644
--- a/runtime/hsa-runtime/core/runtime/runtime.cpp
+++ b/runtime/hsa-runtime/core/runtime/runtime.cpp
@@ -1266,9 +1266,9 @@ hsa_status_t Runtime::IPCCreate(void* ptr, size_t len, hsa_amd_ipc_memory_t* han
   // deferred export will not run into this problem.
   int dmabuf_fd;
   uint64_t dmabufOffset;
-  HSAKMT_STATUS err = HSAKMT_CALL(hsaKmtExportDMABufHandle(ptr, len, &dmabuf_fd, &dmabufOffset));
+  hsa_status_t err = agent->driver().ExportDMABuf(ptr, len, &dmabuf_fd, &dmabufOffset);
   assert(dmabufOffset/pageSize == fragOffset && "DMA Buf inconsistent with pointer offset.");
-  if (err != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR;
+  if (err != HSA_STATUS_SUCCESS) return err;
   close(dmabuf_fd);
 
   ScopedAcquire<KernelMutex> lock(&ipc_sock_server_lock_);
@@ -3104,17 +3104,18 @@ hsa_status_t Runtime::DmaBufExport(const void* ptr, size_t size, int* dmabuf,
         }
         int fd;
         uint64_t off;
-        HSAKMT_STATUS err = HSAKMT_CALL(hsaKmtExportDMABufHandle(const_cast<void*>(ptr), size, &fd, &off));
-        if (err == HSAKMT_STATUS_SUCCESS) {
-          *dmabuf = fd;
-          *offset = off;
-          return HSA_STATUS_SUCCESS;
+        hsa_status_t err = mem->second.region->owner()->driver().ExportDMABuf(
+            const_cast<void*>(ptr), size, &fd, &off);
+
+        if (err != HSA_STATUS_SUCCESS) {
+          assert((err != HSA_STATUS_ERROR_INVALID_ARGUMENT) &&
+                 "Thunk does not recognize an expected allocation.");
+          return err;
         }
 
-        assert((err != HSAKMT_STATUS_INVALID_PARAMETER) &&
-               "Thunk does not recognize an expected allocation.");
-        if (err == HSAKMT_STATUS_ERROR) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
-        return HSA_STATUS_ERROR;
+        *dmabuf = fd;
+        *offset = off;
+        return HSA_STATUS_SUCCESS;
       }
     }
   }
@@ -3653,11 +3654,11 @@ hsa_status_t Runtime::VMemoryExportShareableHandle(int* dmabuf_fd,
     return HSA_STATUS_ERROR_INVALID_ALLOCATION;
   }
 
-  uint64_t offset, ret;
+  uint64_t offset;
 
-  ret = HSAKMT_CALL(hsaKmtExportDMABufHandle(memoryHandle->second.thunk_handle, memoryHandle->second.size,
-                                 dmabuf_fd, &offset));
-  if (ret != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  hsa_status_t err = memoryHandle->second.region->owner()->driver().ExportDMABuf(
+      memoryHandle->second.thunk_handle, memoryHandle->second.size, dmabuf_fd, &offset);
+  if (err != HSA_STATUS_SUCCESS) return err;
 
   return HSA_STATUS_SUCCESS;
 }