diff --git a/hw/xbox/nv2a/pgraph/gl/shaders.c b/hw/xbox/nv2a/pgraph/gl/shaders.c
index 28476aa66a..64c6f4e561 100644
--- a/hw/xbox/nv2a/pgraph/gl/shaders.c
+++ b/hw/xbox/nv2a/pgraph/gl/shaders.c
@@ -783,7 +783,7 @@ static void shader_update_constants(PGRAPHState *pg, ShaderBinding *binding,
         }
         if (binding->color_key_mask_loc[i] != -1) {
             glUniform1ui(binding->color_key_mask_loc[i],
-                         state->psh.colorkey_mask[i]);
+                         pgraph_get_color_key_mask_for_texture(pg, i));
         }
     }
 
diff --git a/hw/xbox/nv2a/pgraph/psh.h b/hw/xbox/nv2a/pgraph/psh.h
index 0b4df4c280..53a74c8781 100644
--- a/hw/xbox/nv2a/pgraph/psh.h
+++ b/hw/xbox/nv2a/pgraph/psh.h
@@ -70,7 +70,6 @@ typedef struct PshState {
     bool compare_mode[4][4];
     bool alphakill[4];
     int colorkey_mode[4];
-    uint32_t colorkey_mask[4];
     enum ConvolutionFilter conv_tex[4];
     bool tex_x8y24[4];
     int dim_tex[4];
diff --git a/hw/xbox/nv2a/pgraph/shaders.c b/hw/xbox/nv2a/pgraph/shaders.c
index 9fc28e773e..73b0357690 100644
--- a/hw/xbox/nv2a/pgraph/shaders.c
+++ b/hw/xbox/nv2a/pgraph/shaders.c
@@ -41,6 +41,14 @@ static uint32_t get_colorkey_mask(unsigned int color_format)
     }
 }
 
+uint32_t pgraph_get_color_key_mask_for_texture(PGRAPHState *pg, int i)
+{
+    assert(i < NV2A_MAX_TEXTURES);
+    uint32_t fmt = pgraph_reg_r(pg, NV_PGRAPH_TEXFMT0 + i*4);
+    unsigned int color_format = GET_MASK(fmt, NV_PGRAPH_TEXFMT0_COLOR);
+    return get_colorkey_mask(color_format);
+}
+
 ShaderState pgraph_get_shader_state(PGRAPHState *pg)
 {
     bool vertex_program = GET_MASK(pgraph_reg_r(pg, NV_PGRAPH_CSV0_D),
@@ -246,7 +254,6 @@ ShaderState pgraph_get_shader_state(PGRAPHState *pg)
         state.psh.rect_tex[i] = f.linear;
         state.psh.tex_x8y24[i] = color_format == NV097_SET_TEXTURE_FORMAT_COLOR_LU_IMAGE_DEPTH_X8_Y24_FIXED ||
                                 color_format == NV097_SET_TEXTURE_FORMAT_COLOR_LU_IMAGE_DEPTH_X8_Y24_FLOAT;
-        state.psh.colorkey_mask[i] = get_colorkey_mask(color_format);
 
         uint32_t border_source =
             GET_MASK(tex_fmt, NV_PGRAPH_TEXFMT0_BORDER_SOURCE);
diff --git a/hw/xbox/nv2a/pgraph/shaders.h b/hw/xbox/nv2a/pgraph/shaders.h
index 4cc07e3a9d..f7bd9d80fa 100644
--- a/hw/xbox/nv2a/pgraph/shaders.h
+++ b/hw/xbox/nv2a/pgraph/shaders.h
@@ -112,5 +112,6 @@ typedef struct ShaderState {
 typedef struct PGRAPHState PGRAPHState;
 
 ShaderState pgraph_get_shader_state(PGRAPHState *pg);
+uint32_t pgraph_get_color_key_mask_for_texture(PGRAPHState *pg, int i);
 
 #endif
diff --git a/hw/xbox/nv2a/pgraph/vk/shaders.c b/hw/xbox/nv2a/pgraph/vk/shaders.c
index 466789a8ad..1cffbd0a52 100644
--- a/hw/xbox/nv2a/pgraph/vk/shaders.c
+++ b/hw/xbox/nv2a/pgraph/vk/shaders.c
@@ -506,10 +506,7 @@ static void shader_update_constants(PGRAPHState *pg, ShaderBinding *binding,
         uniform1uiv(&binding->fragment->uniforms, binding->color_key_loc, 4,
                     color_key_colors);
     }
-    if (binding->color_key_mask_loc != -1) {
-        uniform1uiv(&binding->fragment->uniforms, binding->color_key_mask_loc,
-                    4, state->psh.colorkey_mask);
-    }
+    uint32_t color_key_mask[4] = { 0 };
 
     /* For each texture stage */
     for (int i = 0; i < NV2A_MAX_TEXTURES; i++) {
@@ -557,6 +554,13 @@ static void shader_update_constants(PGRAPHState *pg, ShaderBinding *binding,
             }
             uniform1f(&binding->fragment->uniforms, loc, scale);
         }
+
+        color_key_mask[i] = pgraph_get_color_key_mask_for_texture(pg, i);
+    }
+
+    if (binding->color_key_mask_loc != -1) {
+        uniform1uiv(&binding->fragment->uniforms, binding->color_key_mask_loc,
+                    4, color_key_mask);
     }
 
     if (binding->fog_color_loc != -1) {