diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index b2058435b..3a57ea60a 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -1372,9 +1372,9 @@ impl MemoryManager {
         Ok(())
     }
 
-    #[cfg(target_arch = "aarch64")]
+    #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
     pub fn add_uefi_flash(&mut self) -> Result<(), Error> {
-        // On AArch64, the UEFI binary requires a flash device at address 0.
+        // The UEFI binary requires a flash device at address 0.
         // 4 MiB memory is mapped to simulate the flash.
         let uefi_mem_slot = self.allocate_memory_slot();
         let uefi_region = GuestRegionMmap::new(
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 1f3bcfeeb..04bd2d595 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -899,14 +899,6 @@ impl Vm {
             .allocate_address_space()
             .map_err(Error::MemoryManager)?;
 
-        // Add UEFI flash for aarch64
-        #[cfg(target_arch = "aarch64")]
-        memory_manager
-            .lock()
-            .unwrap()
-            .add_uefi_flash()
-            .map_err(Error::MemoryManager)?;
-
         // Load payload asynchronously
         let load_payload_handle = if snapshot.is_none() {
             Self::load_payload_async(
@@ -1405,7 +1397,11 @@ impl Vm {
         mut firmware: &File,
         memory_manager: Arc<Mutex<MemoryManager>>,
     ) -> Result<EntryPoint> {
-        let uefi_flash = memory_manager.lock().as_ref().unwrap().uefi_flash();
+        let mut memory_manager = memory_manager.lock().unwrap();
+        memory_manager
+            .add_uefi_flash()
+            .map_err(Error::MemoryManager)?;
+        let uefi_flash = memory_manager.uefi_flash();
         let mem = uefi_flash.memory();
         arch::uefi::load_uefi(mem.deref(), arch::layout::UEFI_START, &mut firmware)
             .map_err(Error::UefiLoad)?;