diff --git a/drivers/gpu/nvgpu/common/gr/ctx.c b/drivers/gpu/nvgpu/common/gr/ctx.c
index c4ddf3a1f..7f92ff983 100644
--- a/drivers/gpu/nvgpu/common/gr/ctx.c
+++ b/drivers/gpu/nvgpu/common/gr/ctx.c
@@ -614,52 +614,51 @@ u32 nvgpu_gr_ctx_get_compute_preemption_mode(struct nvgpu_gr_ctx *gr_ctx)
 	return gr_ctx->compute_preempt_mode;
 }
 
-bool nvgpu_gr_ctx_check_valid_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
-	u32 graphics_preempt_mode, u32 compute_preempt_mode)
+bool nvgpu_gr_ctx_check_valid_preemption_mode(struct gk20a *g,
+		struct nvgpu_gr_ctx *gr_ctx,
+		u32 graphics_preempt_mode, u32 compute_preempt_mode)
 {
-#ifdef CONFIG_NVGPU_GRAPHICS
+	u32 supported_graphics_preempt_mode = 0U;
+	u32 supported_compute_preempt_mode = 0U;
+
 	if ((graphics_preempt_mode == 0U) && (compute_preempt_mode == 0U)) {
 		return false;
 	}
-#else
-	if (graphics_preempt_mode != 0U) {
-		return false;
-	}
-#endif
 
-#ifdef CONFIG_NVGPU_CILP
-	if (compute_preempt_mode > NVGPU_PREEMPTION_MODE_COMPUTE_CILP) {
-		return false;
+	g->ops.gr.init.get_supported__preemption_modes(
+			&supported_graphics_preempt_mode,
+			&supported_compute_preempt_mode);
+
+	if (graphics_preempt_mode != 0U) {
+		if ((graphics_preempt_mode & supported_graphics_preempt_mode) == 0U) {
+			return false;
+		}
+
+		/* Do not allow lower preemption modes than current ones */
+		if (graphics_preempt_mode < gr_ctx->graphics_preempt_mode) {
+			return false;
+		}
 	}
-#else
-	if (compute_preempt_mode > NVGPU_PREEMPTION_MODE_COMPUTE_CTA) {
-		return false;
+
+	if (compute_preempt_mode != 0U) {
+		if ((compute_preempt_mode & supported_compute_preempt_mode) == 0U) {
+			return false;
+		}
+
+		/* Do not allow lower preemption modes than current ones */
+		if (compute_preempt_mode < gr_ctx->compute_preempt_mode) {
+			return false;
+		}
 	}
-#endif
 
 #if defined(CONFIG_NVGPU_CILP) && defined(CONFIG_NVGPU_GRAPHICS)
+	/* Invalid combination */
 	if ((graphics_preempt_mode == NVGPU_PREEMPTION_MODE_GRAPHICS_GFXP) &&
 		   (compute_preempt_mode == NVGPU_PREEMPTION_MODE_COMPUTE_CILP)) {
 		return false;
 	}
 #endif
 
-#ifdef CONFIG_NVGPU_GRAPHICS
-	/* Do not allow lower preemption modes than current ones */
-	if ((graphics_preempt_mode != 0U) &&
-	    (graphics_preempt_mode < gr_ctx->graphics_preempt_mode)) {
-		return false;
-	}
-#endif
-
-	if (
-#ifdef CONFIG_NVGPU_GRAPHICS
-	    (compute_preempt_mode != 0U) &&
-#endif
-	    (compute_preempt_mode < gr_ctx->compute_preempt_mode)) {
-		return false;
-	}
-
 	return true;
 }
 
diff --git a/drivers/gpu/nvgpu/common/gr/obj_ctx.c b/drivers/gpu/nvgpu/common/gr/obj_ctx.c
index 48b0b31dd..55c919b2e 100644
--- a/drivers/gpu/nvgpu/common/gr/obj_ctx.c
+++ b/drivers/gpu/nvgpu/common/gr/obj_ctx.c
@@ -76,37 +76,57 @@ static int nvgpu_gr_obj_ctx_init_ctxsw_preemption_mode(struct gk20a *g,
 	int err;
 	u32 graphics_preempt_mode = 0U;
 	u32 compute_preempt_mode = 0U;
+	u32 default_graphics_preempt_mode = 0U;
+	u32 default_compute_preempt_mode = 0U;
 
 	nvgpu_log(g, gpu_dbg_fn | gpu_dbg_gr, " ");
 
-#ifdef CONFIG_NVGPU_GRAPHICS
-	if (!nvgpu_is_enabled(g, NVGPU_SUPPORT_PREEMPTION_GFXP)) {
-		if (g->ops.gpu_class.is_valid_compute(class_num)) {
-			nvgpu_gr_ctx_init_compute_preemption_mode(gr_ctx,
-				NVGPU_PREEMPTION_MODE_COMPUTE_CTA);
-		}
+	/* Skip for engines other than GR */
+	if (!g->ops.gpu_class.is_valid_compute(class_num) &&
+	    !g->ops.gpu_class.is_valid_gfx(class_num)) {
 		return 0;
 	}
 
+	g->ops.gr.init.get_default_preemption_modes(
+			&default_graphics_preempt_mode,
+			&default_compute_preempt_mode);
+
+#ifdef CONFIG_NVGPU_GRAPHICS
 	if ((flags & NVGPU_OBJ_CTX_FLAGS_SUPPORT_GFXP) != 0U) {
 		graphics_preempt_mode = NVGPU_PREEMPTION_MODE_GRAPHICS_GFXP;
 	}
+
+	if (g->ops.gpu_class.is_valid_gfx(class_num) &&
+			nvgpu_gr_ctx_desc_force_preemption_gfxp(gr_ctx_desc)) {
+		graphics_preempt_mode = NVGPU_PREEMPTION_MODE_GRAPHICS_GFXP;
+	}
+
+	if (graphics_preempt_mode == 0U) {
+		graphics_preempt_mode = default_graphics_preempt_mode;
+	}
 #endif
 
 #ifdef CONFIG_NVGPU_CILP
 	if ((flags & NVGPU_OBJ_CTX_FLAGS_SUPPORT_CILP) != 0U) {
 		compute_preempt_mode = NVGPU_PREEMPTION_MODE_COMPUTE_CILP;
 	}
+
+	if (g->ops.gpu_class.is_valid_compute(class_num) &&
+			nvgpu_gr_ctx_desc_force_preemption_cilp(gr_ctx_desc)) {
+		compute_preempt_mode = NVGPU_PREEMPTION_MODE_COMPUTE_CILP;
+	}
 #endif
 
-	if ((graphics_preempt_mode != 0U) || (compute_preempt_mode != 0U)) {
-		err = nvgpu_gr_obj_ctx_set_ctxsw_preemption_mode(g, config,
-			gr_ctx_desc, gr_ctx, vm, class_num, graphics_preempt_mode,
-			compute_preempt_mode);
-		if (err != 0) {
-			nvgpu_err(g, "set_ctxsw_preemption_mode failed");
-			return err;
-		}
+	if (compute_preempt_mode == 0U) {
+		compute_preempt_mode = default_compute_preempt_mode;
+	}
+
+	err = nvgpu_gr_obj_ctx_set_ctxsw_preemption_mode(g, config,
+		gr_ctx_desc, gr_ctx, vm, class_num, graphics_preempt_mode,
+		compute_preempt_mode);
+	if (err != 0) {
+		nvgpu_err(g, "set_ctxsw_preemption_mode failed");
+		return err;
 	}
 
 	nvgpu_log(g, gpu_dbg_fn | gpu_dbg_gr, "done");
@@ -206,31 +226,16 @@ int nvgpu_gr_obj_ctx_set_ctxsw_preemption_mode(struct gk20a *g,
 {
 	int err = 0;
 
-#ifdef CONFIG_NVGPU_GRAPHICS
-	if (!nvgpu_is_enabled(g, NVGPU_SUPPORT_PREEMPTION_GFXP)) {
-		return 0;
-	}
-
-	if (g->ops.gpu_class.is_valid_gfx(class_num) &&
-			nvgpu_gr_ctx_desc_force_preemption_gfxp(gr_ctx_desc)) {
-		graphics_preempt_mode = NVGPU_PREEMPTION_MODE_GRAPHICS_GFXP;
-	}
-#endif
-
-#ifdef CONFIG_NVGPU_CILP
-	if (g->ops.gpu_class.is_valid_compute(class_num) &&
-			nvgpu_gr_ctx_desc_force_preemption_cilp(gr_ctx_desc)) {
-		compute_preempt_mode = NVGPU_PREEMPTION_MODE_COMPUTE_CILP;
-	}
-#endif
-
 	/* check for invalid combinations */
-	if (nvgpu_gr_ctx_check_valid_preemption_mode(gr_ctx,
+	if (nvgpu_gr_ctx_check_valid_preemption_mode(g, gr_ctx,
 			graphics_preempt_mode, compute_preempt_mode) == false) {
 		err = -EINVAL;
 		goto fail;
 	}
 
+	nvgpu_log(g, gpu_dbg_gr, "graphics_preempt_mode=%u compute_preempt_mode=%u",
+			graphics_preempt_mode, compute_preempt_mode);
+
 #ifdef CONFIG_NVGPU_GRAPHICS
 	err = nvgpu_gr_obj_ctx_set_graphics_preemption_mode(g, config,
 				gr_ctx_desc, gr_ctx, vm, graphics_preempt_mode);
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
index 89f248a9f..5b2b4c6c1 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
@@ -437,6 +437,7 @@ u32 nvgpu_gr_ctx_get_compute_preemption_mode(struct nvgpu_gr_ctx *gr_ctx);
 /**
  * @brief Check if given preemption modes are valid.
  *
+ * @param g [in]			Pointer to GPU driver struct.
  * @param gr_ctx [in]			Pointer to graphics context struct.
  * @param graphics_preempt_mode		Requested graphics preemption mode.
  * @param compute_preempt_mode		Requested compute preemption mode.
@@ -455,7 +456,8 @@ u32 nvgpu_gr_ctx_get_compute_preemption_mode(struct nvgpu_gr_ctx *gr_ctx);
  *
  * @return true if requested preemption modes are valid, false otherwise.
  */
-bool nvgpu_gr_ctx_check_valid_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
+bool nvgpu_gr_ctx_check_valid_preemption_mode(struct gk20a *g,
+	struct nvgpu_gr_ctx *gr_ctx,
 	u32 graphics_preempt_mode, u32 compute_preempt_mode);
 
 /**
diff --git a/userspace/required_tests.ini b/userspace/required_tests.ini
index cfaef958b..075e20db3 100644
--- a/userspace/required_tests.ini
+++ b/userspace/required_tests.ini
@@ -670,7 +670,7 @@ test_gr_init_setup_ready.gr_setup_setup=0
 test_gr_setup_alloc_obj_ctx.gr_setup_alloc_obj_ctx=0
 test_gr_setup_alloc_obj_ctx_error_injections.gr_setup_alloc_obj_ctx_error_injections=0
 test_gr_setup_free_obj_ctx.gr_setup_free_obj_ctx=0
-test_gr_setup_preemption_mode_errors.gr_setup_preemption_mode_errors=0
+test_gr_setup_preemption_mode_errors.gr_setup_preemption_mode_errors=2
 test_gr_setup_set_preemption_mode.gr_setup_set_preemption_mode=0
 
 [nvgpu_mem]
diff --git a/userspace/units/gr/setup/nvgpu-gr-setup.c b/userspace/units/gr/setup/nvgpu-gr-setup.c
index acb019fa1..cc3a1a54a 100644
--- a/userspace/units/gr/setup/nvgpu-gr-setup.c
+++ b/userspace/units/gr/setup/nvgpu-gr-setup.c
@@ -691,7 +691,7 @@ struct unit_module_test nvgpu_gr_setup_tests[] = {
 	UNIT_TEST(gr_setup_set_preemption_mode,
 			test_gr_setup_set_preemption_mode, NULL, 0),
 	UNIT_TEST(gr_setup_preemption_mode_errors,
-			test_gr_setup_preemption_mode_errors, NULL, 0),
+			test_gr_setup_preemption_mode_errors, NULL, 2),
 	UNIT_TEST(gr_setup_free_obj_ctx, test_gr_setup_free_obj_ctx, NULL, 0),
 	UNIT_TEST(gr_setup_alloc_obj_ctx_error_injections,
 			test_gr_setup_alloc_obj_ctx_error_injections, NULL, 0),