diff --git a/src/amd/vulkan/radv_pipeline.c b/src/amd/vulkan/radv_pipeline.c
index 3b546851700..886b97976b9 100644
--- a/src/amd/vulkan/radv_pipeline.c
+++ b/src/amd/vulkan/radv_pipeline.c
@@ -626,11 +626,15 @@ radv_postprocess_nir(struct radv_device *device, const struct radv_pipeline_layo
    }
 
    /* Lower I/O intrinsics to memory instructions. */
+   bool is_last_vgt_stage =
+      (stage->info.stage == MESA_SHADER_VERTEX || stage->info.stage == MESA_SHADER_TESS_EVAL ||
+       stage->info.stage == MESA_SHADER_GEOMETRY || stage->info.stage == MESA_SHADER_MESH) &&
+      (stage->info.next_stage == MESA_SHADER_FRAGMENT || stage->info.next_stage == MESA_SHADER_NONE);
    bool io_to_mem = radv_nir_lower_io_to_mem(device, stage);
-   bool lowered_ngg = stage->info.is_ngg && stage->info.next_stage == MESA_SHADER_FRAGMENT;
+   bool lowered_ngg = stage->info.is_ngg && is_last_vgt_stage;
    if (lowered_ngg) {
       radv_lower_ngg(device, stage, pipeline_key);
-   } else if (stage->info.next_stage == MESA_SHADER_FRAGMENT) {
+   } else if (is_last_vgt_stage) {
       if (stage->stage != MESA_SHADER_GEOMETRY) {
          NIR_PASS_V(stage->nir, ac_nir_lower_legacy_vs, gfx_level,
                     stage->info.outinfo.clip_dist_mask | stage->info.outinfo.cull_dist_mask,
diff --git a/src/amd/vulkan/radv_pipeline_graphics.c b/src/amd/vulkan/radv_pipeline_graphics.c
index c8c48f0c1b9..403e357574c 100644
--- a/src/amd/vulkan/radv_pipeline_graphics.c
+++ b/src/amd/vulkan/radv_pipeline_graphics.c
@@ -2084,20 +2084,28 @@ radv_get_next_stage(gl_shader_stage stage, VkShaderStageFlagBits active_nir_stag
          return MESA_SHADER_TESS_CTRL;
       } else if (active_nir_stages & VK_SHADER_STAGE_GEOMETRY_BIT) {
          return MESA_SHADER_GEOMETRY;
-      } else {
+      } else if (active_nir_stages & VK_SHADER_STAGE_FRAGMENT_BIT) {
          return MESA_SHADER_FRAGMENT;
+      } else {
+         return MESA_SHADER_NONE;
       }
    case MESA_SHADER_TESS_CTRL:
       return MESA_SHADER_TESS_EVAL;
    case MESA_SHADER_TESS_EVAL:
       if (active_nir_stages & VK_SHADER_STAGE_GEOMETRY_BIT) {
          return MESA_SHADER_GEOMETRY;
-      } else {
+      } else if (active_nir_stages & VK_SHADER_STAGE_FRAGMENT_BIT) {
          return MESA_SHADER_FRAGMENT;
+      } else {
+         return MESA_SHADER_NONE;
       }
    case MESA_SHADER_GEOMETRY:
    case MESA_SHADER_MESH:
-      return MESA_SHADER_FRAGMENT;
+      if (active_nir_stages & VK_SHADER_STAGE_FRAGMENT_BIT) {
+         return MESA_SHADER_FRAGMENT;
+      } else {
+         return MESA_SHADER_NONE;
+      }
    case MESA_SHADER_TASK:
       return MESA_SHADER_MESH;
    case MESA_SHADER_FRAGMENT: