diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index 001a9d277..284811e77 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -806,8 +806,7 @@ restore_fe_go_idle:
 
 	gr->ctx_vars.golden_image_initialized = true;
 
-	gk20a_writel(g, gr_fecs_current_ctx_r(),
-		gr_fecs_current_ctx_valid_false_f());
+	g->ops.gr.falcon.set_current_ctx_invalid(g);
 
 clean_up:
 	if (err != 0) {
diff --git a/drivers/gpu/nvgpu/gm20b/hal_gm20b.c b/drivers/gpu/nvgpu/gm20b/hal_gm20b.c
index 3e822953e..94183ee3b 100644
--- a/drivers/gpu/nvgpu/gm20b/hal_gm20b.c
+++ b/drivers/gpu/nvgpu/gm20b/hal_gm20b.c
@@ -484,6 +484,8 @@ static const struct gpu_ops gm20b_ops = {
 		.falcon = {
 			.fecs_base_addr = gm20b_gr_falcon_fecs_base_addr,
 			.gpccs_base_addr = gm20b_gr_falcon_gpccs_base_addr,
+			.set_current_ctx_invalid =
+				gm20b_gr_falcon_set_current_ctx_invalid,
 			.dump_stats = gm20b_gr_falcon_fecs_dump_stats,
 			.fecs_ctxsw_mailbox_size =
 				gm20b_gr_falcon_get_fecs_ctxsw_mailbox_size,
diff --git a/drivers/gpu/nvgpu/gp10b/hal_gp10b.c b/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
index 2a1f75eaf..30fc4588c 100644
--- a/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
@@ -569,6 +569,8 @@ static const struct gpu_ops gp10b_ops = {
 		.falcon = {
 			.fecs_base_addr = gm20b_gr_falcon_fecs_base_addr,
 			.gpccs_base_addr = gm20b_gr_falcon_gpccs_base_addr,
+			.set_current_ctx_invalid =
+				gm20b_gr_falcon_set_current_ctx_invalid,
 			.dump_stats = gm20b_gr_falcon_fecs_dump_stats,
 			.fecs_ctxsw_mailbox_size =
 				gm20b_gr_falcon_get_fecs_ctxsw_mailbox_size,
diff --git a/drivers/gpu/nvgpu/gv100/hal_gv100.c b/drivers/gpu/nvgpu/gv100/hal_gv100.c
index 22b34527b..b36115dd1 100644
--- a/drivers/gpu/nvgpu/gv100/hal_gv100.c
+++ b/drivers/gpu/nvgpu/gv100/hal_gv100.c
@@ -713,6 +713,8 @@ static const struct gpu_ops gv100_ops = {
 		.falcon = {
 			.fecs_base_addr = gm20b_gr_falcon_fecs_base_addr,
 			.gpccs_base_addr = gm20b_gr_falcon_gpccs_base_addr,
+			.set_current_ctx_invalid =
+				gm20b_gr_falcon_set_current_ctx_invalid,
 			.dump_stats = gm20b_gr_falcon_fecs_dump_stats,
 			.fecs_ctxsw_mailbox_size =
 				gm20b_gr_falcon_get_fecs_ctxsw_mailbox_size,
diff --git a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
index b4644d084..b937f70b2 100644
--- a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
@@ -672,6 +672,8 @@ static const struct gpu_ops gv11b_ops = {
 		.falcon = {
 			.fecs_base_addr = gm20b_gr_falcon_fecs_base_addr,
 			.gpccs_base_addr = gm20b_gr_falcon_gpccs_base_addr,
+			.set_current_ctx_invalid =
+				gm20b_gr_falcon_set_current_ctx_invalid,
 			.dump_stats = gm20b_gr_falcon_fecs_dump_stats,
 			.fecs_ctxsw_mailbox_size =
 				gm20b_gr_falcon_get_fecs_ctxsw_mailbox_size,
diff --git a/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.c b/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.c
index 86f51a0d6..3c0a4e2ba 100644
--- a/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.c
+++ b/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.c
@@ -456,3 +456,10 @@ u32 gm20b_gr_falcon_get_fecs_ctxsw_mailbox_size(void)
 {
 	return gr_fecs_ctxsw_mailbox__size_1_v();
 }
+
+void gm20b_gr_falcon_set_current_ctx_invalid(struct gk20a *g)
+{
+	nvgpu_writel(g, gr_fecs_current_ctx_r(),
+		gr_fecs_current_ctx_valid_false_f());
+}
+
diff --git a/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.h b/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.h
index 5a991275e..e69a82dd1 100644
--- a/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.h
+++ b/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.h
@@ -54,4 +54,6 @@ void gm20b_gr_falcon_load_ctxsw_ucode_boot(struct gk20a *g,
 	u32 reg_offset, u32 boot_entry, u32 addr_load32, u32 blocks,
 	u32 dst);
 
+void gm20b_gr_falcon_set_current_ctx_invalid(struct gk20a *g);
+
 #endif /* NVGPU_GR_FALCON_GM20B_H */
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
index 721a07228..82b244003 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
@@ -559,6 +559,7 @@ struct gpu_ops {
 		struct {
 			u32 (*fecs_base_addr)(void);
 			u32 (*gpccs_base_addr)(void);
+			void (*set_current_ctx_invalid)(struct gk20a *g);
 			void (*dump_stats)(struct gk20a *g);
 			u32 (*fecs_ctxsw_mailbox_size)(void);
 			u32 (*get_fecs_ctx_state_store_major_rev_id)(
diff --git a/drivers/gpu/nvgpu/tu104/hal_tu104.c b/drivers/gpu/nvgpu/tu104/hal_tu104.c
index 9ed6a4686..c26c6bea7 100644
--- a/drivers/gpu/nvgpu/tu104/hal_tu104.c
+++ b/drivers/gpu/nvgpu/tu104/hal_tu104.c
@@ -746,6 +746,8 @@ static const struct gpu_ops tu104_ops = {
 		.falcon = {
 			.fecs_base_addr = gm20b_gr_falcon_fecs_base_addr,
 			.gpccs_base_addr = gm20b_gr_falcon_gpccs_base_addr,
+			.set_current_ctx_invalid =
+				gm20b_gr_falcon_set_current_ctx_invalid,
 			.dump_stats = gm20b_gr_falcon_fecs_dump_stats,
 			.fecs_ctxsw_mailbox_size =
 				gm20b_gr_falcon_get_fecs_ctxsw_mailbox_size,