diff --git a/drivers/gpu/nvgpu/Makefile b/drivers/gpu/nvgpu/Makefile
index bb7de80dd..82aa25d36 100644
--- a/drivers/gpu/nvgpu/Makefile
+++ b/drivers/gpu/nvgpu/Makefile
@@ -167,6 +167,8 @@ nvgpu-y += \
 	hal/bus/bus_gv100.o \
 	hal/bus/bus_tu104.o \
 	hal/gr/init/gr_init_gm20b.o \
+	hal/gr/init/gr_init_gp10b.o \
+	hal/gr/init/gr_init_gv11b.o \
 	hal/gr/hwpm_map/hwpm_map_gv100.o \
 	hal/gr/zbc/zbc_gm20b.o \
 	hal/gr/zbc/zbc_gp10b.o \
diff --git a/drivers/gpu/nvgpu/Makefile.sources b/drivers/gpu/nvgpu/Makefile.sources
index 422228a0d..07bbd5f81 100644
--- a/drivers/gpu/nvgpu/Makefile.sources
+++ b/drivers/gpu/nvgpu/Makefile.sources
@@ -332,6 +332,8 @@ srcs += common/sim.c \
 	hal/bus/bus_gv100.c	\
 	hal/bus/bus_tu104.c	\
 	hal/gr/init/gr_init_gm20b.c \
+	hal/gr/init/gr_init_gp10b.c \
+	hal/gr/init/gr_init_gv11b.c \
 	hal/gr/hwpm_map/hwpm_map_gv100.c \
 	hal/gr/zbc/zbc_gm20b.c \
 	hal/gr/zbc/zbc_gp10b.c \
