diff --git a/drivers/gpu/nvgpu/common/as.c b/drivers/gpu/nvgpu/common/as.c
index 1b5e2b563..8876cf6d6 100644
--- a/drivers/gpu/nvgpu/common/as.c
+++ b/drivers/gpu/nvgpu/common/as.c
@@ -57,6 +57,9 @@ static int gk20a_vm_alloc_share(struct gk20a_as_share *as_share,
 	char name[32];
 	const bool userspace_managed =
 		(flags & NVGPU_AS_ALLOC_USERSPACE_MANAGED) != 0;
+	const bool unified_va =
+		nvgpu_is_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES) ||
+		(flags & NVGPU_AS_ALLOC_UNIFIED_VA) != 0;
 
 	nvgpu_log_fn(g, " ");
 
@@ -79,7 +82,8 @@ static int gk20a_vm_alloc_share(struct gk20a_as_share *as_share,
 			   U64(big_page_size) << U64(10),
 			   mm->channel.kernel_size,
 			   mm->channel.user_size + mm->channel.kernel_size,
-			   !mm->disable_bigpage, userspace_managed, name);
+			   !mm->disable_bigpage,
+			   userspace_managed, unified_va, name);
 	if (vm == NULL) {
 		return -ENOMEM;
 	}
diff --git a/drivers/gpu/nvgpu/common/mm/mm.c b/drivers/gpu/nvgpu/common/mm/mm.c
index b78085efb..8cc333b75 100644
--- a/drivers/gpu/nvgpu/common/mm/mm.c
+++ b/drivers/gpu/nvgpu/common/mm/mm.c
@@ -99,7 +99,7 @@ u32 nvgpu_vm_get_pte_size(struct vm_gk20a *vm, u64 base, u64 size)
 		return GMMU_PAGE_SIZE_SMALL;
 	}
 
-	if (!nvgpu_is_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES)) {
+	if (!vm->unified_va) {
 		return nvgpu_vm_get_pte_size_split_addr(vm, base, size);
 	}
 
@@ -237,6 +237,7 @@ static int nvgpu_init_system_vm(struct mm_gk20a *mm)
 				   aperture_size,
 				   true,
 				   false,
+				   false,
 				   "system");
 	if (mm->pmu.vm == NULL) {
 		return -ENOMEM;
@@ -279,7 +280,7 @@ static int nvgpu_init_cde_vm(struct mm_gk20a *mm)
 				   U64(big_page_size) << U64(10),
 				   NV_MM_DEFAULT_KERNEL_SIZE,
 				   NV_MM_DEFAULT_KERNEL_SIZE + NV_MM_DEFAULT_USER_SIZE,
-				   false, false, "cde");
+				   false, false, false, "cde");
 	if (mm->cde.vm == NULL) {
 		return -ENOMEM;
 	}
@@ -295,7 +296,7 @@ static int nvgpu_init_ce_vm(struct mm_gk20a *mm)
 				  U64(big_page_size) << U64(10),
 				  NV_MM_DEFAULT_KERNEL_SIZE,
 				  NV_MM_DEFAULT_KERNEL_SIZE + NV_MM_DEFAULT_USER_SIZE,
-				  false, false, "ce");
+				  false, false, false, "ce");
 	if (mm->ce.vm == NULL) {
 		return -ENOMEM;
 	}
@@ -389,7 +390,7 @@ static int nvgpu_init_bar1_vm(struct mm_gk20a *mm)
 				    SZ_64K,
 				    mm->bar1.aperture_size - SZ_64K,
 				    mm->bar1.aperture_size,
