diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index 5e0ef6974..1c251b695 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -69,9 +69,6 @@
 #include <nvgpu/hw/gk20a/hw_gr_gk20a.h>
 #include <nvgpu/hw/gk20a/hw_ram_gk20a.h>
 
-#define CTXSW_MEM_SCRUBBING_TIMEOUT_MAX 1000U
-#define CTXSW_MEM_SCRUBBING_TIMEOUT_DEFAULT 10U
-
 static struct channel_gk20a *gk20a_gr_get_channel_from_ctx(
 	struct gk20a *g, u32 curr_ctx, u32 *curr_tsgid);
 
@@ -1478,39 +1475,6 @@ static void gk20a_init_gr_prepare(struct gk20a *g)
 	g->ops.gr.init.fifo_access(g, true);
 }
 
-static int gr_gk20a_wait_mem_scrubbing(struct gk20a *g)
-{
-	struct nvgpu_timeout timeout;
-	bool fecs_scrubbing;
-	bool gpccs_scrubbing;
-
-	nvgpu_log_fn(g, " ");
-
-	nvgpu_timeout_init(g, &timeout,
-			   CTXSW_MEM_SCRUBBING_TIMEOUT_MAX /
-				CTXSW_MEM_SCRUBBING_TIMEOUT_DEFAULT,
-			   NVGPU_TIMER_RETRY_TIMER);
-	do {
-		fecs_scrubbing = (gk20a_readl(g, gr_fecs_dmactl_r()) &
-			(gr_fecs_dmactl_imem_scrubbing_m() |
-			 gr_fecs_dmactl_dmem_scrubbing_m())) != 0U;
-
-		gpccs_scrubbing = (gk20a_readl(g, gr_gpccs_dmactl_r()) &
-			(gr_gpccs_dmactl_imem_scrubbing_m() |
-			 gr_gpccs_dmactl_imem_scrubbing_m())) != 0U;
-
-		if (!fecs_scrubbing && !gpccs_scrubbing) {
-			nvgpu_log_fn(g, "done");
-			return 0;
-		}
-
-		nvgpu_udelay(CTXSW_MEM_SCRUBBING_TIMEOUT_DEFAULT);
-	} while (nvgpu_timeout_expired(&timeout) == 0);
-
-	nvgpu_err(g, "Falcon mem scrubbing timeout");
-	return -ETIMEDOUT;
-}
-
 static int gr_gk20a_init_ctxsw(struct gk20a *g)
 {
 	int err = 0;
@@ -1552,7 +1516,7 @@ static int gk20a_init_gr_reset_enable_hw(struct gk20a *g)
 			sw_non_ctx_load->l[i].value);
 	}
 
-	err = gr_gk20a_wait_mem_scrubbing(g);
+	err = g->ops.gr.falcon.wait_mem_scrubbing(g);
 	if (err != 0) {
 		goto out;
 	}
diff --git a/drivers/gpu/nvgpu/gm20b/hal_gm20b.c b/drivers/gpu/nvgpu/gm20b/hal_gm20b.c
index 94183ee3b..42e24126c 100644
--- a/drivers/gpu/nvgpu/gm20b/hal_gm20b.c
+++ b/drivers/gpu/nvgpu/gm20b/hal_gm20b.c
@@ -508,6 +508,8 @@ static const struct gpu_ops gm20b_ops = {
 				gm20b_gr_falcon_load_ctxsw_ucode_boot,
 			.load_ctxsw_ucode =
 					nvgpu_gr_falcon_load_ctxsw_ucode,
+			.wait_mem_scrubbing =
+					gm20b_gr_falcon_wait_mem_scrubbing,
 		},
 	},
 	.fb = {
diff --git a/drivers/gpu/nvgpu/gp10b/hal_gp10b.c b/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
index 30fc4588c..c0cb6a1d0 100644
--- a/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
@@ -593,6 +593,8 @@ static const struct gpu_ops gp10b_ops = {
 				gm20b_gr_falcon_load_ctxsw_ucode_boot,
 			.load_ctxsw_ucode =
 					nvgpu_gr_falcon_load_ctxsw_ucode,
+			.wait_mem_scrubbing =
+					gm20b_gr_falcon_wait_mem_scrubbing,
 		},
 	},
 	.fb = {
diff --git a/drivers/gpu/nvgpu/gv100/hal_gv100.c b/drivers/gpu/nvgpu/gv100/hal_gv100.c
index a15ca5f10..fe4411fa3 100644
--- a/drivers/gpu/nvgpu/gv100/hal_gv100.c
+++ b/drivers/gpu/nvgpu/gv100/hal_gv100.c
@@ -738,6 +738,8 @@ static const struct gpu_ops gv100_ops = {
 				gm20b_gr_falcon_load_ctxsw_ucode_boot,
 			.load_ctxsw_ucode =
 					nvgpu_gr_falcon_load_secure_ctxsw_ucode,
+			.wait_mem_scrubbing =
+					gm20b_gr_falcon_wait_mem_scrubbing,
 		},
 	},
 	.fb = {
diff --git a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
index 00dd81314..637b8337c 100644
--- a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
@@ -697,6 +697,8 @@ static const struct gpu_ops gv11b_ops = {
 				gm20b_gr_falcon_load_ctxsw_ucode_boot,
 			.load_ctxsw_ucode =
 					nvgpu_gr_falcon_load_ctxsw_ucode,
+			.wait_mem_scrubbing =
+					gm20b_gr_falcon_wait_mem_scrubbing,
 		},
 	},
 	.fb = {
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 3c0a4e2ba..7265d57c3 100644
--- a/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.c
+++ b/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.c
@@ -31,6 +31,8 @@
 
 #define FECS_ARB_CMD_TIMEOUT_MAX_US 40U
 #define FECS_ARB_CMD_TIMEOUT_DEFAULT_US 2U
+#define CTXSW_MEM_SCRUBBING_TIMEOUT_MAX_US 1000U
+#define CTXSW_MEM_SCRUBBING_TIMEOUT_DEFAULT_US 10U
 
 void gm20b_gr_falcon_load_gpccs_dmem(struct gk20a *g,
 			const u32 *ucode_u32_data, u32 ucode_u32_size)
@@ -425,6 +427,45 @@ void gm20b_gr_falcon_load_ctxsw_ucode_boot(struct gk20a *g, u32 reg_offset,
 	}
 }
 
