diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index f34608013..e5367c246 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -773,7 +773,7 @@ impl MemoryManager {
         }
     }
 
-    fn allocate_address_space(&mut self) -> Result<(), Error> {
+    pub fn allocate_address_space(&mut self) -> Result<(), Error> {
         let mut list = Vec::new();
 
         for (zone_id, memory_zone) in self.memory_zones.iter() {
@@ -1135,10 +1135,17 @@ impl MemoryManager {
             thp: config.thp,
         };
 
-        memory_manager.allocate_address_space()?;
-
         #[cfg(target_arch = "aarch64")]
-        memory_manager.add_uefi_flash()?;
+        {
+            // For Aarch64 we cannot lazily allocate the address space like we
+            // do for x86, because while restoring a VM from snapshot we would
+            // need the address space to be allocated to properly restore VGIC.
+            // And the restore of VGIC happens before we attempt to run the vCPUs
+            // for the first time, thus we need to allocate the address space
+            // beforehand.
+            memory_manager.allocate_address_space()?;
+            memory_manager.add_uefi_flash()?;
+        }
 
         #[cfg(target_arch = "x86_64")]
         if let Some(sgx_epc_config) = sgx_epc_config {
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index a03758298..214173093 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -2116,6 +2116,18 @@ impl Vm {
             self.vm.tdx_finalize().map_err(Error::FinalizeTdx)?;
         }
 
+        #[cfg(target_arch = "x86_64")]
+        // Note: For x86, always call this function before invoking start boot vcpus.
+        // Otherwise guest would fail to boot because we haven't created the
+        // userspace mappings to update the hypervisor about the memory mappings.
+        // These mappings must be created before we start the vCPU threads for
+        // the very first time.
+        self.memory_manager
+            .lock()
+            .unwrap()
+            .allocate_address_space()
+            .map_err(Error::MemoryManager)?;
+
         self.cpu_manager
             .lock()
             .unwrap()
@@ -2131,6 +2143,18 @@ impl Vm {
     pub fn restore(&mut self) -> Result<()> {
         event!("vm", "restoring");
 
+        #[cfg(target_arch = "x86_64")]
+        // Note: For x86, always call this function before invoking start boot vcpus.
+        // Otherwise guest would fail to boot because we haven't created the
+        // userspace mappings to update the hypervisor about the memory mappings.
+        // These mappings must be created before we start the vCPU threads for
+        // the very first time for the restored VM.
+        self.memory_manager
+            .lock()
+            .unwrap()
+            .allocate_address_space()
+            .map_err(Error::MemoryManager)?;
+
         // Now we can start all vCPUs from here.
         self.cpu_manager
             .lock()