diff --git a/runtime/src/iree-amd-aie/driver/xrt-lite/allocator.cc b/runtime/src/iree-amd-aie/driver/xrt-lite/allocator.cc
index 275781e67..f19896076 100644
--- a/runtime/src/iree-amd-aie/driver/xrt-lite/allocator.cc
+++ b/runtime/src/iree-amd-aie/driver/xrt-lite/allocator.cc
@@ -95,10 +95,14 @@ static iree_status_t iree_hal_xrt_lite_allocator_allocate_buffer(
   shim_xdna::bo* bo =
       allocator->shim_device->alloc_bo(allocation_size, flags).release();
   iree_hal_buffer_t* buffer = nullptr;
+  const iree_hal_buffer_placement_t placement = {
+      .queue_affinity = params->queue_affinity ? params->queue_affinity
+                                               : IREE_HAL_QUEUE_AFFINITY_ANY,
+      .flags = IREE_HAL_BUFFER_PLACEMENT_FLAG_NONE,
+  };
   iree_status_t status = iree_hal_xrt_lite_buffer_wrap(
-      bo, reinterpret_cast<iree_hal_allocator_t*>(allocator),
-      compat_params.type, compat_params.access, compat_params.usage,
-      allocation_size,
+      bo, 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(), allocator->host_allocator,
       &buffer);
diff --git a/runtime/src/iree-amd-aie/driver/xrt-lite/buffer.cc b/runtime/src/iree-amd-aie/driver/xrt-lite/buffer.cc
index 6c0dff4f1..a753add12 100644
--- a/runtime/src/iree-amd-aie/driver/xrt-lite/buffer.cc
+++ b/runtime/src/iree-amd-aie/driver/xrt-lite/buffer.cc
@@ -16,6 +16,7 @@ extern const iree_hal_buffer_vtable_t iree_hal_xrt_lite_buffer_vtable;
 struct iree_hal_xrt_lite_buffer {
   iree_hal_buffer_t base;
   shim_xdna::bo* bo;
+  iree_allocator_t host_allocator;
   iree_hal_buffer_release_callback_t release_callback;
 };
 
@@ -111,7 +112,7 @@ static iree_status_t iree_hal_xrt_lite_buffer_unmap_range(
 }
 
 iree_status_t iree_hal_xrt_lite_buffer_wrap(
-    shim_xdna::bo* bo, iree_hal_allocator_t* allocator,
+    shim_xdna::bo* bo, iree_hal_buffer_placement_t placement,
     iree_hal_memory_type_t memory_type, iree_hal_memory_access_t allowed_access,
     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,
@@ -125,9 +126,9 @@ iree_status_t iree_hal_xrt_lite_buffer_wrap(
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0, iree_allocator_malloc(host_allocator, sizeof(*buffer),
                                 reinterpret_cast<void**>(&buffer)));
-  iree_hal_buffer_initialize(host_allocator, allocator, &buffer->base,
-                             allocation_size, byte_offset, byte_length,
-                             memory_type, allowed_access, allowed_usage,
+  iree_hal_buffer_initialize(placement, &buffer->base, allocation_size,
+                             byte_offset, byte_length, memory_type,
+                             allowed_access, allowed_usage,
                              &iree_hal_xrt_lite_buffer_vtable, &buffer->base);
   buffer->release_callback = release_callback;
   buffer->bo = bo;
@@ -142,7 +143,7 @@ static void iree_hal_xrt_lite_buffer_destroy(iree_hal_buffer_t* base_buffer) {
       base_buffer, iree_hal_xrt_lite_buffer_vtable, iree_hal_xrt_lite_buffer);
   IREE_TRACE_ZONE_BEGIN(z0);
 
-  iree_allocator_t host_allocator = base_buffer->host_allocator;
+  iree_allocator_t host_allocator = buffer->host_allocator;
   if (buffer->release_callback.fn) {
     buffer->release_callback.fn(buffer->release_callback.user_data,
                                 base_buffer);
diff --git a/runtime/src/iree-amd-aie/driver/xrt-lite/buffer.h b/runtime/src/iree-amd-aie/driver/xrt-lite/buffer.h
index b70517d1c..7e429e4dd 100644
--- a/runtime/src/iree-amd-aie/driver/xrt-lite/buffer.h
+++ b/runtime/src/iree-amd-aie/driver/xrt-lite/buffer.h
@@ -12,7 +12,7 @@
 #include "iree/hal/api.h"
 
 iree_status_t iree_hal_xrt_lite_buffer_wrap(
-    shim_xdna::bo* bo, iree_hal_allocator_t* allocator,
+    shim_xdna::bo* bo, iree_hal_buffer_placement_t placement,
     iree_hal_memory_type_t memory_type, iree_hal_memory_access_t allowed_access,
     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,
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 be2370d2f..f88dd8b58 100644
--- a/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc
+++ b/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc
@@ -187,9 +187,14 @@ static iree_status_t iree_hal_xrt_allocator_allocate_buffer(
   }
 
   iree_hal_buffer_t* buffer = nullptr;
+  const iree_hal_buffer_placement_t placement = {
+      .queue_affinity = params->queue_affinity ? params->queue_affinity
+                                               : IREE_HAL_QUEUE_AFFINITY_ANY,
+      .flags = IREE_HAL_BUFFER_PLACEMENT_FLAG_NONE,
+  };
   if (iree_status_is_ok(status)) {
     status = iree_hal_xrt_buffer_wrap(
-        xrt_buffer.release(), base_allocator, compat_params.type,
+        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);
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 66429230a..52e740ec7 100644
--- a/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.cc
+++ b/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.cc
@@ -13,6 +13,7 @@
 typedef struct iree_hal_xrt_buffer_t {
   iree_hal_buffer_t base;
   xrt::bo* buffer;
+  iree_allocator_t host_allocator;
   iree_hal_buffer_release_callback_t release_callback;
 } iree_hal_xrt_buffer_t;
 
@@ -33,7 +34,7 @@ static const iree_hal_xrt_buffer_t* iree_hal_xrt_buffer_const_cast(
 }
 
 iree_status_t iree_hal_xrt_buffer_wrap(
-    xrt::bo* xrt_buffer, iree_hal_allocator_t* allocator,
+    xrt::bo* xrt_buffer, iree_hal_buffer_placement_t placement,
     iree_hal_memory_type_t memory_type, iree_hal_memory_access_t allowed_access,
     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,
@@ -49,9 +50,9 @@ iree_status_t iree_hal_xrt_buffer_wrap(
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0,
       iree_allocator_malloc(host_allocator, sizeof(*buffer), (void**)&buffer));
-  iree_hal_buffer_initialize(host_allocator, allocator, &buffer->base,
-                             allocation_size, byte_offset, byte_length,
-                             memory_type, allowed_access, allowed_usage,
+  iree_hal_buffer_initialize(placement, &buffer->base, allocation_size,
+                             byte_offset, byte_length, memory_type,
+                             allowed_access, allowed_usage,
                              &iree_hal_xrt_buffer_vtable, &buffer->base);
   buffer->buffer = xrt_buffer;
   buffer->release_callback = release_callback;
@@ -62,7 +63,7 @@ iree_status_t iree_hal_xrt_buffer_wrap(
 
 static void iree_hal_xrt_buffer_destroy(iree_hal_buffer_t* base_buffer) {
   iree_hal_xrt_buffer_t* buffer = iree_hal_xrt_buffer_cast(base_buffer);
-  iree_allocator_t host_allocator = base_buffer->host_allocator;
+  iree_allocator_t host_allocator = buffer->host_allocator;
   IREE_TRACE_ZONE_BEGIN(z0);
   IREE_TRACE_ZONE_APPEND_VALUE_I64(
       z0, (int64_t)iree_hal_buffer_allocation_size(base_buffer));
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 a31bcf204..27bcaa03e 100644
--- a/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.h
+++ b/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.h
@@ -18,7 +18,7 @@ extern "C" {
 // Wraps a XRT allocation in an iree_hal_buffer_t by retaining |xrt_buffer|.
 // |out_buffer| must be released by the caller (see iree_hal_buffer_release).
 iree_status_t iree_hal_xrt_buffer_wrap(
-    xrt::bo* xrt_buffer, iree_hal_allocator_t* allocator,
+    xrt::bo* xrt_buffer, iree_hal_buffer_placement_t placement,
     iree_hal_memory_type_t memory_type, iree_hal_memory_access_t allowed_access,
     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,
diff --git a/third_party/iree b/third_party/iree
index ecd87d868..eff067156 160000
--- a/third_party/iree
+++ b/third_party/iree
@@ -1 +1 @@
-Subproject commit ecd87d8680c1e79819d4b5914c092c064af32787
+Subproject commit eff06715692fe9642e4b8902e3dbc843c375218b