diff --git a/src/gallium/drivers/zink/zink_program.c b/src/gallium/drivers/zink/zink_program.c
index 7db2b505cfd..ec2d5e1b639 100644
--- a/src/gallium/drivers/zink/zink_program.c
+++ b/src/gallium/drivers/zink/zink_program.c
@@ -160,12 +160,17 @@ get_shader_module_for_stage(struct zink_context *ctx, struct zink_screen *screen
       ralloc_free(keybox);
       zm = entry->data;
    } else {
-      zm = CALLOC_STRUCT(zink_shader_module);
+      zm = malloc(sizeof(struct zink_shader_module) + key->size);
       if (!zm) {
          ralloc_free(keybox);
          return NULL;
       }
       zm->hash = hash;
+      zm->num_uniforms = base_size;
+      zm->key_size = key->size;
+      if (base_size)
+         memcpy(zm->uniforms, &key->base, base_size * sizeof(uint32_t));
+      memcpy(zm->key, key, key->size);
       mod = zink_shader_compile(screen, zs, prog->nir[stage], key);
       if (!mod) {
          ralloc_free(keybox);
diff --git a/src/gallium/drivers/zink/zink_program.h b/src/gallium/drivers/zink/zink_program.h
index 3ccbaa21a24..bd390a1aab5 100644
--- a/src/gallium/drivers/zink/zink_program.h
+++ b/src/gallium/drivers/zink/zink_program.h
@@ -67,6 +67,10 @@ struct zink_shader_module {
    VkShaderModule shader;
    uint32_t hash;
    bool default_variant;
+   uint8_t num_uniforms;
+   uint8_t key_size;
+   uint32_t uniforms[MAX_INLINABLE_UNIFORMS];
+   uint8_t key[0];
 };
 
 struct zink_program {