diff --git a/drivers/gpu/nvgpu/common/linux/dma.c b/drivers/gpu/nvgpu/common/linux/dma.c
index ca657da2b..b62c4593f 100644
--- a/drivers/gpu/nvgpu/common/linux/dma.c
+++ b/drivers/gpu/nvgpu/common/linux/dma.c
@@ -28,9 +28,11 @@
 #include <nvgpu/vidmem.h>
 
 #include <nvgpu/linux/dma.h>
+#include <nvgpu/linux/vidmem.h>
 
 #include "gk20a/gk20a.h"
 #include "gk20a/platform_gk20a.h"
+
 #include "os_linux.h"
 
 /*
@@ -348,6 +350,7 @@ int nvgpu_dma_alloc_flags_vid_at(struct gk20a *g, unsigned long flags,
 
 	mem->aligned_size = size;
 	mem->aperture = APERTURE_VIDMEM;
+	mem->vidmem_alloc = (struct nvgpu_page_alloc *)(uintptr_t)addr;
 	mem->allocator = vidmem_alloc;
 	mem->priv.flags = flags;
 
diff --git a/drivers/gpu/nvgpu/common/linux/nvgpu_mem.c b/drivers/gpu/nvgpu/common/linux/nvgpu_mem.c
index e1f8a5987..8740ac3d1 100644
--- a/drivers/gpu/nvgpu/common/linux/nvgpu_mem.c
+++ b/drivers/gpu/nvgpu/common/linux/nvgpu_mem.c
@@ -25,6 +25,7 @@
 #include <nvgpu/vidmem.h>
 
 #include <nvgpu/linux/dma.h>
+#include <nvgpu/linux/vidmem.h>
 
 #include "os_linux.h"
 
@@ -355,7 +356,7 @@ u64 nvgpu_mem_get_addr(struct gk20a *g, struct nvgpu_mem *mem)
 	/*
 	 * Otherwise get the vidmem address.
 	 */
-	alloc = nvgpu_vidmem_get_page_alloc(mem->priv.sgt->sgl);
+	alloc = mem->vidmem_alloc;
 
 	/* This API should not be used with > 1 chunks */
 	WARN_ON(alloc->nr_chunks != 1);
diff --git a/drivers/gpu/nvgpu/common/mm/vidmem.c b/drivers/gpu/nvgpu/common/mm/vidmem.c
index 18f46c034..d1c5a2e80 100644
--- a/drivers/gpu/nvgpu/common/mm/vidmem.c
+++ b/drivers/gpu/nvgpu/common/mm/vidmem.c
@@ -195,7 +195,7 @@ int nvgpu_vidmem_clear(struct gk20a *g, struct nvgpu_mem *mem)
 	if (g->mm.vidmem.ce_ctx_id == (u32)~0)
 		return -EINVAL;
 
