diff --git a/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h b/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
index fa83c2fa7b9d..e6f93ac38b6a 100644
--- a/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
+++ b/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
@@ -53,8 +53,7 @@ extern struct host_kvm host_kvm;
 
 typedef u32 pkvm_id;
 static const pkvm_id pkvm_host_id	= 0;
-static const pkvm_id pkvm_hyp_id	= 1;
-static const pkvm_id pkvm_guest_id	= 2;
+static const pkvm_id pkvm_hyp_id	= (1 << 16);
 
 extern unsigned long hyp_nr_cpus;
 
diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index 062258f91fa3..da75c35ef4e8 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -29,6 +29,12 @@ struct host_kvm host_kvm;
 
 static struct hyp_pool host_s2_pool;
 
+static pkvm_id pkvm_guest_id(struct kvm_vcpu *vcpu)
+{
+	return vcpu->arch.hw_mmu->vmid.vmid;
+
+}
+
 static DEFINE_PER_CPU(struct kvm_shadow_vm *, __current_vm);
 #define current_vm (*this_cpu_ptr(&__current_vm))
 
@@ -687,6 +693,8 @@ static pkvm_id completer_owner_id(const struct pkvm_mem_transition *tx)
 		return pkvm_host_id;
 	case PKVM_ID_HYP:
 		return pkvm_hyp_id;
+	case PKVM_ID_GUEST:
+		return pkvm_guest_id(tx->completer.guest.vcpu);
 	default:
 		WARN_ON(1);
 		return -1;