diff --git a/drivers/gpu/nvgpu/common/fb/fb_gm20b.c b/drivers/gpu/nvgpu/common/fb/fb_gm20b.c
index 65b7336c9..00d0fba42 100644
--- a/drivers/gpu/nvgpu/common/fb/fb_gm20b.c
+++ b/drivers/gpu/nvgpu/common/fb/fb_gm20b.c
@@ -55,9 +55,32 @@ void gm20b_fb_reset(struct gk20a *g)
 
 void gm20b_fb_init_hw(struct gk20a *g)
 {
-	u32 addr = nvgpu_mem_get_addr(g, &g->mm.sysmem_flush) >> 8;
+	u64 addr = nvgpu_mem_get_addr(g, &g->mm.sysmem_flush) >> 8;
 
 	gk20a_writel(g, fb_niso_flush_sysmem_addr_r(), addr);
+
+	/* init mmu debug buffer */
+	addr = nvgpu_mem_get_addr(g, &g->mm.mmu_wr_mem);
+	addr >>= fb_mmu_debug_wr_addr_alignment_v();
+
+	gk20a_writel(g, fb_mmu_debug_wr_r(),
+		     nvgpu_aperture_mask(g, &g->mm.mmu_wr_mem,
+				fb_mmu_debug_wr_aperture_sys_mem_ncoh_f(),
+				fb_mmu_debug_wr_aperture_sys_mem_coh_f(),
+				fb_mmu_debug_wr_aperture_vid_mem_f()) |
+		     fb_mmu_debug_wr_vol_false_f() |
+		     fb_mmu_debug_wr_addr_f(addr));
+
+	addr = nvgpu_mem_get_addr(g, &g->mm.mmu_rd_mem);
+	addr >>= fb_mmu_debug_rd_addr_alignment_v();
+
+	gk20a_writel(g, fb_mmu_debug_rd_r(),
+		     nvgpu_aperture_mask(g, &g->mm.mmu_rd_mem,
+				fb_mmu_debug_wr_aperture_sys_mem_ncoh_f(),
+				fb_mmu_debug_wr_aperture_sys_mem_coh_f(),
+				fb_mmu_debug_rd_aperture_vid_mem_f()) |
+		     fb_mmu_debug_rd_vol_false_f() |
+		     fb_mmu_debug_rd_addr_f(addr));
 }
 
 int gm20b_fb_tlb_invalidate(struct gk20a *g, struct nvgpu_mem *pdb)
diff --git a/drivers/gpu/nvgpu/common/mm/mm.c b/drivers/gpu/nvgpu/common/mm/mm.c
index 54f621aee..988b1e5ca 100644
--- a/drivers/gpu/nvgpu/common/mm/mm.c
+++ b/drivers/gpu/nvgpu/common/mm/mm.c
@@ -173,6 +173,9 @@ static void nvgpu_remove_mm_support(struct mm_gk20a *mm)
 {
 	struct gk20a *g = gk20a_from_mm(mm);
 
+	nvgpu_dma_free(g, &mm->mmu_wr_mem);
+	nvgpu_dma_free(g, &mm->mmu_rd_mem);
+
 	if (g->ops.mm.fault_info_mem_destroy) {
 		g->ops.mm.fault_info_mem_destroy(g);
 	}
@@ -294,6 +297,32 @@ static int nvgpu_init_ce_vm(struct mm_gk20a *mm)
 	return 0;
 }
 
+static int nvgpu_init_mmu_debug(struct mm_gk20a *mm)
+{
+	struct gk20a *g = gk20a_from_mm(mm);
+	int err;
+
+	if (!nvgpu_mem_is_valid(&mm->mmu_wr_mem)) {
+		err = nvgpu_dma_alloc_sys(g, SZ_4K, &mm->mmu_wr_mem);
+		if (err) {
+			goto err;
+		}
+	}
+
+	if (!nvgpu_mem_is_valid(&mm->mmu_rd_mem)) {
+		err = nvgpu_dma_alloc_sys(g, SZ_4K, &mm->mmu_rd_mem);
+		if (err) {
+			goto err_free_wr_mem;
+		}
+	}
+	return 0;
+
+ err_free_wr_mem:
+	nvgpu_dma_free(g, &mm->mmu_wr_mem);
+ err:
+	return -ENOMEM;
+}
+
 void nvgpu_init_mm_ce_context(struct gk20a *g)
 {
 #if defined(CONFIG_GK20A_VIDMEM)
@@ -459,6 +488,10 @@ static int nvgpu_init_mm_setup_sw(struct gk20a *g)
 		return err;
 	}
 
+	err = nvgpu_init_mmu_debug(mm);
+	if (err)
+		return err;
+
 	mm->remove_support = nvgpu_remove_mm_support;
 	mm->remove_ce_support = nvgpu_remove_mm_ce_support;
 
diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index f3b580e46..2a871b5a6 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -60,7 +60,6 @@
 #include <nvgpu/hw/gk20a/hw_ram_gk20a.h>
 #include <nvgpu/hw/gk20a/hw_pri_ringmaster_gk20a.h>
 #include <nvgpu/hw/gk20a/hw_top_gk20a.h>
