diff --git a/src/virtio/vulkan/vn_android.c b/src/virtio/vulkan/vn_android.c
index f4062d2e547..e0794e153e0 100644
--- a/src/virtio/vulkan/vn_android.c
+++ b/src/virtio/vulkan/vn_android.c
@@ -766,10 +766,9 @@ vn_android_get_drm_format_modifier_info(
 }
 
 VkResult
-vn_android_device_import_ahb(
-   struct vn_device *dev,
-   struct vn_device_memory *mem,
-   const struct VkMemoryDedicatedAllocateInfo *dedicated_info)
+vn_android_device_import_ahb(struct vn_device *dev,
+                             struct vn_device_memory *mem,
+                             const struct VkMemoryAllocateInfo *alloc_info)
 {
    const struct vk_device_memory *mem_vk = &mem->base.vk;
    const native_handle_t *handle = NULL;
@@ -790,6 +789,9 @@ vn_android_device_import_ahb(
    if (result != VK_SUCCESS)
       return result;
 
+   const VkMemoryDedicatedAllocateInfo *dedicated_info =
+      vk_find_struct_const(alloc_info->pNext, MEMORY_DEDICATED_ALLOCATE_INFO);
+
    /* If ahb is for an image, finish the deferred image creation first */
    if (dedicated_info && dedicated_info->image != VK_NULL_HANDLE) {
       struct vn_image *img = vn_image_from_handle(dedicated_info->image);
diff --git a/src/virtio/vulkan/vn_android.h b/src/virtio/vulkan/vn_android.h
index 1d85ed77de7..e71cca6f4ed 100644
--- a/src/virtio/vulkan/vn_android.h
+++ b/src/virtio/vulkan/vn_android.h
@@ -43,10 +43,9 @@ vn_android_get_ahb_usage(const VkImageUsageFlags usage,
                          const VkImageCreateFlags flags);
 
 VkResult
-vn_android_device_import_ahb(
-   struct vn_device *dev,
-   struct vn_device_memory *mem,
-   const struct VkMemoryDedicatedAllocateInfo *dedicated_info);
+vn_android_device_import_ahb(struct vn_device *dev,
+                             struct vn_device_memory *mem,
+                             const struct VkMemoryAllocateInfo *alloc_info);
 
 VkFormat
 vn_android_drm_format_to_vk_format(uint32_t format);
@@ -100,7 +99,7 @@ static inline VkResult
 vn_android_device_import_ahb(
    UNUSED struct vn_device *dev,
    UNUSED struct vn_device_memory *mem,
-   UNUSED const struct VkMemoryDedicatedAllocateInfo *dedicated_info)
+   UNUSED const struct VkMemoryAllocateInfo *alloc_info)
 {
    return VK_ERROR_OUT_OF_HOST_MEMORY;
 }
diff --git a/src/virtio/vulkan/vn_device_memory.c b/src/virtio/vulkan/vn_device_memory.c
index 2225c04a63a..f3859524cd4 100644
--- a/src/virtio/vulkan/vn_device_memory.c
+++ b/src/virtio/vulkan/vn_device_memory.c
@@ -360,21 +360,6 @@ vn_AllocateMemory(VkDevice device,
 {
    struct vn_device *dev = vn_device_from_handle(device);
 
-   const VkImportMemoryFdInfoKHR *import_fd_info = NULL;
-   const VkMemoryDedicatedAllocateInfo *dedicated_info = NULL;
-   vk_foreach_struct_const(pnext, pAllocateInfo->pNext) {
-      switch (pnext->sType) {
-      case VK_STRUCTURE_TYPE_IMPORT_MEMORY_FD_INFO_KHR:
-         import_fd_info = (const void *)pnext;
-         break;
-      case VK_STRUCTURE_TYPE_MEMORY_DEDICATED_ALLOCATE_INFO:
-         dedicated_info = (const void *)pnext;
-         break;
-      default:
-         break;
-      }
-   }
-
    struct vn_device_memory *mem = vk_device_memory_create(
       &dev->base.vk, pAllocateInfo, pAllocator, sizeof(*mem));
    if (!mem)
@@ -382,9 +367,12 @@ vn_AllocateMemory(VkDevice device,
 
    vn_object_set_id(mem, vn_get_next_obj_id(), VK_OBJECT_TYPE_DEVICE_MEMORY);
 
+   const VkImportMemoryFdInfoKHR *import_fd_info =
+      vk_find_struct_const(pAllocateInfo->pNext, IMPORT_MEMORY_FD_INFO_KHR);
+
    VkResult result;
    if (mem->base.vk.ahardware_buffer) {
-      result = vn_android_device_import_ahb(dev, mem, dedicated_info);
+      result = vn_android_device_import_ahb(dev, mem, pAllocateInfo);
    } else if (import_fd_info) {
       result = vn_device_memory_import_dma_buf(dev, mem, pAllocateInfo,
                                                import_fd_info->fd);