diff --git a/arch/um/kernel/main.c b/arch/um/kernel/main.c
index e59f58152678..1e1a87f1c510 100644
--- a/arch/um/kernel/main.c
+++ b/arch/um/kernel/main.c
@@ -69,7 +69,6 @@ static __init void do_uml_initcalls(void)
 
 static void last_ditch_exit(int sig)
 {
-        kmalloc_ok = 0;
 	signal(SIGINT, SIG_DFL);
 	signal(SIGTERM, SIG_DFL);
 	signal(SIGHUP, SIG_DFL);
diff --git a/arch/um/kernel/reboot.c b/arch/um/kernel/reboot.c
index 207f89d74908..fcec51da1d37 100644
--- a/arch/um/kernel/reboot.c
+++ b/arch/um/kernel/reboot.c
@@ -38,14 +38,14 @@ static void kill_off_processes(void)
 
 void uml_cleanup(void)
 {
-	kill_off_processes();
+        kmalloc_ok = 0;
 	do_uml_exitcalls();
+	kill_off_processes();
 }
 
 void machine_restart(char * __unused)
 {
-	do_uml_exitcalls();
-	kill_off_processes();
+        uml_cleanup();
 	CHOOSE_MODE(reboot_tt(), reboot_skas());
 }
 
@@ -53,8 +53,7 @@ EXPORT_SYMBOL(machine_restart);
 
 void machine_power_off(void)
 {
-	do_uml_exitcalls();
-	kill_off_processes();
+        uml_cleanup();
 	CHOOSE_MODE(halt_tt(), halt_skas());
 }