diff --git a/arch/arm64/kvm/hyp/nvhe/pkvm.c b/arch/arm64/kvm/hyp/nvhe/pkvm.c
index 97281f8c3b41..4316111f8334 100644
--- a/arch/arm64/kvm/hyp/nvhe/pkvm.c
+++ b/arch/arm64/kvm/hyp/nvhe/pkvm.c
@@ -560,13 +560,13 @@ static void init_pkvm_hyp_vm(struct kvm *host_kvm, struct pkvm_hyp_vm *hyp_vm,
 	hyp_vm->kvm.created_vcpus = nr_vcpus;
 	hyp_vm->kvm.arch.mmu.vtcr = host_mmu.arch.mmu.vtcr;
 	hyp_vm->kvm.arch.pkvm.enabled = READ_ONCE(host_kvm->arch.pkvm.enabled);
+	hyp_vm->kvm.arch.flags = 0;
 
 	if (hyp_vm->kvm.arch.pkvm.enabled)
 		pvmfw_load_addr = READ_ONCE(host_kvm->arch.pkvm.pvmfw_load_addr);
 	hyp_vm->kvm.arch.pkvm.pvmfw_load_addr = pvmfw_load_addr;
 
 	hyp_vm->kvm.arch.mmu.last_vcpu_ran = (int __percpu *)last_ran;
-	hyp_vm->kvm.arch.flags = 0;
 	memset(last_ran, -1, pkvm_get_last_ran_size());
 	pkvm_init_features_from_host(hyp_vm, host_kvm);
 	hyp_spin_lock_init(&hyp_vm->vcpus_lock);
@@ -656,8 +656,10 @@ static int init_pkvm_hyp_vcpu(struct pkvm_hyp_vcpu *hyp_vcpu,
 	hyp_vcpu->vcpu.arch.debug_ptr = &host_vcpu->arch.vcpu_debug_state;
 	hyp_vcpu->vcpu.arch.hyp_reqs->type = KVM_HYP_LAST_REQ;
 
-	kvm_init_pvm_id_regs(&hyp_vcpu->vcpu);
-	kvm_reset_pvm_sys_regs(&hyp_vcpu->vcpu);
+	if (pkvm_hyp_vcpu_is_protected(hyp_vcpu)) {
+		kvm_init_pvm_id_regs(&hyp_vcpu->vcpu);
+		kvm_reset_pvm_sys_regs(&hyp_vcpu->vcpu);
+	}
 
 	ret = pkvm_vcpu_init_traps(hyp_vcpu);
 	if (ret)
diff --git a/arch/arm64/kvm/hyp/nvhe/sys_regs.c b/arch/arm64/kvm/hyp/nvhe/sys_regs.c
index de5c34f1b0d5..ebf7e28c1be9 100644
--- a/arch/arm64/kvm/hyp/nvhe/sys_regs.c
+++ b/arch/arm64/kvm/hyp/nvhe/sys_regs.c
@@ -302,9 +302,10 @@ static u64 read_id_reg(const struct kvm_vcpu *vcpu,
 	struct kvm *kvm = vcpu->kvm;
 	u32 reg = reg_to_encoding(r);
 
-	BUG_ON(!test_bit(KVM_ARCH_FLAG_ID_REGS_INITIALIZED, &kvm->arch.flags));
+	if (WARN_ON_ONCE(!test_bit(KVM_ARCH_FLAG_ID_REGS_INITIALIZED, &kvm->arch.flags)))
+		return 0;
 
-	if (reg >= sys_reg(3, 0, 0, 4, 0) && reg <= sys_reg(3, 0, 0, 7, 7))
+	if (reg >= sys_reg(3, 0, 0, 1, 0) && reg <= sys_reg(3, 0, 0, 7, 7))
 		return kvm->arch.id_regs[IDREG_IDX(reg)];
 
 	return 0;
@@ -618,37 +619,23 @@ static const struct sys_reg_desc_reset pvm_sys_reg_reset_vals[] = {
  */
 void kvm_init_pvm_id_regs(struct kvm_vcpu *vcpu)
 {
-	/* List of feature registers to reset for protected VMs. */
-	const u32 pvm_feat_id_regs[] = {
-		SYS_ID_AA64PFR0_EL1,
-		SYS_ID_AA64PFR1_EL1,
-		SYS_ID_AA64ISAR0_EL1,
-		SYS_ID_AA64ISAR1_EL1,
-		SYS_ID_AA64ISAR2_EL1,
-		SYS_ID_AA64ZFR0_EL1,
-		SYS_ID_AA64MMFR0_EL1,
-		SYS_ID_AA64MMFR1_EL1,
-		SYS_ID_AA64MMFR2_EL1,
-		SYS_ID_AA64MMFR4_EL1,
-		SYS_ID_AA64DFR0_EL1,
-	};
 	struct kvm *kvm = vcpu->kvm;
-	int i;
+	struct kvm_arch *ka = &kvm->arch;
+	u32 r;
 
 	if (test_bit(KVM_ARCH_FLAG_ID_REGS_INITIALIZED, &kvm->arch.flags))
 		return;
 
-	for (i = 0; i < ARRAY_SIZE(pvm_feat_id_regs); i++) {
-		struct kvm_arch *ka = &kvm->arch;
-		u32 reg = pvm_feat_id_regs[i];
-
-		ka->id_regs[IDREG_IDX(reg)] = pvm_calc_id_reg(vcpu, reg);
-	}
+	/*
+	 * Initialize only AArch64 id registers since AArch32 isn't supported
+	 * for protected VMs.
+	 */
+	for (r = sys_reg(3, 0, 0, 1, 0); r <= sys_reg(3, 0, 0, 7, 7); r += sys_reg(0, 0, 0, 0, 1))
+		ka->id_regs[IDREG_IDX(r)] = pvm_calc_id_reg(vcpu, r);
 
 	set_bit(KVM_ARCH_FLAG_ID_REGS_INITIALIZED, &kvm->arch.flags);
 }
 
-
 /*
  * Sets system registers to reset value
  *