-				    true, false,
+				    true, false, false,
 				    "bar1");
 	if (mm->bar1.vm == NULL) {
 		return -ENOMEM;
diff --git a/drivers/gpu/nvgpu/common/mm/vm.c b/drivers/gpu/nvgpu/common/mm/vm.c
index 723ecedd7..3c34456e9 100644
--- a/drivers/gpu/nvgpu/common/mm/vm.c
+++ b/drivers/gpu/nvgpu/common/mm/vm.c
@@ -282,6 +282,7 @@ int __nvgpu_vm_init(struct mm_gk20a *mm,
 			   u64 aperture_size,
 			   bool big_pages,
 			   bool userspace_managed,
+			   bool unified_va,
 			   const char *name)
 {
 	int err = 0;
@@ -315,7 +316,7 @@ int __nvgpu_vm_init(struct mm_gk20a *mm,
 	vm->vma[GMMU_PAGE_SIZE_SMALL]  = &vm->user;
 	vm->vma[GMMU_PAGE_SIZE_BIG]    = &vm->user;
 	vm->vma[GMMU_PAGE_SIZE_KERNEL] = &vm->kernel;
-	if (!nvgpu_is_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES)) {
+	if (!unified_va) {
 		vm->vma[GMMU_PAGE_SIZE_BIG] = &vm->user_lp;
 	}
 
@@ -324,6 +325,7 @@ int __nvgpu_vm_init(struct mm_gk20a *mm,
 
 	vm->big_page_size     = vm->gmmu_page_sizes[GMMU_PAGE_SIZE_BIG];
 	vm->userspace_managed = userspace_managed;
+	vm->unified_va        = unified_va;
 	vm->mmu_levels        = g->ops.mm.get_mmu_levels(g, vm->big_page_size);
 
 #ifdef CONFIG_TEGRA_GR_VIRTUALIZATION
@@ -351,8 +353,7 @@ int __nvgpu_vm_init(struct mm_gk20a *mm,
 		 * sense to make one VM, same as if the unified address flag
 		 * is set.
 		 */
-		if (!big_pages ||
-		    nvgpu_is_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES)) {
+		if (!big_pages || unified_va) {
 			user_vma_start = low_hole;
 			user_vma_limit = vm->va_limit - kernel_reserved;
 			user_lp_vma_start = user_vma_limit;
@@ -405,7 +406,7 @@ int __nvgpu_vm_init(struct mm_gk20a *mm,
 	 * Determine if big pages are possible in this VM. If a split address
 	 * space is used then check the user_lp vma instead of the user vma.
 	 */
-	if (nvgpu_is_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES)) {
+	if (unified_va) {
 		vm->big_pages = big_pages &&
 			nvgpu_big_pages_possible(vm, user_vma_start,
 					user_vma_limit - user_vma_start);
@@ -577,6 +578,7 @@ struct vm_gk20a *nvgpu_vm_init(struct gk20a *g,
 			       u64 aperture_size,
 			       bool big_pages,
 			       bool userspace_managed,
+			       bool unified_va,
 			       const char *name)
 {
 	struct vm_gk20a *vm = nvgpu_kzalloc(g, sizeof(*vm));
@@ -587,7 +589,7 @@ struct vm_gk20a *nvgpu_vm_init(struct gk20a *g,
 
 	if (__nvgpu_vm_init(&g->mm, vm, big_page_size, low_hole,
 			    kernel_reserved, aperture_size, big_pages,
-			    userspace_managed, name) != 0) {
+			    userspace_managed, unified_va, name) != 0) {
 		nvgpu_kfree(g, vm);
 		return NULL;
 	}
diff --git a/drivers/gpu/nvgpu/gp10b/mm_gp10b.c b/drivers/gpu/nvgpu/gp10b/mm_gp10b.c
index 9e0bf185d..b381aec86 100644
--- a/drivers/gpu/nvgpu/gp10b/mm_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/mm_gp10b.c
@@ -59,7 +59,7 @@ int gp10b_init_bar2_vm(struct gk20a *g)
 
 	mm->bar2.vm = nvgpu_vm_init(g, big_page_size, SZ_4K,
 		mm->bar2.aperture_size - SZ_4K,
-		mm->bar2.aperture_size, false, false, "bar2");
+		mm->bar2.aperture_size, false, false, false, "bar2");
 	if (mm->bar2.vm == NULL) {
 		return -ENOMEM;
 	}
diff --git a/drivers/gpu/nvgpu/include/nvgpu/as.h b/drivers/gpu/nvgpu/include/nvgpu/as.h
index f2249f9e3..46256b3d4 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/as.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/as.h
@@ -42,7 +42,8 @@ struct gk20a_as_share {
 /*
  * AS allocation flags.
  */
-#define NVGPU_AS_ALLOC_USERSPACE_MANAGED	(1 << 0)
+#define NVGPU_AS_ALLOC_USERSPACE_MANAGED	BIT32(0)
+#define NVGPU_AS_ALLOC_UNIFIED_VA		BIT32(1)
 
 int gk20a_as_release_share(struct gk20a_as_share *as_share);
 
diff --git a/drivers/gpu/nvgpu/include/nvgpu/vm.h b/drivers/gpu/nvgpu/include/nvgpu/vm.h
index 5b07a1857..dfe8c3f9c 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/vm.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/vm.h
@@ -159,6 +159,7 @@ struct vm_gk20a {
 	u32 big_page_size;
 
 	bool userspace_managed;
+	bool unified_va;
 
 	const struct gk20a_mmu_level *mmu_levels;
 
@@ -302,6 +303,7 @@ int __nvgpu_vm_init(struct mm_gk20a *mm,
 			   u64 aperture_size,
 			   bool big_pages,
 			   bool userspace_managed,
+			   bool unified_va,
 			   const char *name);
 
 struct vm_gk20a *nvgpu_vm_init(struct gk20a *g,
@@ -311,6 +313,7 @@ struct vm_gk20a *nvgpu_vm_init(struct gk20a *g,
 			       u64 aperture_size,
 			       bool big_pages,
 			       bool userspace_managed,
+			       bool unified_va,
 			       const char *name);
 
 /*
diff --git a/drivers/gpu/nvgpu/os/linux/ioctl_ctrl.c b/drivers/gpu/nvgpu/os/linux/ioctl_ctrl.c
index 42cbfcafc..4381c378a 100644
--- a/drivers/gpu/nvgpu/os/linux/ioctl_ctrl.c
+++ b/drivers/gpu/nvgpu/os/linux/ioctl_ctrl.c
@@ -77,6 +77,8 @@ static u32 gk20a_as_translate_as_alloc_flags(struct gk20a *g, u32 flags)
 
 	if (flags & NVGPU_GPU_IOCTL_ALLOC_AS_FLAGS_USERSPACE_MANAGED)
 		core_flags |= NVGPU_AS_ALLOC_USERSPACE_MANAGED;
+	if (flags & NVGPU_GPU_IOCTL_ALLOC_AS_FLAGS_UNIFIED_VA)
+		core_flags |= NVGPU_AS_ALLOC_UNIFIED_VA;
 
 	return core_flags;
 }
diff --git a/drivers/gpu/nvgpu/os/linux/ioctl_dbg.c b/drivers/gpu/nvgpu/os/linux/ioctl_dbg.c
index fc15d4971..dd6dae5cb 100644
--- a/drivers/gpu/nvgpu/os/linux/ioctl_dbg.c
+++ b/drivers/gpu/nvgpu/os/linux/ioctl_dbg.c
@@ -1354,7 +1354,7 @@ static int gk20a_perfbuf_map(struct dbg_session_gk20a *dbg_s,
 			big_page_size << 10,
 			NV_MM_DEFAULT_KERNEL_SIZE,
 			NV_MM_DEFAULT_KERNEL_SIZE + NV_MM_DEFAULT_USER_SIZE,
-			false, false, "perfbuf");
+			false, false, false, "perfbuf");
 	if (!mm->perfbuf.vm) {
 		nvgpu_mutex_release(&g->dbg_sessions_lock);
 		return -ENOMEM;
diff --git a/include/uapi/linux/nvgpu.h b/include/uapi/linux/nvgpu.h
index 93afdde43..031321e34 100644
--- a/include/uapi/linux/nvgpu.h
+++ b/include/uapi/linux/nvgpu.h
@@ -357,7 +357,8 @@ struct nvgpu_alloc_as_args {
  *      increments at kickoffs and decrements at job completion are
  *      bypassed.
  */
-#define NVGPU_GPU_IOCTL_ALLOC_AS_FLAGS_USERSPACE_MANAGED (1 << 0)
+#define NVGPU_GPU_IOCTL_ALLOC_AS_FLAGS_USERSPACE_MANAGED	(1 << 0)
+#define NVGPU_GPU_IOCTL_ALLOC_AS_FLAGS_UNIFIED_VA	 	(1 << 1)
 	__u32 flags;
 
 	__u32 reserved;			/* must be zero */
diff --git a/userspace/units/mm/page_table/page_table.c b/userspace/units/mm/page_table/page_table.c
index 729f7d396..d84fe8830 100644
--- a/userspace/units/mm/page_table/page_table.c
+++ b/userspace/units/mm/page_table/page_table.c
@@ -240,6 +240,7 @@ static int init_mm(struct unit_module *m, struct gk20a *g)
 				   aperture_size,
 				   true,
 				   false,
+				   false,
 				   "system");
 	if (mm->pmu.vm == NULL) {
 		unit_return_fail(m, "nvgpu_vm_init failed\n");