diff --git a/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c b/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
index 41d14a823..8e9e18188 100644
--- a/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
@@ -188,6 +188,9 @@ int channel_gv11b_setup_ramfc(struct channel_gk20a *c,
 	nvgpu_mem_wr32(g, mem, ram_fc_set_channel_info_w(),
 			pbdma_set_channel_info_veid_f(c->subctx_id));
 
+	nvgpu_mem_wr32(g, mem, ram_in_engine_wfi_veid_w(),
+			ram_in_engine_wfi_veid_f(c->subctx_id));
+
 	gv11b_fifo_init_ramfc_eng_method_buffer(g, c, mem);
 
 	if (c->is_privileged_channel) {
diff --git a/drivers/gpu/nvgpu/gv11b/mm_gv11b.c b/drivers/gpu/nvgpu/gv11b/mm_gv11b.c
index 8d307d7c6..ade1d9fec 100644
--- a/drivers/gpu/nvgpu/gv11b/mm_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/mm_gv11b.c
@@ -35,6 +35,7 @@
 
 #include "mm_gv11b.h"
 #include "fb_gv11b.h"
+#include "subctx_gv11b.h"
 
 #include <nvgpu/hw/gv11b/hw_fb_gv11b.h>
 #include <nvgpu/hw/gv11b/hw_gmmu_gv11b.h>
@@ -59,6 +60,8 @@ void gv11b_init_inst_block(struct nvgpu_mem *inst_block,
 
 	if (big_page_size && g->ops.mm.set_big_page_size)
 		g->ops.mm.set_big_page_size(g, inst_block, big_page_size);
+
+	gv11b_init_subcontext_pdb(vm, inst_block);
 }
 
 bool gv11b_mm_mmu_fault_pending(struct gk20a *g)
diff --git a/drivers/gpu/nvgpu/gv11b/subctx_gv11b.c b/drivers/gpu/nvgpu/gv11b/subctx_gv11b.c
index 4f98d82a5..05d7dee02 100644
--- a/drivers/gpu/nvgpu/gv11b/subctx_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/subctx_gv11b.c
@@ -33,12 +33,9 @@
 #include <nvgpu/hw/gv11b/hw_ram_gv11b.h>
 #include <nvgpu/hw/gv11b/hw_ctxsw_prog_gv11b.h>
 
-static void gv11b_init_subcontext_pdb(struct channel_gk20a *c,
+static void gv11b_subctx_commit_valid_mask(struct vm_gk20a *vm,
 				struct nvgpu_mem *inst_block);
-
-static void gv11b_subctx_commit_valid_mask(struct channel_gk20a *c,
-				struct nvgpu_mem *inst_block);
-static void gv11b_subctx_commit_pdb(struct channel_gk20a *c,
+static void gv11b_subctx_commit_pdb(struct vm_gk20a *vm,
 				struct nvgpu_mem *inst_block);
 
 void gv11b_free_subctx_header(struct channel_gk20a *c)
@@ -89,23 +86,15 @@ int gv11b_alloc_subctx_header(struct channel_gk20a *c)
 
 		nvgpu_memset(g, &ctx->mem, 0, 0, ctx->mem.size);
 		nvgpu_mem_end(g, &ctx->mem);
-
-		gv11b_init_subcontext_pdb(c, &c->inst_block);
 	}
 	return ret;
 }
 
-static void gv11b_init_subcontext_pdb(struct channel_gk20a *c,
+void gv11b_init_subcontext_pdb(struct vm_gk20a *vm,
 				struct nvgpu_mem *inst_block)
 {
-	struct gk20a *g = c->g;
-
-	gv11b_subctx_commit_pdb(c, inst_block);
-	gv11b_subctx_commit_valid_mask(c, inst_block);
-
-	nvgpu_log(g, gpu_dbg_info, " subctx %d instblk set", c->subctx_id);
-	nvgpu_mem_wr32(g, inst_block, ram_in_engine_wfi_veid_w(),
-			ram_in_engine_wfi_veid_f(c->subctx_id));
+	gv11b_subctx_commit_pdb(vm, inst_block);
+	gv11b_subctx_commit_valid_mask(vm, inst_block);
 
 }
 
@@ -167,22 +156,21 @@ int gv11b_update_subctx_header(struct channel_gk20a *c, u64 gpu_va)
 	return ret;
 }
 
-void gv11b_subctx_commit_valid_mask(struct channel_gk20a *c,
+void gv11b_subctx_commit_valid_mask(struct vm_gk20a *vm,
 				struct nvgpu_mem *inst_block)
 {
-	struct gk20a *g = c->g;
+	struct gk20a *g = gk20a_from_vm(vm);
 
 	/* Make all subctx pdbs valid */
 	nvgpu_mem_wr32(g, inst_block, 166, 0xffffffff);
 	nvgpu_mem_wr32(g, inst_block, 167, 0xffffffff);
 }
 
-void gv11b_subctx_commit_pdb(struct channel_gk20a *c,
+void gv11b_subctx_commit_pdb(struct vm_gk20a *vm,
 				struct nvgpu_mem *inst_block)
 {
-	struct gk20a *g = c->g;
+	struct gk20a *g = gk20a_from_vm(vm);
 	struct fifo_gk20a *f = &g->fifo;
-	struct vm_gk20a *vm = c->vm;
 	u32 lo, hi;
 	u32 subctx_id = 0;
 	u32 format_word;
diff --git a/drivers/gpu/nvgpu/gv11b/subctx_gv11b.h b/drivers/gpu/nvgpu/gv11b/subctx_gv11b.h
index 10dc0ba5d..3dcb58b34 100644
--- a/drivers/gpu/nvgpu/gv11b/subctx_gv11b.h
+++ b/drivers/gpu/nvgpu/gv11b/subctx_gv11b.h
@@ -2,7 +2,7 @@
  *
  * Volta GPU series Subcontext
  *
- * Copyright (c) 2016 - 2017, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2016 - 2018, 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"),
@@ -31,4 +31,7 @@ void gv11b_free_subctx_header(struct channel_gk20a *c);
 
 int gv11b_update_subctx_header(struct channel_gk20a *c, u64 gpu_va);
 
+void gv11b_init_subcontext_pdb(struct vm_gk20a *vm,
+				struct nvgpu_mem *inst_block);
+
 #endif /* __SUBCONTEXT_GV11B_H__ */