diff --git a/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc b/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc
index f88dd8b58..54c9c7b90 100644
--- a/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc
+++ b/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc
@@ -197,7 +197,8 @@ static iree_status_t iree_hal_xrt_allocator_allocate_buffer(
         xrt_buffer.release(), placement, compat_params.type,
         compat_params.access, compat_params.usage, allocation_size,
         /*byte_offset=*/0, /*byte_length=*/allocation_size,
-        iree_hal_buffer_release_callback_null(), &buffer);
+        iree_hal_buffer_release_callback_null(), allocator->host_allocator,
+        &buffer);
   }
 
   if (iree_status_is_ok(status)) {
diff --git a/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.cc b/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.cc
index 52e740ec7..3cfdfffd8 100644
--- a/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.cc
+++ b/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.cc
@@ -39,13 +39,10 @@ iree_status_t iree_hal_xrt_buffer_wrap(
     iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size,
     iree_device_size_t byte_offset, iree_device_size_t byte_length,
     iree_hal_buffer_release_callback_t release_callback,
-    iree_hal_buffer_t** out_buffer) {
-  IREE_ASSERT_ARGUMENT(allocator);
+    iree_allocator_t host_allocator, iree_hal_buffer_t** out_buffer) {
   IREE_ASSERT_ARGUMENT(out_buffer);
   IREE_TRACE_ZONE_BEGIN(z0);
 
-  iree_allocator_t host_allocator =
-      iree_hal_allocator_host_allocator(allocator);
   iree_hal_xrt_buffer_t* buffer = nullptr;
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0,
diff --git a/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.h b/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.h
index 27bcaa03e..e7f39bca9 100644
--- a/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.h
+++ b/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.h
@@ -23,7 +23,7 @@ iree_status_t iree_hal_xrt_buffer_wrap(
     iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size,
     iree_device_size_t byte_offset, iree_device_size_t byte_length,
     iree_hal_buffer_release_callback_t release_callback,
-    iree_hal_buffer_t** out_buffer);
+    iree_allocator_t host_allocator, iree_hal_buffer_t** out_buffer);
 
 // Returns the underlying XRT buffer handle for the given |buffer|.
 xrt::bo* iree_hal_xrt_buffer_handle(const iree_hal_buffer_t* buffer);