diff --git a/drivers/gpu/nvgpu/common/vgpu/gp10b/vgpu_hal_gp10b.c b/drivers/gpu/nvgpu/common/vgpu/gp10b/vgpu_hal_gp10b.c
index d4b7b7ad5..c0d45e19e 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gp10b/vgpu_hal_gp10b.c
+++ b/drivers/gpu/nvgpu/common/vgpu/gp10b/vgpu_hal_gp10b.c
@@ -200,7 +200,6 @@ static const struct gpu_ops vgpu_gp10b_ops = {
 		.set_preemption_mode = vgpu_gr_set_preemption_mode,
 		.pre_process_sm_exception = NULL,
 		.set_preemption_buffer_va = gr_gp10b_set_preemption_buffer_va,
-		.init_preemption_state = NULL,
 		.set_bes_crop_debug3 = NULL,
 		.set_bes_crop_debug4 = NULL,
 		.set_ctxsw_preemption_mode =
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 01f927f23..7a840b971 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gv11b/vgpu_hal_gv11b.c
+++ b/drivers/gpu/nvgpu/common/vgpu/gv11b/vgpu_hal_gv11b.c
@@ -218,7 +218,6 @@ static const struct gpu_ops vgpu_gv11b_ops = {
 		.set_preemption_mode = vgpu_gr_set_preemption_mode,
 		.pre_process_sm_exception = NULL,
 		.set_preemption_buffer_va = gr_gv11b_set_preemption_buffer_va,
-		.init_preemption_state = NULL,
 		.set_bes_crop_debug3 = NULL,
 		.set_bes_crop_debug4 = NULL,
 		.set_ctxsw_preemption_mode = vgpu_gr_set_ctxsw_preemption_mode,
diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index afac60db3..7cf5865ec 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -1120,8 +1120,13 @@ int gr_gk20a_init_golden_ctx_image(struct gk20a *g,
 			     sw_ctx_load->l[i].value);
 	}
 
-	if (g->ops.gr.init_preemption_state != NULL) {
-		g->ops.gr.init_preemption_state(g);
+	if (g->ops.gr.init.preemption_state != NULL) {
+		err = g->ops.gr.init.preemption_state(g,
+			gr->gfxp_wfi_timeout_count,
+			gr->gfxp_wfi_timeout_unit_usec);
+		if (err != 0) {
+			goto clean_up;
+		}
 	}
 
 	nvgpu_cg_blcg_gr_load_enable(g);
@@ -2662,8 +2667,9 @@ static int gk20a_init_gr_setup_hw(struct gk20a *g)
 		g->ops.gr.disable_rd_coalesce(g);
 	}
 
-	if (g->ops.gr.init_preemption_state != NULL) {
-		err = g->ops.gr.init_preemption_state(g);
+	if (g->ops.gr.init.preemption_state != NULL) {
+		err = g->ops.gr.init.preemption_state(g, gr->gfxp_wfi_timeout_count,
+			gr->gfxp_wfi_timeout_unit_usec);
 		if (err != 0) {
 			goto out;
 		}
diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.h b/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
index e9ab6bdee..6b1e19adf 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
@@ -224,8 +224,9 @@ struct gr_gk20a {
 	u32 attrib_cb_gfxp_size;
 	u32 alpha_cb_default_size;
 	u32 alpha_cb_size;
+
 	u32 gfxp_wfi_timeout_count;
-	u32 gfxp_wfi_timeout_unit;
+	bool gfxp_wfi_timeout_unit_usec;
 
 	/*
 	 * The deductible memory size for max_comptag_mem (in MBytes)
diff --git a/drivers/gpu/nvgpu/gp10b/gr_gp10b.c b/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
index 1f9210ac5..712059a6d 100644
--- a/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
@@ -2148,23 +2148,6 @@ int gr_gp10b_get_preemption_mode_flags(struct gk20a *g,
 	return 0;
 }
 
-int gr_gp10b_init_preemption_state(struct gk20a *g)
-{
-	u32 debug_2;
-	struct gr_gk20a *gr = &g->gr;
-	u32 sysclk_cycles = gr->gfxp_wfi_timeout_count;
-	gk20a_writel(g, gr_fe_gfxp_wfi_timeout_r(),
-			gr_fe_gfxp_wfi_timeout_count_f(sysclk_cycles));
-
-	debug_2 = gk20a_readl(g, gr_debug_2_r());
-	debug_2 = set_field(debug_2,
-			gr_debug_2_gfxp_wfi_always_injects_wfi_m(),
-			gr_debug_2_gfxp_wfi_always_injects_wfi_enabled_f());
-	gk20a_writel(g, gr_debug_2_r(), debug_2);
-
-	return 0;
-}
-
 void gr_gp10b_set_preemption_buffer_va(struct gk20a *g,
 			struct nvgpu_mem *mem, u64 gpu_va)
 {
diff --git a/drivers/gpu/nvgpu/gp10b/hal_gp10b.c b/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
index 71814d3ad..217db9ca3 100644
--- a/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/hal_gp10b.c
@@ -53,6 +53,7 @@
 #include "hal/fifo/engines_gm20b.h"
 #include "hal/gr/zbc/zbc_gp10b.h"
 #include "hal/gr/init/gr_init_gm20b.h"
+#include "hal/gr/init/gr_init_gp10b.h"
 #include "hal/falcon/falcon_gk20a.h"
 
 #include "common/ptimer/ptimer_gk20a.h"
@@ -344,7 +345,6 @@ static const struct gpu_ops gp10b_ops = {
 		.set_preemption_mode = gr_gp10b_set_preemption_mode,
 		.pre_process_sm_exception = gr_gp10b_pre_process_sm_exception,
 		.set_preemption_buffer_va = gr_gp10b_set_preemption_buffer_va,
-		.init_preemption_state = gr_gp10b_init_preemption_state,
 		.set_bes_crop_debug3 = gr_gp10b_set_bes_crop_debug3,
 		.set_ctxsw_preemption_mode = gr_gp10b_set_ctxsw_preemption_mode,
 		.init_ecc = gp10b_ecc_init,
@@ -495,15 +495,16 @@ static const struct gpu_ops gp10b_ops = {
 				gp10b_gr_zbc_get_gpcs_swdx_dss_zbc_z_format_reg,
 		},
 		.init = {
+			.pd_tpc_per_gpc = gm20b_gr_init_pd_tpc_per_gpc,
+			.pd_skip_table_gpc = gm20b_gr_init_pd_skip_table_gpc,
+			.cwd_gpcs_tpcs_num = gm20b_gr_init_cwd_gpcs_tpcs_num,
+			.wait_idle = gm20b_gr_init_wait_idle,
+			.wait_fe_idle = gm20b_gr_init_wait_fe_idle,
 			.fe_pwr_mode_force_on =
 				gm20b_gr_init_fe_pwr_mode_force_on,
 			.override_context_reset =
 				gm20b_gr_init_override_context_reset,
-			.wait_idle = gm20b_gr_init_wait_idle,
-			.wait_fe_idle = gm20b_gr_init_wait_fe_idle,
-			.pd_tpc_per_gpc = gm20b_gr_init_pd_tpc_per_gpc,
-			.pd_skip_table_gpc = gm20b_gr_init_pd_skip_table_gpc,
-			.cwd_gpcs_tpcs_num = gm20b_gr_init_cwd_gpcs_tpcs_num,
+			.preemption_state = gp10b_gr_init_preemption_state,
 		},
 	},
 	.fb = {
diff --git a/drivers/gpu/nvgpu/gv100/hal_gv100.c b/drivers/gpu/nvgpu/gv100/hal_gv100.c
index ed73c37aa..b6d9f03c9 100644
--- a/drivers/gpu/nvgpu/gv100/hal_gv100.c
+++ b/drivers/gpu/nvgpu/gv100/hal_gv100.c
@@ -450,7 +450,6 @@ static const struct gpu_ops gv100_ops = {
 		.set_preemption_mode = gr_gp10b_set_preemption_mode,
 		.pre_process_sm_exception = gr_gv11b_pre_process_sm_exception,
 		.set_preemption_buffer_va = gr_gv11b_set_preemption_buffer_va,
-		.init_preemption_state = NULL,
 		.set_bes_crop_debug3 = gr_gp10b_set_bes_crop_debug3,
 		.set_bes_crop_debug4 = gr_gp10b_set_bes_crop_debug4,
 		.set_ctxsw_preemption_mode = gr_gp10b_set_ctxsw_preemption_mode,
diff --git a/drivers/gpu/nvgpu/gv11b/gr_gv11b.c b/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
index 47ae65e18..1a415fbd7 100644
--- a/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
@@ -4170,38 +4170,16 @@ void gr_gv11b_init_gpc_mmu(struct gk20a *g)
 			g->ops.fb.mmu_debug_rd(g));
 }
 
-int gr_gv11b_init_preemption_state(struct gk20a *g)
-{
-	u32 debug_2;
-	struct gr_gk20a *gr = &g->gr;
-	u32 unit;
-
-	nvgpu_log_fn(g, " ");
-
-	if (gr->gfxp_wfi_timeout_unit == GFXP_WFI_TIMEOUT_UNIT_USEC) {
-		unit = gr_debug_2_gfxp_wfi_timeout_unit_usec_f();
-	} else {
-		unit = gr_debug_2_gfxp_wfi_timeout_unit_sysclk_f();
-	}
-
-	debug_2 = gk20a_readl(g, gr_debug_2_r());
-	debug_2 = set_field(debug_2,
-		gr_debug_2_gfxp_wfi_timeout_unit_m(),
-		unit);
-	gk20a_writel(g, gr_debug_2_r(), debug_2);
-
-	return 0;
-}
 void gr_gv11b_init_gfxp_wfi_timeout_count(struct gk20a *g)
 {
 	struct gr_gk20a *gr = &g->gr;
-	gr->gfxp_wfi_timeout_unit = GFXP_WFI_TIMEOUT_UNIT_USEC;
+	gr->gfxp_wfi_timeout_unit_usec = true;
 	gr->gfxp_wfi_timeout_count = GFXP_WFI_TIMEOUT_COUNT_IN_USEC_DEFAULT;
 }
 
 unsigned long gr_gv11b_get_max_gfxp_wfi_timeout_count(struct gk20a *g)
 {
-	if (g->gr.gfxp_wfi_timeout_unit == GFXP_WFI_TIMEOUT_UNIT_USEC) {
+	if (g->gr.gfxp_wfi_timeout_unit_usec) {
 		/* 100 msec in usec count */
 		return (100UL * 1000UL);
 	} else {
diff --git a/drivers/gpu/nvgpu/gv11b/gr_gv11b.h b/drivers/gpu/nvgpu/gv11b/gr_gv11b.h
index 1a5415b1b..1b31a97db 100644
--- a/drivers/gpu/nvgpu/gv11b/gr_gv11b.h
+++ b/drivers/gpu/nvgpu/gv11b/gr_gv11b.h
@@ -30,9 +30,6 @@
 
 #define PRI_BROADCAST_FLAGS_SMPC  BIT32(17)
 
-#define GFXP_WFI_TIMEOUT_UNIT_SYSCLK      0U
-#define GFXP_WFI_TIMEOUT_UNIT_USEC        1U
-
 struct gk20a;
 struct gr_gk20a;
 struct nvgpu_gr_ctx;
diff --git a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
index 261df8c25..0c7eee4ec 100644
--- a/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/hal_gv11b.c
@@ -43,6 +43,7 @@
 #include "hal/gr/zbc/zbc_gp10b.h"
 #include "hal/gr/zbc/zbc_gv11b.h"
 #include "hal/gr/init/gr_init_gm20b.h"
+#include "hal/gr/init/gr_init_gv11b.h"
 #include "hal/gr/hwpm_map/hwpm_map_gv100.h"
 #include "hal/falcon/falcon_gk20a.h"
 
@@ -403,7 +404,6 @@ static const struct gpu_ops gv11b_ops = {
 		.set_preemption_mode = gr_gp10b_set_preemption_mode,
 		.pre_process_sm_exception = gr_gv11b_pre_process_sm_exception,
 		.set_preemption_buffer_va = gr_gv11b_set_preemption_buffer_va,
-		.init_preemption_state = gr_gv11b_init_preemption_state,
 		.set_bes_crop_debug3 = gr_gp10b_set_bes_crop_debug3,
 		.set_bes_crop_debug4 = gr_gp10b_set_bes_crop_debug4,
 		.init_ecc = gv11b_ecc_init,
@@ -588,15 +588,16 @@ static const struct gpu_ops gv11b_ops = {
 				gv100_gr_hwpm_map_align_regs_perf_pma,
 		},
 		.init = {
+			.pd_tpc_per_gpc = gm20b_gr_init_pd_tpc_per_gpc,
+			.pd_skip_table_gpc = gm20b_gr_init_pd_skip_table_gpc,
+			.cwd_gpcs_tpcs_num = gm20b_gr_init_cwd_gpcs_tpcs_num,
+			.wait_idle = gm20b_gr_init_wait_idle,
+			.wait_fe_idle = gm20b_gr_init_wait_fe_idle,
 			.fe_pwr_mode_force_on =
 				gm20b_gr_init_fe_pwr_mode_force_on,
 			.override_context_reset =
 				gm20b_gr_init_override_context_reset,
-			.wait_idle = gm20b_gr_init_wait_idle,
-			.wait_fe_idle = gm20b_gr_init_wait_fe_idle,
-			.pd_tpc_per_gpc = gm20b_gr_init_pd_tpc_per_gpc,
-			.pd_skip_table_gpc = gm20b_gr_init_pd_skip_table_gpc,
-			.cwd_gpcs_tpcs_num = gm20b_gr_init_cwd_gpcs_tpcs_num,
+			.preemption_state = gv11b_gr_init_preemption_state,
 		},
 	},
 	.fb = {
diff --git a/drivers/gpu/nvgpu/hal/gr/init/gr_init_gp10b.c b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gp10b.c
new file mode 100644
index 000000000..6cee4c7ee
--- /dev/null
+++ b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gp10b.c
@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include <nvgpu/gk20a.h>
+#include <nvgpu/io.h>
+
+#include "gr_init_gp10b.h"
+
+#include <nvgpu/hw/gp10b/hw_gr_gp10b.h>
+
+int gp10b_gr_init_preemption_state(struct gk20a *g, u32 gfxp_wfi_timeout_count,
+	bool gfxp_wfi_timeout_unit_usec)
+{
+	u32 debug_2;
+
+	nvgpu_writel(g, gr_fe_gfxp_wfi_timeout_r(),
+			gr_fe_gfxp_wfi_timeout_count_f(gfxp_wfi_timeout_count));
+
+	debug_2 = nvgpu_readl(g, gr_debug_2_r());
+	debug_2 = set_field(debug_2,
+			gr_debug_2_gfxp_wfi_always_injects_wfi_m(),
+			gr_debug_2_gfxp_wfi_always_injects_wfi_enabled_f());
+	nvgpu_writel(g, gr_debug_2_r(), debug_2);
+
+	return 0;
+}
+
diff --git a/drivers/gpu/nvgpu/hal/gr/init/gr_init_gp10b.h b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gp10b.h
new file mode 100644
index 000000000..ec11cd90b
--- /dev/null
+++ b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gp10b.h
@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef NVGPU_GR_INIT_GP10B_H
+#define NVGPU_GR_INIT_GP10B_H
+
+#include <nvgpu/types.h>
+
+struct gk20a;
+
+int gp10b_gr_init_preemption_state(struct gk20a *g, u32 gfxp_wfi_timeout_count,
+	bool gfxp_wfi_timeout_unit_usec);
+
+#endif /* NVGPU_GR_INIT_GP10B_H */
diff --git a/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.c b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.c
new file mode 100644
index 000000000..d17d0a687
--- /dev/null
+++ b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.c
@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include <nvgpu/gk20a.h>
+#include <nvgpu/io.h>
+
+#include "gr_init_gv11b.h"
+
+#include <nvgpu/hw/gv11b/hw_gr_gv11b.h>
+
+int gv11b_gr_init_preemption_state(struct gk20a *g, u32 gfxp_wfi_timeout_count,
+	bool gfxp_wfi_timeout_unit_usec)
+{
+	u32 debug_2;
+	u32 unit;
+
+	nvgpu_log_fn(g, " ");
+
+	if (gfxp_wfi_timeout_unit_usec) {
+		unit = gr_debug_2_gfxp_wfi_timeout_unit_usec_f();
+	} else {
+		unit = gr_debug_2_gfxp_wfi_timeout_unit_sysclk_f();
+	}
+
+	debug_2 = nvgpu_readl(g, gr_debug_2_r());
+	debug_2 = set_field(debug_2,
+		gr_debug_2_gfxp_wfi_timeout_unit_m(),
+		unit);
+	nvgpu_writel(g, gr_debug_2_r(), debug_2);
+
+	return 0;
+}
+
diff --git a/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.h b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.h
new file mode 100644
index 000000000..afa52010c
--- /dev/null
+++ b/drivers/gpu/nvgpu/hal/gr/init/gr_init_gv11b.h
@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef NVGPU_GR_INIT_GV11B_H
+#define NVGPU_GR_INIT_GV11B_H
+
+#include <nvgpu/types.h>
+
+struct gk20a;
+
+int gv11b_gr_init_preemption_state(struct gk20a *g, u32 gfxp_wfi_timeout_count,
+	bool gfxp_wfi_timeout_unit_usec);
+
+#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 0ba2b693e..f297c80c6 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gk20a.h
@@ -252,7 +252,6 @@ struct gpu_ops {
 	struct {
 		u32 (*get_patch_slots)(struct gk20a *g);
 		int (*init_fs_state)(struct gk20a *g);
-		int (*init_preemption_state)(struct gk20a *g);
 		void (*access_smpc_reg)(struct gk20a *g, u32 quad, u32 offset);
 		void (*bundle_cb_defaults)(struct gk20a *g);
 		void (*cb_size_default)(struct gk20a *g);
@@ -682,6 +681,9 @@ struct gpu_ops {
 			int (*fe_pwr_mode_force_on)(struct gk20a *g,
 				bool force_on);
 			void (*override_context_reset)(struct gk20a *g);
+			int (*preemption_state)(struct gk20a *g,
+				u32 gfxp_wfi_timeout_count,
+				bool gfxp_wfi_timeout_unit_usec);
 		} init;
 
 		u32 (*fecs_falcon_base_addr)(void);
diff --git a/drivers/gpu/nvgpu/os/linux/sysfs.c b/drivers/gpu/nvgpu/os/linux/sysfs.c
index 21261f3b0..afce42a24 100644
--- a/drivers/gpu/nvgpu/os/linux/sysfs.c
+++ b/drivers/gpu/nvgpu/os/linux/sysfs.c
@@ -1002,13 +1002,15 @@ static ssize_t gfxp_wfi_timeout_count_store(struct device *dev,
 
 	gr->gfxp_wfi_timeout_count = val;
 
-	if (g->ops.gr.init_preemption_state && g->power_on) {
+	if (g->ops.gr.init.preemption_state && g->power_on) {
 		err = gk20a_busy(g);
 		if (err)
 			return err;
 
 		err = nvgpu_pg_elpg_protected_call(g,
-			g->ops.gr.init_preemption_state(g));
+			g->ops.gr.init.preemption_state(g,
+				gr->gfxp_wfi_timeout_count,
+				gr->gfxp_wfi_timeout_unit_usec));
 
 		gk20a_idle(g);
 
@@ -1027,18 +1029,20 @@ static ssize_t gfxp_wfi_timeout_unit_store(struct device *dev,
 
 	if (count > 0 && buf[0] == 's')
 		/* sysclk */
-		gr->gfxp_wfi_timeout_unit = GFXP_WFI_TIMEOUT_UNIT_SYSCLK;
+		gr->gfxp_wfi_timeout_unit_usec = false;
 	else
 		/* usec */
-		gr->gfxp_wfi_timeout_unit = GFXP_WFI_TIMEOUT_UNIT_USEC;
+		gr->gfxp_wfi_timeout_unit_usec = true;
 
-	if (g->ops.gr.init_preemption_state && g->power_on) {
+	if (g->ops.gr.init.preemption_state && g->power_on) {
 		err = gk20a_busy(g);
 		if (err)
 			return err;
 
 		err = nvgpu_pg_elpg_protected_call(g,
-			g->ops.gr.init_preemption_state(g));
+			g->ops.gr.init.preemption_state(g,
+				gr->gfxp_wfi_timeout_count,
+				gr->gfxp_wfi_timeout_unit_usec));
 
 		gk20a_idle(g);
 
@@ -1065,7 +1069,7 @@ static ssize_t gfxp_wfi_timeout_unit_read(struct device *dev,
 	struct gk20a *g = get_gk20a(dev);
 	struct gr_gk20a *gr = &g->gr;
 
-	if (gr->gfxp_wfi_timeout_unit == GFXP_WFI_TIMEOUT_UNIT_USEC)
+	if (gr->gfxp_wfi_timeout_unit_usec)
 		return snprintf(buf, PAGE_SIZE, "usec\n");
 	else
 		return snprintf(buf, PAGE_SIZE, "sysclk\n");
diff --git a/drivers/gpu/nvgpu/tu104/hal_tu104.c b/drivers/gpu/nvgpu/tu104/hal_tu104.c
index 7637f2c06..aec810ade 100644
--- a/drivers/gpu/nvgpu/tu104/hal_tu104.c
+++ b/drivers/gpu/nvgpu/tu104/hal_tu104.c
@@ -43,6 +43,7 @@
 #include "hal/gr/zbc/zbc_gp10b.h"
 #include "hal/gr/zbc/zbc_gv11b.h"
 #include "hal/gr/init/gr_init_gm20b.h"
+#include "hal/gr/init/gr_init_gv11b.h"
 #include "hal/gr/hwpm_map/hwpm_map_gv100.h"
 #include "hal/falcon/falcon_gk20a.h"
 
@@ -468,7 +469,6 @@ static const struct gpu_ops tu104_ops = {
 		.set_preemption_mode = gr_gp10b_set_preemption_mode,
 		.pre_process_sm_exception = gr_gv11b_pre_process_sm_exception,
 		.set_preemption_buffer_va = gr_gv11b_set_preemption_buffer_va,
-		.init_preemption_state = gr_gv11b_init_preemption_state,
 		.set_bes_crop_debug3 = gr_gp10b_set_bes_crop_debug3,
 		.set_bes_crop_debug4 = gr_gp10b_set_bes_crop_debug4,
 		.init_ecc = tu104_ecc_init,
@@ -656,15 +656,16 @@ static const struct gpu_ops tu104_ops = {
 				gv100_gr_hwpm_map_get_active_fbpa_mask,
 		},
 		.init = {
+			.pd_tpc_per_gpc = gm20b_gr_init_pd_tpc_per_gpc,
+			.pd_skip_table_gpc = gm20b_gr_init_pd_skip_table_gpc,
+			.cwd_gpcs_tpcs_num = gm20b_gr_init_cwd_gpcs_tpcs_num,
+			.wait_idle = gm20b_gr_init_wait_idle,
+			.wait_fe_idle = gm20b_gr_init_wait_fe_idle,
 			.fe_pwr_mode_force_on =
 				gm20b_gr_init_fe_pwr_mode_force_on,
 			.override_context_reset =
 				gm20b_gr_init_override_context_reset,
-			.wait_idle = gm20b_gr_init_wait_idle,
-			.wait_fe_idle = gm20b_gr_init_wait_fe_idle,
-			.pd_tpc_per_gpc = gm20b_gr_init_pd_tpc_per_gpc,
-			.pd_skip_table_gpc = gm20b_gr_init_pd_skip_table_gpc,
-			.cwd_gpcs_tpcs_num = gm20b_gr_init_cwd_gpcs_tpcs_num,
+			.preemption_state = gv11b_gr_init_preemption_state,
 		},
 	},
 	.fb = {