diff --git a/src/intel/compiler/jay/jay_from_nir.c b/src/intel/compiler/jay/jay_from_nir.c
index f6a47870501..ed86e381656 100644
--- a/src/intel/compiler/jay/jay_from_nir.c
+++ b/src/intel/compiler/jay/jay_from_nir.c
@@ -2456,6 +2456,24 @@ read_vector_payload(struct payload_builder *b, enum jay_file file, unsigned len)
    return jay_collect_vectors(b->b, defs, len);
 }
 
+static void
+setup_payload_dispatch_start(struct nir_to_jay_state *nj,
+                             struct payload_builder *p)
+{
+   jay_shader *s = nj->s;
+
+   const unsigned start_grf = p->offsets[GPR] * jay_grf_per_gpr(nj->s) +
+                              p->offsets[UGPR] / jay_ugpr_per_grf(nj->s);
+
+   if (s->stage == MESA_SHADER_FRAGMENT && s->dispatch_width == 32) {
+      s->prog_data->fs.dispatch_grf_start_reg_32 = start_grf;
+   } else if (s->stage == MESA_SHADER_FRAGMENT && s->dispatch_width == 16) {
+      s->prog_data->fs.dispatch_grf_start_reg_16 = start_grf;
+   } else {
+      s->prog_data->base.dispatch_grf_start_reg = start_grf;
+   }
+}
+
 static void
 setup_payload_push(struct nir_to_jay_state *nj, struct payload_builder *p)
 {
@@ -2482,12 +2500,13 @@ setup_vertex_payload(struct nir_to_jay_state *nj, struct payload_builder *p)
     */
    p->offsets[GPR] += 7;
 
+   setup_payload_dispatch_start(nj, p);
+   setup_payload_push(nj, p);
+
    for (unsigned i = 0; i < (8 * nj->s->prog_data->vue.urb_read_length); ++i) {
       assert(i < ARRAY_SIZE(nj->payload.vs.attributes));
       nj->payload.vs.attributes[i] = read_payload(p, GPR);
    }
-
-   setup_payload_push(nj, p);
 }
 
 static void
@@ -2498,6 +2517,8 @@ setup_compute_payload(struct nir_to_jay_state *nj, struct payload_builder *p)
 
    nj->payload.inline_data =
       read_vector_payload(p, UGPR, jay_ugpr_per_grf(nj->s));
+
+   setup_payload_dispatch_start(nj, p);
 }
 
 static void
@@ -2553,8 +2574,6 @@ setup_fragment_payload(struct nir_to_jay_state *nj, struct payload_builder *p)
       nj->payload.u1 = read_vector_payload(p, UGPR, jay_ugpr_per_grf(nj->s));
    }
 
-   setup_payload_push(nj, p);
-
    u_foreach_bit(i, nj->s->prog_data->fs.barycentric_interp_modes) {
       fs->bary[i] = read_vector_payload(p, GPR, 2);
    }
@@ -2567,6 +2586,9 @@ setup_fragment_payload(struct nir_to_jay_state *nj, struct payload_builder *p)
       fs->coord.w = read_payload(p, GPR);
    }
 
+   setup_payload_dispatch_start(nj, p);
+   setup_payload_push(nj, p);
+
    unsigned nr_attribs = 16 * 4; /* TODO */
    for (unsigned i = 0; i < nr_attribs; ++i) {
       jay_def comps[] = { read_payload(p, UGPR), read_payload(p, UGPR),
diff --git a/src/intel/compiler/jay/jay_register_allocate.c b/src/intel/compiler/jay/jay_register_allocate.c
index 58dafccf60d..380a40d6814 100644
--- a/src/intel/compiler/jay/jay_register_allocate.c
+++ b/src/intel/compiler/jay/jay_register_allocate.c
@@ -1493,7 +1493,6 @@ jay_partition_grf(jay_shader *shader)
    assert((uniform_grfs + nonuniform_grfs) == JAY_NUM_PHYS_GRF);
 
    /* Partition GRFs between GPR & UGPR */
-   unsigned dispatch_grf = 0;
    unsigned stride4_header_size = 0;
 
    if (shader->stage == MESA_SHADER_VERTEX) {
@@ -1508,7 +1507,6 @@ jay_partition_grf(jay_shader *shader)
       };
 
       build_partition(shader, blocks, ARRAY_SIZE(blocks));
-      dispatch_grf = blocks[0] + blocks[1];
       stride4_header_size = blocks[1] + blocks[3];
    } else if (shader->stage == MESA_SHADER_FRAGMENT) {
       unsigned len0 = jay_grf_per_gpr(shader);
@@ -1521,7 +1519,6 @@ jay_partition_grf(jay_shader *shader)
          nonuniform_grfs - payload_grfs, /* GPR: General & EOT */
       };
       build_partition(shader, blocks, ARRAY_SIZE(blocks));
-      dispatch_grf = blocks[0] + blocks[1];
       stride4_header_size = blocks[1];
    } else {
       unsigned blocks[] = { uniform_grfs - 4, nonuniform_grfs, 4 };
@@ -1538,15 +1535,6 @@ jay_partition_grf(jay_shader *shader)
    // print_partition(p);
    validate_partition(p, stride4_header_size, nonuniform_gprs);
 
-   if (shader->stage == MESA_SHADER_FRAGMENT && shader->dispatch_width == 32) {
-      shader->prog_data->fs.dispatch_grf_start_reg_32 = dispatch_grf;
-   } else if (shader->stage == MESA_SHADER_FRAGMENT &&
-              shader->dispatch_width == 16) {
-      shader->prog_data->fs.dispatch_grf_start_reg_16 = dispatch_grf;
-   } else {
-      shader->prog_data->base.dispatch_grf_start_reg = dispatch_grf;
-   }
-
    /* By construction of our partition, the entire GRF is used. */
    shader->prog_data->base.grf_used = JAY_NUM_PHYS_GRF;