+int gm20b_gr_falcon_wait_mem_scrubbing(struct gk20a *g)
+{
+	struct nvgpu_timeout timeout;
+	int err;
+	bool fecs_scrubbing;
+	bool gpccs_scrubbing;
+
+	nvgpu_log_fn(g, " ");
+
+	err = nvgpu_timeout_init(g, &timeout,
+			   CTXSW_MEM_SCRUBBING_TIMEOUT_MAX_US /
+				CTXSW_MEM_SCRUBBING_TIMEOUT_DEFAULT_US,
+			   NVGPU_TIMER_RETRY_TIMER);
+	if (err != 0) {
+		nvgpu_err(g, "ctxsw mem scrub timeout_init failed: %d", err);
+		return err;
+	}
+
+	do {
+		fecs_scrubbing = (nvgpu_readl(g, gr_fecs_dmactl_r()) &
+			(gr_fecs_dmactl_imem_scrubbing_m() |
+			 gr_fecs_dmactl_dmem_scrubbing_m())) != 0U;
+
+		gpccs_scrubbing = (nvgpu_readl(g, gr_gpccs_dmactl_r()) &
+			(gr_gpccs_dmactl_imem_scrubbing_m() |
+			 gr_gpccs_dmactl_imem_scrubbing_m())) != 0U;
+
+		if (!fecs_scrubbing && !gpccs_scrubbing) {
+			nvgpu_log_fn(g, "done");
+			return 0;
+		}
+
+		nvgpu_udelay(CTXSW_MEM_SCRUBBING_TIMEOUT_DEFAULT_US);
+	} while (nvgpu_timeout_expired(&timeout) == 0);
+
+	nvgpu_err(g, "Falcon mem scrubbing timeout");
+	return -ETIMEDOUT;
+}
+
 u32 gm20b_gr_falcon_fecs_base_addr(void)
 {
 	return gr_fecs_irqsset_r();
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 e69a82dd1..dc63f8a03 100644
--- a/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.h
+++ b/drivers/gpu/nvgpu/hal/gr/falcon/gr_falcon_gm20b.h
@@ -53,6 +53,7 @@ void gm20b_gr_falcon_load_ctxsw_ucode_header(struct gk20a *g,
 void gm20b_gr_falcon_load_ctxsw_ucode_boot(struct gk20a *g,
 	u32 reg_offset, u32 boot_entry, u32 addr_load32, u32 blocks,
 	u32 dst);
+int gm20b_gr_falcon_wait_mem_scrubbing(struct gk20a *g);
 
 void gm20b_gr_falcon_set_current_ctx_invalid(struct gk20a *g);
 
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
index b98c4dae7..b2334c766 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
@@ -585,6 +585,7 @@ struct gpu_ops {
 				u32 reg_offset, u32 boot_entry,
 				u32 addr_load32, u32 blocks, u32 dst);
 			int (*load_ctxsw_ucode)(struct gk20a *g);
+			int (*wait_mem_scrubbing)(struct gk20a *g);
 		} falcon;
 
 #ifdef CONFIG_GK20A_CTXSW_TRACE
diff --git a/drivers/gpu/nvgpu/tu104/hal_tu104.c b/drivers/gpu/nvgpu/tu104/hal_tu104.c
index 1f38fbded..ea5f6ed29 100644
--- a/drivers/gpu/nvgpu/tu104/hal_tu104.c
+++ b/drivers/gpu/nvgpu/tu104/hal_tu104.c
@@ -771,6 +771,8 @@ static const struct gpu_ops tu104_ops = {
 				gm20b_gr_falcon_load_ctxsw_ucode_boot,
 			.load_ctxsw_ucode =
 				nvgpu_gr_falcon_load_secure_ctxsw_ucode,
+			.wait_mem_scrubbing =
+					gm20b_gr_falcon_wait_mem_scrubbing,
 		},
 	},
 	.fb = {