diff --git a/drivers/gpu/nvgpu/common/gr/gr_setup.c b/drivers/gpu/nvgpu/common/gr/gr_setup.c
index bcd4afe0b..b6f4d9801 100644
--- a/drivers/gpu/nvgpu/common/gr/gr_setup.c
+++ b/drivers/gpu/nvgpu/common/gr/gr_setup.c
@@ -295,8 +295,8 @@ static bool nvgpu_gr_setup_validate_preemption_mode(u32 *graphics_preempt_mode,
 
 
 int nvgpu_gr_setup_set_preemption_mode(struct nvgpu_channel *ch,
-					u32 graphics_preempt_mode,
-					u32 compute_preempt_mode)
+		u32 graphics_preempt_mode, u32 compute_preempt_mode,
+		u32 gr_instance_id)
 {
 	struct nvgpu_gr_ctx *gr_ctx;
 	struct gk20a *g = ch->g;
@@ -306,7 +306,7 @@ int nvgpu_gr_setup_set_preemption_mode(struct nvgpu_channel *ch,
 	u32 class_num;
 	int err = 0;
 
-	gr = nvgpu_gr_get_cur_instance_ptr(g);
+	gr = &g->gr[gr_instance_id];
 
 	class_num = ch->obj_class;
 	if (class_num == 0U) {
diff --git a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
index 51872c896..a37329da3 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
+++ b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
@@ -1470,8 +1470,8 @@ fail:
 }
 
 int vgpu_gr_set_preemption_mode(struct nvgpu_channel *ch,
-				u32 graphics_preempt_mode,
-				u32 compute_preempt_mode)
+		u32 graphics_preempt_mode, u32 compute_preempt_mode,
+		u32 gr_instance_id)
 {
 	struct nvgpu_gr_ctx *gr_ctx;
 	struct gk20a *g = ch->g;
diff --git a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.h b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.h
index 6ab71f7a4..c6bf6f4da 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.h
+++ b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.h
@@ -86,8 +86,8 @@ int vgpu_gr_init_sm_id_table(struct gk20a *g,
 int vgpu_gr_update_pc_sampling(struct nvgpu_channel *ch, bool enable);
 void vgpu_gr_init_cyclestats(struct gk20a *g);
 int vgpu_gr_set_preemption_mode(struct nvgpu_channel *ch,
-				u32 graphics_preempt_mode,
-				u32 compute_preempt_mode);
+		u32 graphics_preempt_mode, u32 compute_preempt_mode,
+		u32 gr_instance_id);
 int vgpu_gr_isr(struct gk20a *g, struct tegra_vgpu_gr_intr_info *info);
 void vgpu_gr_handle_sm_esr_event(struct gk20a *g,
 			struct tegra_vgpu_sm_esr_info *info);
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gops/gr.h b/drivers/gpu/nvgpu/include/nvgpu/gops/gr.h
index c825feff9..8e4dfe5a2 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gops/gr.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gops/gr.h
@@ -204,6 +204,8 @@ struct gops_gr_setup {
 	 *					preemption mode.
 	 * @param compute_preempt_mode [in]	Requested compute
 	 *					preemption mode.
+	 * @param gr_instance_id[in]		Instance ID of
+	 * 					graphics engine.
 	 *
 	 * This function will program newly requested preemption modes
 	 * into GR engine context image.
@@ -223,7 +225,8 @@ struct gops_gr_setup {
 	 */
 	int (*set_preemption_mode)(struct nvgpu_channel *ch,
 				   u32 graphics_preempt_mode,
-				   u32 compute_preempt_mode);
+				   u32 compute_preempt_mode,
+				   u32 gr_instance_id);
 
 	/** @cond DOXYGEN_SHOULD_SKIP_THIS */
 #ifdef CONFIG_NVGPU_GRAPHICS
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gr/setup.h b/drivers/gpu/nvgpu/include/nvgpu/gr/setup.h
index d6db27bf9..0a3e52fda 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gr/setup.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gr/setup.h
@@ -108,6 +108,7 @@ void nvgpu_gr_setup_free_subctx(struct nvgpu_channel *c);
  * @param ch [in]			Pointer to GPU channel.
  * @param graphics_preempt_mode [in]	Requested graphics preemption mode.
  * @param compute_preempt_mode [in]	Requested compute preemption mode.
+ * @param gr_instance_id [in]		Instance ID of the graphics engine.
  *
  * This function will program newly requested preemption modes into
  * GR engine context image. This function is typically needed if user
@@ -122,8 +123,8 @@ void nvgpu_gr_setup_free_subctx(struct nvgpu_channel *c);
  * @retval -EINVAL if invalid GPU channel pointer is provided.
  */
 int nvgpu_gr_setup_set_preemption_mode(struct nvgpu_channel *ch,
-					u32 graphics_preempt_mode,
-					u32 compute_preempt_mode);
+		u32 graphics_preempt_mode, u32 compute_preempt_mode,
+		u32 gr_instance_id);
 
 #ifdef CONFIG_NVGPU_GRAPHICS
 int nvgpu_gr_setup_bind_ctxsw_zcull(struct gk20a *g, struct nvgpu_channel *c,
diff --git a/drivers/gpu/nvgpu/os/linux/ioctl_channel.c b/drivers/gpu/nvgpu/os/linux/ioctl_channel.c
index bffe6a242..8e417d85c 100644
--- a/drivers/gpu/nvgpu/os/linux/ioctl_channel.c
+++ b/drivers/gpu/nvgpu/os/linux/ioctl_channel.c
@@ -1086,7 +1086,8 @@ static u32 nvgpu_get_common_compute_preempt_mode(u32 compute_preempt_mode)
 }
 
 static int nvgpu_ioctl_channel_set_preemption_mode(struct nvgpu_channel *ch,
-	u32 graphics_preempt_mode, u32 compute_preempt_mode)
+		u32 graphics_preempt_mode, u32 compute_preempt_mode,
+		u32 gr_instance_id)
 {
 	int err;
 
@@ -1098,7 +1099,8 @@ static int nvgpu_ioctl_channel_set_preemption_mode(struct nvgpu_channel *ch,
 		}
 		err = ch->g->ops.gr.setup.set_preemption_mode(ch,
 			nvgpu_get_common_graphics_preempt_mode(graphics_preempt_mode),
-			nvgpu_get_common_compute_preempt_mode(compute_preempt_mode));
+			nvgpu_get_common_compute_preempt_mode(compute_preempt_mode),
+			gr_instance_id);
 		gk20a_idle(ch->g);
 	} else {
 		err = -EINVAL;
@@ -1468,7 +1470,8 @@ long gk20a_channel_ioctl(struct file *filp,
 	case NVGPU_IOCTL_CHANNEL_SET_PREEMPTION_MODE:
 		err = nvgpu_ioctl_channel_set_preemption_mode(ch,
 		     ((struct nvgpu_preemption_mode_args *)buf)->graphics_preempt_mode,
-		     ((struct nvgpu_preemption_mode_args *)buf)->compute_preempt_mode);
+		     ((struct nvgpu_preemption_mode_args *)buf)->compute_preempt_mode,
+			gr_instance_id);
 		break;
 	case NVGPU_IOCTL_CHANNEL_SET_BOOSTED_CTX:
 		if (ch->g->ops.gr.set_boosted_ctx) {
diff --git a/userspace/units/gr/setup/nvgpu-gr-setup.c b/userspace/units/gr/setup/nvgpu-gr-setup.c
index 649d1e224..36c8fc9a7 100644
--- a/userspace/units/gr/setup/nvgpu-gr-setup.c
+++ b/userspace/units/gr/setup/nvgpu-gr-setup.c
@@ -288,7 +288,7 @@ int test_gr_setup_preemption_mode_errors(struct unit_module *m,
 	for (i = 0; i < arry_cnt; i++) {
 		err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch,
 				preemp_mode_types[i].graphics_mode,
-				preemp_mode_types[i].compute_mode);
+				preemp_mode_types[i].compute_mode, 0);
 		if (err != preemp_mode_types[i].result) {
 			unit_return_fail(m, "Fail Preemp_mode Error Test-1\n");
 		}
@@ -299,7 +299,7 @@ int test_gr_setup_preemption_mode_errors(struct unit_module *m,
 			NVGPU_PREEMPTION_MODE_COMPUTE_WFI;
 	g->ops.fifo.preempt_tsg = stub_gr_fifo_preempt_tsg;
 	err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch, 0,
-				NVGPU_PREEMPTION_MODE_COMPUTE_CTA);
+				NVGPU_PREEMPTION_MODE_COMPUTE_CTA, 0);
 	if (err == 0) {
 		unit_return_fail(m, "Fail Preemp_mode Error Test-2\n");
 	}
@@ -308,7 +308,7 @@ int test_gr_setup_preemption_mode_errors(struct unit_module *m,
 	tsgid = gr_setup_ch->tsgid;
 	/* Unset the tsgid */
 	gr_setup_ch->tsgid = NVGPU_INVALID_TSG_ID;
-	err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch, 0, 0);
+	err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch, 0, 0, 0);
 	if (err == 0) {
 		unit_return_fail(m, "Fail Preemp_mode Error Test-2\n");
 	}
