diff --git a/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h b/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
index ae14492be415..db48cfb89e00 100644
--- a/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
+++ b/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
@@ -80,7 +80,9 @@ u64 __pkvm_ptdump_walk_range(pkvm_handle_t handle, struct pkvm_ptdump_log_hdr *l
 int hyp_check_range_owned(u64 addr, u64 size);
 int __pkvm_install_guest_mmio(struct pkvm_hyp_vcpu *hyp_vcpu, u64 pfn, u64 gfn);
 
-int __pkvm_guest_get_valid_phys_page(struct pkvm_hyp_vm *vm, u64 *phys, u64 ipa);
+int pkvm_get_guest_pa_request(struct pkvm_hyp_vcpu *hyp_vcpu, u64 ipa,
+			      size_t ipa_size_request, u64 *out_pa, s8 *out_level,
+			      u64 *exit_code);
 
 bool addr_is_memory(phys_addr_t phys);
 int host_stage2_idmap_locked(phys_addr_t addr, u64 size,
diff --git a/arch/arm64/kvm/hyp/nvhe/device/device.c b/arch/arm64/kvm/hyp/nvhe/device/device.c
index dfcc6ec221c0..2a2efa2a2ab4 100644
--- a/arch/arm64/kvm/hyp/nvhe/device/device.c
+++ b/arch/arm64/kvm/hyp/nvhe/device/device.c
@@ -267,29 +267,6 @@ out_ret:
 	return ret;
 }
 
-static int pkvm_get_device_pa(struct pkvm_hyp_vcpu *hyp_vcpu, u64 ipa, u64 *pa, u64 *exit_code)
-{
-	struct kvm_hyp_req *req;
-	int ret;
-	struct pkvm_hyp_vm *vm = pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu);
-
-	ret = __pkvm_guest_get_valid_phys_page(vm, pa, ipa);
-	if (ret == -ENOENT) {
-		/* Page not mapped, create a request*/
-		req = pkvm_hyp_req_reserve(hyp_vcpu, KVM_HYP_REQ_TYPE_MAP);
-		if (!req)
-			return -ENOMEM;
-
-		req->map.guest_ipa = ipa;
-		req->map.size = PAGE_SIZE;
-		*exit_code = ARM_EXCEPTION_HYP_REQ;
-		/* Repeat next time. */
-		write_sysreg_el2(read_sysreg_el2(SYS_ELR) - 4, SYS_ELR);
-	}
-
-	return ret;
-}
-
 bool pkvm_device_request_mmio(struct pkvm_hyp_vcpu *hyp_vcpu, u64 *exit_code)
 {
 	int i, j, ret;
@@ -297,17 +274,22 @@ bool pkvm_device_request_mmio(struct pkvm_hyp_vcpu *hyp_vcpu, u64 *exit_code)
 	struct pkvm_hyp_vm *vm = pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu);
 	u64 ipa = smccc_get_arg1(vcpu);
 	u64 token;
+	s8 level;
 
 	/* arg2 and arg3 reserved for future use. */
 	if (smccc_get_arg2(vcpu) || smccc_get_arg3(vcpu) || !PAGE_ALIGNED(ipa))
 		goto out_inval;
 
-	ret = pkvm_get_device_pa(hyp_vcpu, ipa, &token, exit_code);
+	ret = pkvm_get_guest_pa_request(hyp_vcpu, ipa, PAGE_SIZE,
+					&token, &level, exit_code);
 	if (ret == -ENOENT)
 		return false;
 	else if (ret)
 		goto out_inval;
 
+	/* It's expected the address is mapped as page for MMIO */
+	WARN_ON(level != KVM_PGTABLE_LAST_LEVEL);
+
 	hyp_spin_lock(&device_spinlock);
 	for (i = 0 ; i < registered_devices_nr ; ++i) {
 		struct pkvm_device *dev = &registered_devices[i];
diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index e5470cc98d68..47604640c32c 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -2105,18 +2105,6 @@ static int guest_get_valid_pte(struct pkvm_hyp_vm *vm, u64 *phys, u64 ipa, u8 or
 	return 0;
 }
 
-int __pkvm_guest_get_valid_phys_page(struct pkvm_hyp_vm *vm, u64 *phys, u64 ipa)
-{
-	kvm_pte_t pte;
-	int ret;
-
-	guest_lock_component(vm);
-	ret = guest_get_valid_pte(vm, phys, ipa, 0, &pte);
-	guest_unlock_component(vm);
-
-	return ret;
-}
-
 int __pkvm_host_relax_perms_guest(u64 gfn, struct pkvm_hyp_vcpu *vcpu, enum kvm_pgtable_prot prot)
 {
 	struct pkvm_hyp_vm *vm = pkvm_hyp_vcpu_to_hyp_vm(vcpu);
@@ -2743,6 +2731,42 @@ teardown:
 	return ret;
 }
 
+/* Return PA for an owned guest IPA or request it, and repeat the guest HVC */
+int pkvm_get_guest_pa_request(struct pkvm_hyp_vcpu *hyp_vcpu, u64 ipa,
+			      size_t ipa_size_request, u64 *out_pa, s8 *out_level,
+			      u64 *exit_code)
+{
+	struct kvm_hyp_req *req;
+	kvm_pte_t pte;
+	enum pkvm_page_state state;
+	struct pkvm_hyp_vm *vm = pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu);
+
+	guest_lock_component(vm);
+	WARN_ON(kvm_pgtable_get_leaf(&vm->pgt, ipa, &pte, out_level));
+	guest_unlock_component(vm);
+	if (!kvm_pte_valid(pte)) {
+		/* Page not mapped, create a request*/
+		req = pkvm_hyp_req_reserve(hyp_vcpu, KVM_HYP_REQ_TYPE_MAP);
+		if (!req)
+			return -ENOMEM;
+
+		req->map.guest_ipa = ipa;
+		req->map.size = ipa_size_request;
+		*exit_code = ARM_EXCEPTION_HYP_REQ;
+		/* Repeat next time. */
+		write_sysreg_el2(read_sysreg_el2(SYS_ELR) - 4, SYS_ELR);
+		return -ENOENT;
+	}
+
+	state = pkvm_getstate(kvm_pgtable_stage2_pte_prot(pte));
+	if (state != PKVM_PAGE_OWNED)
+		return -EPERM;
+
+	*out_pa = kvm_pte_to_phys(pte);
+	*out_pa |= (ipa & kvm_granule_size(*out_level) - 1) & PAGE_MASK;
+	return 0;
+}
+
 #ifdef CONFIG_PKVM_SELFTESTS
 struct pkvm_expected_state {
 	enum pkvm_page_state host;