diff --git a/src/vulkan/runtime/vk_acceleration_structure.c b/src/vulkan/runtime/vk_acceleration_structure.c
index c87746b8379..eb82ad35a10 100644
--- a/src/vulkan/runtime/vk_acceleration_structure.c
+++ b/src/vulkan/runtime/vk_acceleration_structure.c
@@ -251,7 +251,7 @@ get_scratch_layout(struct vk_device *device,
    if (build_info->type == VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR &&
        device->as_build_ops->update_as[0]) {
       scratch->update_size =
-         device->as_build_ops->get_update_scratch_size(device, leaf_count);
+         device->as_build_ops->get_update_scratch_size(device, build_info, leaf_count);
    } else {
       scratch->update_size = offset;
    }
@@ -1117,7 +1117,7 @@ vk_cmd_build_acceleration_structures(VkCommandBuffer commandBuffer,
          VK_FROM_HANDLE(vk_acceleration_structure, dst_as, pInfos[i].dstAccelerationStructure);
 
          ops->init_update_scratch(commandBuffer, pInfos[i].scratchData.deviceAddress,
-                                  leaf_node_count, src_as, dst_as);
+                                  &pInfos[i], ppBuildRangeInfos[i], leaf_node_count, src_as, dst_as);
       }
    }
 
diff --git a/src/vulkan/runtime/vk_acceleration_structure.h b/src/vulkan/runtime/vk_acceleration_structure.h
index 14f9de7c98e..11e8e16d780 100644
--- a/src/vulkan/runtime/vk_acceleration_structure.h
+++ b/src/vulkan/runtime/vk_acceleration_structure.h
@@ -76,7 +76,9 @@ struct vk_acceleration_structure_build_ops {
    VkDeviceSize (*get_as_size)(VkDevice device,
                                const VkAccelerationStructureBuildGeometryInfoKHR *pBuildInfo,
                                uint32_t leaf_count);
-   VkDeviceSize (*get_update_scratch_size)(struct vk_device *device, uint32_t leaf_count);
+   VkDeviceSize (*get_update_scratch_size)(struct vk_device *device,
+                                           const VkAccelerationStructureBuildGeometryInfoKHR *build_info,
+                                           uint32_t leaf_count);
    uint32_t (*get_encode_key[MAX_ENCODE_PASSES])(struct vk_device *device,
                                                  VkAccelerationStructureTypeKHR type,
                                                  VkBuildAccelerationStructureFlagBitsKHR flags);
@@ -93,6 +95,8 @@ struct vk_acceleration_structure_build_ops {
                                         struct vk_acceleration_structure *dst);
    void (*init_update_scratch)(VkCommandBuffer cmd_buffer,
                                VkDeviceAddress scratch,
+                               const VkAccelerationStructureBuildGeometryInfoKHR *build_info,
+                               const VkAccelerationStructureBuildRangeInfoKHR *build_range_infos,
                                uint32_t leaf_count,
                                struct vk_acceleration_structure *src_as,
                                struct vk_acceleration_structure *dst_as);