diff --git a/arch/um/kernel/main.c b/arch/um/kernel/main.c
index e59f581526782adca0763540a98bc80505a7a79a..1e1a87f1c510cbc7a2d3a4ea472bd6395f6e4bb2 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 207f89d749080a596954a6ac7adbf2fd68fd88cb..fcec51da1d3771abe84944ca1138853417df6bec 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());
 }