diff --git a/src/gallium/drivers/zink/zink_context.c b/src/gallium/drivers/zink/zink_context.c
index 7ebd1cc665d..2d9a1a6bf5a 100644
--- a/src/gallium/drivers/zink/zink_context.c
+++ b/src/gallium/drivers/zink/zink_context.c
@@ -2573,15 +2573,8 @@ zink_set_framebuffer_state(struct pipe_context *pctx,
    update_framebuffer_state(ctx, w, h);
 
    uint8_t rast_samples = ctx->fb_state.samples - 1;
-   /* update the shader key if applicable:
-    * if gl_SampleMask[] is written to, we have to ensure that we get a shader with the same sample count:
-    * in GL, rast_samples==1 means ignore gl_SampleMask[]
-    * in VK, gl_SampleMask[] is never ignored
-    */
-   if (rast_samples != ctx->gfx_pipeline_state.rast_samples &&
-       (!ctx->gfx_stages[PIPE_SHADER_FRAGMENT] ||
-        ctx->gfx_stages[PIPE_SHADER_FRAGMENT]->nir->info.outputs_written & (1 << FRAG_RESULT_SAMPLE_MASK)))
-      zink_set_fs_key(ctx)->samples = ctx->fb_state.samples > 0;
+   if (rast_samples != ctx->gfx_pipeline_state.rast_samples)
+      zink_update_fs_key_samples(ctx);
    if (ctx->gfx_pipeline_state.rast_samples != rast_samples) {
       ctx->sample_locations_changed |= ctx->gfx_pipeline_state.sample_locations_enabled;
       ctx->gfx_pipeline_state.dirty = true;
diff --git a/src/gallium/drivers/zink/zink_program.c b/src/gallium/drivers/zink/zink_program.c
index 393f2c10a78..fa5461dc990 100644
--- a/src/gallium/drivers/zink/zink_program.c
+++ b/src/gallium/drivers/zink/zink_program.c
@@ -854,6 +854,23 @@ zink_bind_vs_state(struct pipe_context *pctx,
 
 }
 
+/* if gl_SampleMask[] is written to, we have to ensure that we get a shader with the same sample count:
+ * in GL, samples==1 means ignore gl_SampleMask[]
+ * in VK, gl_SampleMask[] is never ignored
+ */
+void
+zink_update_fs_key_samples(struct zink_context *ctx)
+{
+   if (!ctx->gfx_stages[PIPE_SHADER_FRAGMENT])
+      return;
+   nir_shader *nir = ctx->gfx_stages[PIPE_SHADER_FRAGMENT]->nir;
+   if (nir->info.outputs_written & (1 << FRAG_RESULT_SAMPLE_MASK)) {
+      bool samples = zink_get_fs_key(ctx)->samples;
+      if (samples != (ctx->fb_state.samples > 1))
+         zink_set_fs_key(ctx)->samples = ctx->fb_state.samples > 1;
+   }
+}
+
 static void
 zink_bind_fs_state(struct pipe_context *pctx,
                    void *cso)
@@ -871,6 +888,7 @@ zink_bind_fs_state(struct pipe_context *pctx,
                ctx->fbfetch_outputs |= BITFIELD_BIT(var->data.location - FRAG_RESULT_DATA0);
          }
       }
+      zink_update_fs_key_samples(ctx);
    }
    zink_update_fbfetch(ctx);
 }
diff --git a/src/gallium/drivers/zink/zink_program.h b/src/gallium/drivers/zink/zink_program.h
index 7de44385c10..111207f9541 100644
--- a/src/gallium/drivers/zink/zink_program.h
+++ b/src/gallium/drivers/zink/zink_program.h
@@ -297,6 +297,9 @@ zink_get_fs_key(struct zink_context *ctx)
    return (const struct zink_fs_key *)&ctx->gfx_pipeline_state.shader_keys.key[PIPE_SHADER_FRAGMENT];
 }
 
+void
+zink_update_fs_key_samples(struct zink_context *ctx);
+
 static inline struct zink_vs_key *
 zink_set_vs_key(struct zink_context *ctx)
 {