diff --git a/src/amd/vulkan/radv_cmd_buffer.c b/src/amd/vulkan/radv_cmd_buffer.c
index 0de11824df5..a027ceb6497 100644
--- a/src/amd/vulkan/radv_cmd_buffer.c
+++ b/src/amd/vulkan/radv_cmd_buffer.c
@@ -3803,7 +3803,13 @@ lookup_vs_prolog(struct radv_cmd_buffer *cmd_buffer, const struct radv_shader *v
    key.as_ls = vs_shader->info.vs.as_ls && instance_rate_inputs;
    key.is_ngg = vs_shader->info.is_ngg;
    key.wave32 = vs_shader->info.wave_size == 32;
-   key.next_stage = vs_shader->info.stage;
+
+   if (vs_shader->info.merged_shader_compiled_separately) {
+      assert(vs_shader->info.next_stage == MESA_SHADER_TESS_CTRL || vs_shader->info.next_stage == MESA_SHADER_GEOMETRY);
+      key.next_stage = vs_shader->info.next_stage;
+   } else {
+      key.next_stage = vs_shader->info.stage;
+   }
 
    union vs_prolog_key_header header;
    header.v = 0;