diff --git a/BaseBin/libjailbreak/src/jbserver.c b/BaseBin/libjailbreak/src/jbserver.c
index 07acc3c5d..e94138c0d 100644
--- a/BaseBin/libjailbreak/src/jbserver.c
+++ b/BaseBin/libjailbreak/src/jbserver.c
@@ -1,4 +1,5 @@
 #include "jbserver.h"
+#include "util.h"
 
 int jbserver_received_xpc_message(struct jbserver_impl *server, xpc_object_t xmsg)
 {
@@ -30,6 +31,8 @@ int jbserver_received_xpc_message(struct jbserver_impl *server, xpc_object_t xms
 	}
 	if (!action->handler) return -1;
 
+	thread_caffeinate_start();
+
 	int (*handler)(void *a1, void *a2, void *a3, void *a4, void *a5, void *a6, void *a7, void *a8) = action->handler;
 	void *args[8] = { NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL };
 	void *argsOut[8] = { NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL };
@@ -114,5 +117,7 @@ int jbserver_received_xpc_message(struct jbserver_impl *server, xpc_object_t xms
 	xpc_dictionary_set_int64(xreply, "result", result);
 	xpc_pipe_routine_reply(xreply);
 	xpc_release(xreply);
+
+	thread_caffeinate_stop();
 	return 0;
 }
\ No newline at end of file
diff --git a/BaseBin/libjailbreak/src/jbserver_boomerang.c b/BaseBin/libjailbreak/src/jbserver_boomerang.c
index 0f28939ad..d6b852bf3 100644
--- a/BaseBin/libjailbreak/src/jbserver_boomerang.c
+++ b/BaseBin/libjailbreak/src/jbserver_boomerang.c
@@ -22,14 +22,12 @@ int boomerang_get_physrw(audit_token_t *clientToken, bool singlePTE, uint64_t *s
 	int r = -1;
 	pid_t pid = audit_token_to_pid(*clientToken);
 
-	thread_caffeinate_start();
 	if (singlePTE) {
 		r = physrw_pte_handoff(pid, singlePTEAsidPtr);
 	}
 	else {
 		r = physrw_handoff(pid);
 	}
-	thread_caffeinate_stop();
 
 	return r;
 }
diff --git a/BaseBin/libjailbreak/src/kernel.c b/BaseBin/libjailbreak/src/kernel.c
index 466c06c7e..77dcf4226 100644
--- a/BaseBin/libjailbreak/src/kernel.c
+++ b/BaseBin/libjailbreak/src/kernel.c
@@ -194,10 +194,32 @@ kern_return_t pmap_enter_options_addr(uint64_t pmap, uint64_t pa, uint64_t va)
 	}
 }
 
-uint64_t pmap_remove(uint64_t pmap, uint64_t start, uint64_t end)
+uint64_t pmap_remove_options(uint64_t pmap, uint64_t start, uint64_t end)
 {
-	uint64_t kr = -1;
-	if (!is_kcall_available()) return kr;
-	kcall(&kr, ksymbol(pmap_remove_options), 4, (uint64_t[]){ pmap, start, end, 0x100 });
-	return kr;
+	uint64_t r = -1;
+	if (!is_kcall_available()) return r;
+	kcall(&r, ksymbol(pmap_remove_options), 4, (uint64_t[]){ pmap, start, end, 0x100 });
+	return r;
+}
+
+void pmap_remove(uint64_t pmap, uint64_t start, uint64_t end)
+{
+#ifdef __arm64e__
+	pmap_remove_options(pmap, start, end);
+#else
+    uint64_t remove_count = 0;
+    if (!pmap) {
+        return;
+    }
+    uint64_t va = start;
+    while (va < end) {
+        uint64_t l;
+        l = ((va + L2_BLOCK_SIZE) & ~L2_BLOCK_MASK);
+        if (l > end) {
+            l = end;
+        }
+        remove_count = pmap_remove_options(pmap, va, l);
+        va = remove_count;
+    }
+#endif
 }
\ No newline at end of file
diff --git a/BaseBin/libjailbreak/src/kernel.h b/BaseBin/libjailbreak/src/kernel.h
index 2f9c5590a..44144abda 100644
--- a/BaseBin/libjailbreak/src/kernel.h
+++ b/BaseBin/libjailbreak/src/kernel.h
@@ -58,6 +58,7 @@ void mac_label_set(uint64_t label, int slot, uint64_t value);
 int pmap_cs_allow_invalid(uint64_t pmap);
 int cs_allow_invalid(uint64_t proc, bool emulateFully);
 kern_return_t pmap_enter_options_addr(uint64_t pmap, uint64_t pa, uint64_t va);
-uint64_t pmap_remove(uint64_t pmap, uint64_t start, uint64_t end);
+uint64_t pmap_remove_options(uint64_t pmap, uint64_t start, uint64_t end);
+void pmap_remove(uint64_t pmap, uint64_t start, uint64_t end);
 
 #endif
\ No newline at end of file
diff --git a/BaseBin/libjailbreak/src/util.c b/BaseBin/libjailbreak/src/util.c
index 36fa8dff6..aeca1b103 100644
--- a/BaseBin/libjailbreak/src/util.c
+++ b/BaseBin/libjailbreak/src/util.c
@@ -263,16 +263,7 @@ int pmap_expand_range(uint64_t pmap, uint64_t vaStart, uint64_t size)
 				physwrite8(kvtophys(pmap + koffsetof(pmap, type)), 3);
 
 				// Remove mapping (table will stay cause nested is set)
-				if (vm_real_kernel_page_size == 0x1000) {
-					// 4k devices are fucked, don't ask me why
-					// If this isn't done, you get a panic with "%s: PTE range [%p, %p) in pmap %p crosses page table boundary"
-					for (uint64_t off = unmappedStart; off < (unmappedStart + unmappedSize); off += L2_BLOCK_SIZE) {
-						pmap_remove(pmap, off, off + vm_real_kernel_page_size);
-					}
-				}
-				else {
-					pmap_remove(pmap, unmappedStart, unmappedStart + unmappedSize);
-				}
+				pmap_remove(pmap, unmappedStart, unmappedStart + unmappedSize);
 
 				// Change type back
 				physwrite8(kvtophys(pmap + koffsetof(pmap, type)), 0);