diff --git a/drivers/gpu/nvgpu/common/vgpu/gv11b/vgpu_hal_gv11b.c b/drivers/gpu/nvgpu/common/vgpu/gv11b/vgpu_hal_gv11b.c
index 34422d8a8..b316dd1cb 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gv11b/vgpu_hal_gv11b.c
+++ b/drivers/gpu/nvgpu/common/vgpu/gv11b/vgpu_hal_gv11b.c
@@ -409,6 +409,8 @@ static const struct gpu_ops vgpu_gv11b_ops = {
 			.commit_ctxsw_spill = gv11b_gr_init_commit_ctxsw_spill,
 			.commit_cbes_reserve =
 				gv11b_gr_init_commit_cbes_reserve,
+			.gfxp_wfi_timeout =
+				gv11b_gr_init_commit_gfxp_wfi_timeout,
 		},
 	},
 	.perf = {
diff --git a/drivers/gpu/nvgpu/gv100/hal_gv100.c b/drivers/gpu/nvgpu/gv100/hal_gv100.c
index 18967549d..22b34527b 100644
--- a/drivers/gpu/nvgpu/gv100/hal_gv100.c
+++ b/drivers/gpu/nvgpu/gv100/hal_gv100.c
@@ -698,6 +698,8 @@ static const struct gpu_ops gv100_ops = {
 			.commit_ctxsw_spill = gv11b_gr_init_commit_ctxsw_spill,
 			.commit_cbes_reserve =
 				gv11b_gr_init_commit_cbes_reserve,
+			.gfxp_wfi_timeout =
+				gv11b_gr_init_commit_gfxp_wfi_timeout,
 		},
 		.intr = {
 			.handle_tex_exception = NULL,
diff --git a/drivers/gpu/nvgpu/gv11b/gr_gv11b.c b/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
index 40f7e4bed..67243a740 100644
--- a/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
@@ -1393,11 +1393,8 @@ void gr_gv11b_update_ctxsw_preemption_mode(struct gk20a *g,
 		g->ops.gr.init.commit_ctxsw_spill(g, gr_ctx, addr, size, true);
 
 		g->ops.gr.init.commit_cbes_reserve(g, gr_ctx, true);
-
-		nvgpu_gr_ctx_patch_write(g, gr_ctx,
-				gr_fe_gfxp_wfi_timeout_r(),
-				g->gr.gfxp_wfi_timeout_count,
-				true);
+		g->ops.gr.init.gfxp_wfi_timeout(g, gr_ctx,
+			g->gr.gfxp_wfi_timeout_count, true);
 
 		if (g->ops.gr.init.commit_gfxp_rtv_cb != NULL) {
 			g->ops.gr.init.commit_gfxp_rtv_cb(g, gr_ctx, true);
diff --git a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
index 2296ec539..ad3475f99 100644
--- a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
@@ -657,6 +657,8 @@ static const struct gpu_ops gv11b_ops = {
 			.commit_ctxsw_spill = gv11b_gr_init_commit_ctxsw_spill,
 			.commit_cbes_reserve =
 				gv11b_gr_init_commit_cbes_reserve,
+			.gfxp_wfi_timeout =
+				gv11b_gr_init_commit_gfxp_wfi_timeout,
 		},
 		.intr = {
 			.handle_tex_exception = NULL,
diff --git a/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.c b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.c
index 485ba7228..f0863b58c 100644
--- a/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.c
+++ b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.c
@@ -969,3 +969,10 @@ void gv11b_gr_init_commit_cbes_reserve(struct gk20a *g,
 		patch);
 }
 
+void gv11b_gr_init_commit_gfxp_wfi_timeout(struct gk20a *g,
+	struct nvgpu_gr_ctx *gr_ctx, u32 timeout, bool patch)
+{
+	nvgpu_gr_ctx_patch_write(g, gr_ctx, gr_fe_gfxp_wfi_timeout_r(),
+		timeout, patch);
+}
+
diff --git a/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.h b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.h
index a3a6e4c3c..11cc8ba9b 100644
--- a/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.h
+++ b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.h
@@ -73,5 +73,7 @@ void gv11b_gr_init_commit_ctxsw_spill(struct gk20a *g,
 	struct nvgpu_gr_ctx *gr_ctx, u64 addr, u32 size, bool patch);
 void gv11b_gr_init_commit_cbes_reserve(struct gk20a *g,
 	struct nvgpu_gr_ctx *gr_ctx, bool patch);
+void gv11b_gr_init_commit_gfxp_wfi_timeout(struct gk20a *g,
+	struct nvgpu_gr_ctx *gr_ctx, u32 timeout, bool patch);
 
 #endif /* NVGPU_GR_INIT_GV11B_H */
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
index 7dfedad65..721a07228 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
@@ -762,6 +762,9 @@ struct gpu_ops {
 				bool patch);
 			void (*commit_cbes_reserve)(struct gk20a *g,
 				struct nvgpu_gr_ctx *gr_ctx, bool patch);
+			void (*gfxp_wfi_timeout)(struct gk20a *g,
+				struct nvgpu_gr_ctx *gr_ctx, u32 timeout,
+				bool patch);
 		} init;
 
 		struct {
diff --git a/drivers/gpu/nvgpu/tu104/hal_tu104.c b/drivers/gpu/nvgpu/tu104/hal_tu104.c
index dfb009d3c..07d7937e4 100644
--- a/drivers/gpu/nvgpu/tu104/hal_tu104.c
+++ b/drivers/gpu/nvgpu/tu104/hal_tu104.c
@@ -731,6 +731,8 @@ static const struct gpu_ops tu104_ops = {
 			.commit_ctxsw_spill = gv11b_gr_init_commit_ctxsw_spill,
 			.commit_cbes_reserve =
 				gv11b_gr_init_commit_cbes_reserve,
+			.gfxp_wfi_timeout =
+				gv11b_gr_init_commit_gfxp_wfi_timeout,
 		},
 		.intr = {
 			.handle_tex_exception = NULL,