-	alloc = nvgpu_vidmem_get_page_alloc(mem->priv.sgt->sgl);
+	alloc = mem->vidmem_alloc;
 
 	nvgpu_sgt_for_each_sgl(sgl, &alloc->sgt) {
 		if (gk20a_last_fence)
diff --git a/drivers/gpu/nvgpu/common/pramin.c b/drivers/gpu/nvgpu/common/pramin.c
index b7bc74394..b6166f513 100644
--- a/drivers/gpu/nvgpu/common/pramin.c
+++ b/drivers/gpu/nvgpu/common/pramin.c
@@ -23,7 +23,6 @@
 #include <nvgpu/pramin.h>
 #include <nvgpu/page_allocator.h>
 #include <nvgpu/enabled.h>
-#include <nvgpu/vidmem.h>
 
 #include "gk20a/gk20a.h"
 
@@ -55,7 +54,7 @@ void nvgpu_pramin_access_batched(struct gk20a *g, struct nvgpu_mem *mem,
 	if (!g->regs && nvgpu_is_enabled(g, NVGPU_DRIVER_IS_DYING))
 		return;
 
-	alloc = nvgpu_vidmem_get_page_alloc(mem->priv.sgt->sgl);
+	alloc = mem->vidmem_alloc;
 	sgt = &alloc->sgt;
 
 	nvgpu_sgt_for_each_sgl(sgl, sgt) {
diff --git a/drivers/gpu/nvgpu/include/nvgpu/linux/vidmem.h b/drivers/gpu/nvgpu/include/nvgpu/linux/vidmem.h
index 76bbb05bc..ec02faec8 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/linux/vidmem.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/linux/vidmem.h
@@ -24,13 +24,18 @@ struct dma_buf;
 struct gk20a;
 
 #ifdef CONFIG_GK20A_VIDMEM
+
 struct gk20a *nvgpu_vidmem_buf_owner(struct dma_buf *dmabuf);
 int nvgpu_vidmem_export_linux(struct gk20a *g, size_t bytes);
 
+void nvgpu_vidmem_set_page_alloc(struct scatterlist *sgl, u64 addr);
+struct nvgpu_page_alloc *nvgpu_vidmem_get_page_alloc(struct scatterlist *sgl);
+
 int nvgpu_vidmem_buf_access_memory(struct gk20a *g, struct dma_buf *dmabuf,
 		void *buffer, u64 offset, u64 size, u32 cmd);
 
 #else /* !CONFIG_GK20A_VIDMEM */
+
 static inline struct gk20a *nvgpu_vidmem_buf_owner(struct dma_buf *dmabuf)
 {
 	return NULL;
@@ -41,16 +46,28 @@ static inline int nvgpu_vidmem_export_linux(struct gk20a *g, size_t bytes)
 	return -ENOSYS;
 }
 
+static inline void nvgpu_vidmem_set_page_alloc(struct scatterlist *sgl,
+					       u64 addr)
+{
+}
+
+static inline struct nvgpu_page_alloc *nvgpu_vidmem_get_page_alloc(
+	struct scatterlist *sgl)
+{
+	return NULL;
+}
+
 static inline int nvgpu_vidmem_buf_access_memory(struct gk20a *g,
-					     struct dma_buf *dmabuf,
-					     void *buffer, u64 offset,
-					     u64 size, u32 cmd)
+						 struct dma_buf *dmabuf,
+						 void *buffer, u64 offset,
+						 u64 size, u32 cmd)
 {
 	return -ENOSYS;
 }
 
 #endif
 
+
 struct nvgpu_vidmem_linux {
 	struct dma_buf	*dmabuf;
 	void		*dmabuf_priv;
diff --git a/drivers/gpu/nvgpu/include/nvgpu/nvgpu_mem.h b/drivers/gpu/nvgpu/include/nvgpu/nvgpu_mem.h
index 23a1bad79..537409a82 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/nvgpu_mem.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/nvgpu_mem.h
@@ -39,6 +39,7 @@ struct nvgpu_sgt;
 struct gk20a;
 struct nvgpu_allocator;
 struct nvgpu_gmmu_attrs;
+struct nvgpu_page_alloc;
 
 #define NVGPU_MEM_DMA_ERROR		(~0ULL)
 
@@ -162,6 +163,7 @@ struct nvgpu_mem {
 	/*
 	 * Fields only populated for vidmem allocations.
 	 */
+	struct nvgpu_page_alloc			*vidmem_alloc;
 	struct nvgpu_allocator			*allocator;
 	struct nvgpu_list_node			 clear_list_entry;
 
diff --git a/drivers/gpu/nvgpu/include/nvgpu/vidmem.h b/drivers/gpu/nvgpu/include/nvgpu/vidmem.h
index b89c710dd..9e9f83013 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/vidmem.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/vidmem.h
@@ -27,7 +27,6 @@
 #include <nvgpu/errno.h>
 #include <nvgpu/nvgpu_mem.h>
 
-struct scatterlist;
 struct work_struct;
 
 struct gk20a;
@@ -74,8 +73,6 @@ struct nvgpu_vidmem_buf *nvgpu_vidmem_user_alloc(struct gk20a *g, size_t bytes);
 
 void nvgpu_vidmem_buf_free(struct gk20a *g, struct nvgpu_vidmem_buf *buf);
 
-struct nvgpu_page_alloc *nvgpu_vidmem_get_page_alloc(struct scatterlist *sgl);
-void nvgpu_vidmem_set_page_alloc(struct scatterlist *sgl, u64 addr);
 bool nvgpu_addr_is_vidmem_page_alloc(u64 addr);
 int nvgpu_vidmem_get_space(struct gk20a *g, u64 *space);
 
@@ -93,16 +90,6 @@ int nvgpu_vidmem_clear(struct gk20a *g, struct nvgpu_mem *mem);
  * When VIDMEM support is not present this interface is used.
  */
 
-static inline struct nvgpu_page_alloc *
-nvgpu_vidmem_get_page_alloc(struct scatterlist *sgl)
-{
-	return NULL;
-}
-
-static inline void nvgpu_vidmem_set_page_alloc(struct scatterlist *sgl, u64 addr)
-{
-}
-
 static inline bool nvgpu_addr_is_vidmem_page_alloc(u64 addr)
 {
 	return false;