diff --git a/drivers/gpu/nvgpu/Makefile b/drivers/gpu/nvgpu/Makefile
index 9ac1b9029..e88bae156 100644
--- a/drivers/gpu/nvgpu/Makefile
+++ b/drivers/gpu/nvgpu/Makefile
@@ -478,6 +478,7 @@ nvgpu-y += \
 	gv100/hal_gv100.o \
 	gv100/gsp_gv100.o \
 	gv100/clk_gv100.o \
+	tu104/mm_tu104.o \
 	tu104/hal_tu104.o \
 	tu104/fifo_tu104.o \
 	tu104/gr_tu104.o \
diff --git a/drivers/gpu/nvgpu/Makefile.sources b/drivers/gpu/nvgpu/Makefile.sources
index 45b8a1fe9..660851ac5 100644
--- a/drivers/gpu/nvgpu/Makefile.sources
+++ b/drivers/gpu/nvgpu/Makefile.sources
@@ -291,6 +291,7 @@ srcs += common/sim.c \
 	tu104/fbpa_tu104.c	\
 	tu104/fifo_tu104.c	\
 	tu104/gr_tu104.c	\
+	tu104/mm_tu104.c	\
 	tu104/hal_tu104.c	\
 	tu104/sec2_tu104.c	\
 	tu104/func_tu104.c	\
diff --git a/drivers/gpu/nvgpu/tu104/hal_tu104.c b/drivers/gpu/nvgpu/tu104/hal_tu104.c
index 6bd52564d..d16dc6e80 100644
--- a/drivers/gpu/nvgpu/tu104/hal_tu104.c
+++ b/drivers/gpu/nvgpu/tu104/hal_tu104.c
@@ -137,8 +137,8 @@
 #include "gv100/bios_gv100.h"
 #include "gv100/fifo_gv100.h"
 #include "gv100/gr_gv100.h"
-#include "gv100/mm_gv100.h"
 
+#include "tu104/mm_tu104.h"
 #include "tu104/fifo_tu104.h"
 #include "tu104/gr_tu104.h"
 #include "tu104/bios_tu104.h"
@@ -929,7 +929,7 @@ static const struct gpu_ops tu104_ops = {
 		.remove_bar2_vm = gp10b_remove_bar2_vm,
 		.fault_info_mem_destroy = gv11b_mm_fault_info_mem_destroy,
 		.mmu_fault_disable_hw = gv11b_mm_mmu_fault_disable_hw,
-		.get_flush_retries = gv100_mm_get_flush_retries,
+		.get_flush_retries = tu104_mm_get_flush_retries,
 		.bar1_map_userd = NULL,
 	},
 	.pramin = {
diff --git a/drivers/gpu/nvgpu/tu104/mm_tu104.c b/drivers/gpu/nvgpu/tu104/mm_tu104.c
new file mode 100644
index 000000000..a6f04aeef
--- /dev/null
+++ b/drivers/gpu/nvgpu/tu104/mm_tu104.c
@@ -0,0 +1,47 @@
+/*
+ * TU104 memory management
+ *
+ * Copyright (c) 2019, 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"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include <nvgpu/mm.h>
+
+#include "tu104/mm_tu104.h"
+
+u32 tu104_mm_get_flush_retries(struct gk20a *g, enum nvgpu_flush_op op)
+{
+	u32 retries;
+
+	switch (op) {
+	/* TU104 has a large FB so it needs larger timeouts */
+	case NVGPU_FLUSH_FB:
+		retries = 2500;
+		break;
+	case NVGPU_FLUSH_L2_FLUSH:
+		retries = 2000;
+		break;
+	default:
+		retries = 200; /* Default retry timer */
+		break;
+	}
+
+	return retries;
+}
diff --git a/drivers/gpu/nvgpu/tu104/mm_tu104.h b/drivers/gpu/nvgpu/tu104/mm_tu104.h
new file mode 100644
index 000000000..975ec1a71
--- /dev/null
+++ b/drivers/gpu/nvgpu/tu104/mm_tu104.h
@@ -0,0 +1,33 @@
+/*
+ * TU104 memory management
+ *
+ * Copyright (c) 2019, 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"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef MM_TU104_H
+#define MM_TU104_H
+
+struct gk20a;
+enum nvgpu_flush_op;
+
+u32 tu104_mm_get_flush_retries(struct gk20a *g, enum nvgpu_flush_op op);
+
+#endif