@@ -316,14 +316,14 @@ int test_gr_setup_preemption_mode_errors(struct unit_module *m,
 	gr_setup_ch->tsgid = tsgid;
 	/* Unset the valid Class*/
 	gr_setup_ch->obj_class = 0;
-	err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch, 0, 0);
+	err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch, 0, 0, 0);
 	if (err == 0) {
 		unit_return_fail(m, "Fail Preemp_mode Error Test-2\n");
 	}
 
 	/* Set invalid Class*/
 	gr_setup_ch->obj_class = 0x1234;
-	err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch, 0, 0);
+	err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch, 0, 0, 0);
 	if (err == 0) {
 		unit_return_fail(m, "Fail Preemp_mode Error Test-2\n");
 	}
@@ -606,7 +606,7 @@ int test_gr_setup_set_preemption_mode(struct unit_module *m,
 						&compute_mode);
 
 	err = g->ops.gr.setup.set_preemption_mode(gr_setup_ch, 0,
-		(compute_mode & NVGPU_PREEMPTION_MODE_COMPUTE_CTA));
+		(compute_mode & NVGPU_PREEMPTION_MODE_COMPUTE_CTA), 0);
 	if (err != 0) {
 		unit_return_fail(m, "setup preemption_mode failed\n");
 	}