-#include <nvgpu/hw/gk20a/hw_fb_gk20a.h>
 #include <nvgpu/hw/gk20a/hw_pbdma_gk20a.h>
 
 #define BLK_SIZE (256)
@@ -3153,9 +3152,6 @@ static void gk20a_remove_gr_support(struct gr_gk20a *gr)
 
 	gr_gk20a_free_global_ctx_buffers(g);
 
-	nvgpu_dma_free(g, &gr->mmu_wr_mem);
-	nvgpu_dma_free(g, &gr->mmu_rd_mem);
-
 	nvgpu_dma_free(g, &gr->compbit_store.mem);
 
 	memset(&gr->compbit_store, 0, sizeof(struct compbit_store_desc));
@@ -3495,31 +3491,6 @@ clean_up:
 	return -ENOMEM;
 }
 
-static int gr_gk20a_init_mmu_sw(struct gk20a *g, struct gr_gk20a *gr)
-{
-	int err;
-
-	if (!nvgpu_mem_is_valid(&gr->mmu_wr_mem)) {
-		err = nvgpu_dma_alloc_sys(g, 0x1000, &gr->mmu_wr_mem);
-		if (err) {
-			goto err;
-		}
-	}
-
-	if (!nvgpu_mem_is_valid(&gr->mmu_rd_mem)) {
-		err = nvgpu_dma_alloc_sys(g, 0x1000, &gr->mmu_rd_mem);
-		if (err) {
-			goto err_free_wr_mem;
-		}
-	}
-	return 0;
-
- err_free_wr_mem:
-	nvgpu_dma_free(g, &gr->mmu_wr_mem);
- err:
-	return -ENOMEM;
-}
-
 static u32 prime_set[18] = {
 	2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61 };
 
@@ -4529,35 +4500,11 @@ static int gk20a_init_gr_setup_hw(struct gk20a *g)
 	struct aiv_list_gk20a *sw_ctx_load = &g->gr.ctx_vars.sw_ctx_load;
 	struct av_list_gk20a *sw_method_init = &g->gr.ctx_vars.sw_method_init;
 	u32 data;
-	u64 addr;
 	u32 last_method_data = 0;
 	u32 i, err;
 
 	nvgpu_log_fn(g, " ");
 
-	/* init mmu debug buffer */
-	addr = nvgpu_mem_get_addr(g, &gr->mmu_wr_mem);
-	addr >>= fb_mmu_debug_wr_addr_alignment_v();
-
-	gk20a_writel(g, fb_mmu_debug_wr_r(),
-		     nvgpu_aperture_mask(g, &gr->mmu_wr_mem,
-				fb_mmu_debug_wr_aperture_sys_mem_ncoh_f(),
-				fb_mmu_debug_wr_aperture_sys_mem_coh_f(),
-				fb_mmu_debug_wr_aperture_vid_mem_f()) |
-		     fb_mmu_debug_wr_vol_false_f() |
-		     fb_mmu_debug_wr_addr_f(addr));
-
-	addr = nvgpu_mem_get_addr(g, &gr->mmu_rd_mem);
-	addr >>= fb_mmu_debug_rd_addr_alignment_v();
-
-	gk20a_writel(g, fb_mmu_debug_rd_r(),
-		     nvgpu_aperture_mask(g, &gr->mmu_rd_mem,
-				fb_mmu_debug_wr_aperture_sys_mem_ncoh_f(),
-				fb_mmu_debug_wr_aperture_sys_mem_coh_f(),
-				fb_mmu_debug_rd_aperture_vid_mem_f()) |
-		     fb_mmu_debug_rd_vol_false_f() |
-		     fb_mmu_debug_rd_addr_f(addr));
-
 	if (g->ops.gr.init_gpc_mmu) {
 		g->ops.gr.init_gpc_mmu(g);
 	}
@@ -4940,11 +4887,6 @@ static int gk20a_init_gr_setup_sw(struct gk20a *g)
 		goto clean_up;
 	}
 
-	err = gr_gk20a_init_mmu_sw(g, gr);
-	if (err) {
-		goto clean_up;
-	}
-
 	err = gr_gk20a_init_map_tiles(g, gr);
 	if (err) {
 		goto clean_up;
diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.h b/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
index 617aad34e..c4ce49af4 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.h
@@ -382,9 +382,6 @@ struct gr_gk20a {
 
 	struct gr_ctx_buffer_desc global_ctx_buffer[NR_GLOBAL_CTX_BUF];
 
-	struct nvgpu_mem mmu_wr_mem;
-	struct nvgpu_mem mmu_rd_mem;
-
 	u8 *map_tiles;
 	u32 map_tile_count;
 	u32 map_row_offset;
diff --git a/drivers/gpu/nvgpu/include/nvgpu/mm.h b/drivers/gpu/nvgpu/include/nvgpu/mm.h
index e627296dc..bfce243cb 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/mm.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/mm.h
@@ -176,6 +176,9 @@ struct mm_gk20a {
 
 		nvgpu_atomic64_t bytes_pending;
 	} vidmem;
+
+	struct nvgpu_mem mmu_wr_mem;
+	struct nvgpu_mem mmu_rd_mem;
 };
 
 #define gk20a_from_mm(mm) ((mm)->g)