diff --git a/arch/arm64/kvm/hyp/nvhe/tlb.c b/arch/arm64/kvm/hyp/nvhe/tlb.c
index 3f5601176fab..d296d617f589 100644
--- a/arch/arm64/kvm/hyp/nvhe/tlb.c
+++ b/arch/arm64/kvm/hyp/nvhe/tlb.c
@@ -11,62 +11,26 @@
 #include <nvhe/mem_protect.h>
 
 struct tlb_inv_context {
-	struct kvm_s2_mmu	*mmu;
-	u64			tcr;
-	u64			sctlr;
+	u64		tcr;
 };
 
-static void enter_vmid_context(struct kvm_s2_mmu *mmu,
-			       struct tlb_inv_context *cxt)
+static void __tlb_switch_to_guest(struct kvm_s2_mmu *mmu,
+				  struct tlb_inv_context *cxt)
 {
-	struct kvm_s2_mmu *host_mmu = &host_kvm.arch.mmu;
-	struct kvm_cpu_context *host_ctxt;
-	struct kvm_vcpu *vcpu;
-
-	host_ctxt = &this_cpu_ptr(&kvm_host_data)->host_ctxt;
-	vcpu = host_ctxt->__hyp_running_vcpu;
-	cxt->mmu = NULL;
-
-	/*
-	 * If we're already in the desired context, then there's nothing
-	 * to do.
-	 */
-	if (vcpu) {
-		if (mmu == vcpu->arch.hw_mmu || WARN_ON(mmu != host_mmu))
-			return;
-	} else if (mmu == host_mmu) {
-		return;
-	}
-
-	cxt->mmu = mmu;
 	if (cpus_have_final_cap(ARM64_WORKAROUND_SPECULATIVE_AT)) {
 		u64 val;
 
 		/*
 		 * For CPUs that are affected by ARM 1319367, we need to
-		 * avoid a Stage-1 walk with the old VMID while we have
-		 * the new VMID set in the VTTBR in order to invalidate TLBs.
-		 * We're guaranteed that the host S1 MMU is enabled, so
-		 * we can simply set the EPD bits to avoid any further
-		 * TLB fill. For guests, we ensure that the S1 MMU is
-		 * temporarily enabled in the next context.
+		 * avoid a host Stage-1 walk while we have the guest's
+		 * VMID set in the VTTBR in order to invalidate TLBs.
+		 * We're guaranteed that the S1 MMU is enabled, so we can
+		 * simply set the EPD bits to avoid any further TLB fill.
 		 */
 		val = cxt->tcr = read_sysreg_el1(SYS_TCR);
 		val |= TCR_EPD1_MASK | TCR_EPD0_MASK;
 		write_sysreg_el1(val, SYS_TCR);
 		isb();
-
-		if (vcpu) {
-			val = cxt->sctlr = read_sysreg_el1(SYS_SCTLR);
-			if (!(val & SCTLR_ELx_M)) {
-				val |= SCTLR_ELx_M;
-				write_sysreg_el1(val, SYS_SCTLR);
-				isb();
-			}
-		} else {
-			/* The host S1 MMU is always enabled. */
-			cxt->sctlr = SCTLR_ELx_M;
-		}
 	}
 
 	/*
@@ -75,44 +39,20 @@ static void enter_vmid_context(struct kvm_s2_mmu *mmu,
 	 * ensuring that we always have an ISB, but not two ISBs back
 	 * to back.
 	 */
-	if (vcpu)
-		__load_host_stage2();
-	else
-		__load_stage2(mmu, kern_hyp_va(mmu->arch));
-
+	__load_stage2(mmu, kern_hyp_va(mmu->arch));
 	asm(ALTERNATIVE("isb", "nop", ARM64_WORKAROUND_SPECULATIVE_AT));
 }
 
-static void exit_vmid_context(struct tlb_inv_context *cxt)
+static void __tlb_switch_to_host(struct tlb_inv_context *cxt)
 {
-	struct kvm_s2_mmu *mmu = cxt->mmu;
-	struct kvm_cpu_context *host_ctxt;
-	struct kvm_vcpu *vcpu;
-
-	host_ctxt = &this_cpu_ptr(&kvm_host_data)->host_ctxt;
-	vcpu = host_ctxt->__hyp_running_vcpu;
-
-	if (!mmu)
-		return;
-
-	if (vcpu)
-		__load_stage2(mmu, kern_hyp_va(mmu->arch));
-	else
-		__load_host_stage2();
+	__load_host_stage2();
 
 	if (cpus_have_final_cap(ARM64_WORKAROUND_SPECULATIVE_AT)) {
-		/* Ensure write of the old VMID */
+		/* Ensure write of the host VMID */
 		isb();
-
-		if (!(cxt->sctlr & SCTLR_ELx_M)) {
-			write_sysreg_el1(cxt->sctlr, SYS_SCTLR);
-			isb();
-		}
-
+		/* Restore the host's TCR_EL1 */
 		write_sysreg_el1(cxt->tcr, SYS_TCR);
 	}
-
-	cxt->mmu = NULL;
 }
 
 void __kvm_tlb_flush_vmid_ipa(struct kvm_s2_mmu *mmu,
@@ -123,7 +63,7 @@ void __kvm_tlb_flush_vmid_ipa(struct kvm_s2_mmu *mmu,
 	dsb(ishst);
 
 	/* Switch to requested VMID */
-	enter_vmid_context(mmu, &cxt);
+	__tlb_switch_to_guest(mmu, &cxt);
 
 	/*
 	 * We could do so much better if we had the VA as well.
@@ -166,7 +106,7 @@ void __kvm_tlb_flush_vmid_ipa(struct kvm_s2_mmu *mmu,
 	if (icache_is_vpipt())
 		icache_inval_all_pou();
 
-	exit_vmid_context(&cxt);
+	__tlb_switch_to_host(&cxt);
 }
 
 void __kvm_tlb_flush_vmid(struct kvm_s2_mmu *mmu)
@@ -176,13 +116,13 @@ void __kvm_tlb_flush_vmid(struct kvm_s2_mmu *mmu)
 	dsb(ishst);
 
 	/* Switch to requested VMID */
-	enter_vmid_context(mmu, &cxt);
+	__tlb_switch_to_guest(mmu, &cxt);
 
 	__tlbi(vmalls12e1is);
 	dsb(ish);
 	isb();
 
-	exit_vmid_context(&cxt);
+	__tlb_switch_to_host(&cxt);
 }
 
 void __kvm_flush_cpu_context(struct kvm_s2_mmu *mmu)
@@ -190,14 +130,14 @@ void __kvm_flush_cpu_context(struct kvm_s2_mmu *mmu)
 	struct tlb_inv_context cxt;
 
 	/* Switch to requested VMID */
-	enter_vmid_context(mmu, &cxt);
+	__tlb_switch_to_guest(mmu, &cxt);
 
 	__tlbi(vmalle1);
 	asm volatile("ic iallu");
 	dsb(nsh);
 	isb();
 
-	exit_vmid_context(&cxt);
+	__tlb_switch_to_host(&cxt);
 }
 
 void __kvm_flush_vm_context(void)