diff --git a/Cargo.toml b/Cargo.toml
index 7214779f1a..1882ed0498 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -69,5 +69,5 @@ axdma = { path = "modules/axdma" }
 lto = true
 
 [patch.crates-io]
-page_table_multiarch = { git = "https://github.com/Mivik/page_table_multiarch.git", rev = "19ededd" }
-page_table_entry = { git = "https://github.com/Mivik/page_table_multiarch.git", rev = "19ededd" }
+page_table_multiarch = { git = "https://github.com/Ticonderoga2017/page_table_multiarch.git", rev = "0d0df14" }
+page_table_entry = { git = "https://github.com/Ticonderoga2017/page_table_multiarch.git", rev = "0d0df14" }
diff --git a/modules/axalloc/Cargo.toml b/modules/axalloc/Cargo.toml
index c0314aa1ef..139012efdc 100644
--- a/modules/axalloc/Cargo.toml
+++ b/modules/axalloc/Cargo.toml
@@ -21,6 +21,6 @@ page-alloc-4g = ["allocator/page-alloc-4g"] # Support up to 4G memory capacity
 log = "=0.4.21"
 cfg-if = "1.0"
 kspin = "0.1"
-memory_addr = "0.3"
+memory_addr = "0.4"
 axerrno = "0.1"
 allocator = { git = "https://github.com/arceos-org/allocator.git", tag ="v0.1.1", features = ["bitmap"] }
diff --git a/modules/axdma/Cargo.toml b/modules/axdma/Cargo.toml
index 12095f703a..1dd462595c 100644
--- a/modules/axdma/Cargo.toml
+++ b/modules/axdma/Cargo.toml
@@ -13,7 +13,7 @@ documentation = "https://arceos-org.github.io/arceos/axdma/index.html"
 [dependencies]
 log = "=0.4.21"
 kspin = "0.1"
-memory_addr = "0.3"
+memory_addr = "0.4"
 axerrno = "0.1"
 allocator = { git = "https://github.com/arceos-org/allocator.git", tag = "v0.1.1" }
 axalloc = { workspace = true }
diff --git a/modules/axdma/src/dma.rs b/modules/axdma/src/dma.rs
index 81ae42c97f..d3136f3b48 100644
--- a/modules/axdma/src/dma.rs
+++ b/modules/axdma/src/dma.rs
@@ -97,7 +97,7 @@ impl DmaAllocator {
         let expand_size = num_pages * PAGE_SIZE_4K;
         axmm::kernel_aspace()
             .lock()
-            .protect(vaddr, expand_size, flags, PageSize::Size4K)
+            .protect(vaddr, expand_size, flags)
             .map_err(|e| {
                 error!("change table flag fail: {e:?}");
                 AllocError::NoMemory
diff --git a/modules/axhal/Cargo.toml b/modules/axhal/Cargo.toml
index c6ec41cf15..6633c7744a 100644
--- a/modules/axhal/Cargo.toml
+++ b/modules/axhal/Cargo.toml
@@ -31,7 +31,7 @@ kspin = "0.1"
 int_ratio = "0.1.1"
 lazyinit = "0.2"
 percpu = "0.2"
-memory_addr = "0.3"
+memory_addr = "0.4"
 handler_table = "0.1"
 page_table_entry = "0.5"
 page_table_multiarch = "0.5"
diff --git a/modules/axmm/Cargo.toml b/modules/axmm/Cargo.toml
index 46a7d5d828..7c60d23874 100644
--- a/modules/axmm/Cargo.toml
+++ b/modules/axmm/Cargo.toml
@@ -17,7 +17,7 @@ axconfig = { workspace = true }
 log = "=0.4.21"
 axerrno = "0.1"
 lazyinit = "0.2"
-memory_addr = "0.3"
+memory_addr = "0.4"
 kspin = "0.1"
-memory_set = "0.3"
+memory_set = "0.4"
 page_table_multiarch = "0.5.3"
diff --git a/modules/axmm/src/aspace.rs b/modules/axmm/src/aspace.rs
index f25c7d7b53..95f41459bc 100644
--- a/modules/axmm/src/aspace.rs
+++ b/modules/axmm/src/aspace.rs
@@ -134,36 +134,7 @@ impl AddrSpace {
         limit: VirtAddrRange,
         align: PageSize,
     ) -> Option<VirtAddr> {
-        let mut last_end = hint.max(limit.start).align_up(align);
-        for area in self.areas.iter() {
-            if area.end() <= last_end {
-                last_end = last_end.max(area.end().align_up(align));
-            } else {
-                break;
-            }
-        }
-        for area in self.areas.iter() {
-            let area_start = area.start();
-            if area_start < last_end {
-                continue;
-            }
-            if last_end
-                .checked_add(size)
-                .is_some_and(|end| end <= area_start)
-            {
-                return Some(last_end);
-            }
-            last_end = area.end().align_up(align);
-        }
-
-        if last_end
-            .checked_add(size)
-            .is_some_and(|end| end <= limit.end)
-        {
-            Some(last_end)
-        } else {
-            None
-        }
+        self.areas.find_free_area(hint, size, limit, align.into())
     }
 
     /// Add a new linear mapping.
@@ -223,8 +194,8 @@ impl AddrSpace {
 
     /// Populates the area with physical frames, returning false if the area
     /// contains unmapped area.
-    pub fn populate_area(&mut self, mut start: VirtAddr, size: usize, align: PageSize) -> AxResult {
-        self.validate_region(start, size, align)?;
+    pub fn populate_area(&mut self, mut start: VirtAddr, size: usize) -> AxResult {
+        self.validate_region(start, size, PageSize::Size4K)?;
         let end = start + size;
 
         while let Some(area) = self.areas.find(start) {
@@ -246,7 +217,7 @@ impl AddrSpace {
                 }
             }
             start = area.end();
-            assert!(start.is_aligned(align));
+            assert!(start.is_aligned_4k());
             if start >= end {
                 break;
             }
@@ -387,15 +358,9 @@ impl AddrSpace {
     ///
     /// Returns an error if the address range is out of the address space or not
     /// aligned.
-    pub fn protect(
-        &mut self,
-        start: VirtAddr,
-        size: usize,
-        flags: MappingFlags,
-        align: PageSize,
-    ) -> AxResult {
+    pub fn protect(&mut self, start: VirtAddr, size: usize, flags: MappingFlags) -> AxResult {
         // Populate the area first, which also checks the address range for us.
-        self.populate_area(start, size, align)?;
+        self.populate_area(start, size)?;
 
         self.areas
             .protect(start, size, |_| Some(flags), &mut self.pt)
diff --git a/modules/axtask/Cargo.toml b/modules/axtask/Cargo.toml
index ede617f2f7..575a49e9a4 100644
--- a/modules/axtask/Cargo.toml
+++ b/modules/axtask/Cargo.toml
@@ -43,7 +43,7 @@ axconfig = { workspace = true, optional = true }
 percpu = { version = "0.2", optional = true }
 kspin = { version = "0.1", optional = true }
 lazyinit = { version = "0.2", optional = true }
-memory_addr = { version = "0.3", optional = true }
+memory_addr = { version = "0.4", optional = true }
 timer_list = { version = "0.1", optional = true }
 kernel_guard = { version = "0.1", optional = true }
 crate_interface = { version = "0.1", optional = true }