diff --git a/arch/arm64/kvm/hyp/pgtable.c b/arch/arm64/kvm/hyp/pgtable.c
index 4cf7359c3c4d..b1dfd9832e37 100644
--- a/arch/arm64/kvm/hyp/pgtable.c
+++ b/arch/arm64/kvm/hyp/pgtable.c
@@ -864,17 +864,10 @@ static bool stage2_unmap_defer_tlb_flush(struct kvm_pgtable *pgt)
 	return system_supports_tlb_range() && stage2_has_fwb(pgt);
 }
 
-static void stage2_unmap_put_pte(const struct kvm_pgtable_visit_ctx *ctx,
-				struct kvm_s2_mmu *mmu,
-				struct kvm_pgtable_mm_ops *mm_ops)
+static void stage2_unmap_clear_pte(const struct kvm_pgtable_visit_ctx *ctx,
+				   struct kvm_s2_mmu *mmu)
 {
 	struct kvm_pgtable *pgt = ctx->arg;
-
-	/*
-	 * Clear the existing PTE, and perform break-before-make if it was
-	 * valid. Depending on the system support, defer the TLB maintenance
-	 * for the same until the entire unmap walk is completed.
-	 */
 	if (kvm_pte_valid(ctx->old)) {
 		kvm_clear_pte(ctx->ptep);
 
@@ -886,7 +879,18 @@ static void stage2_unmap_put_pte(const struct kvm_pgtable_visit_ctx *ctx,
 				     ctx->level);
 		}
 	}
+}
 
+static void stage2_unmap_put_pte(const struct kvm_pgtable_visit_ctx *ctx,
+				 struct kvm_s2_mmu *mmu,
+				 struct kvm_pgtable_mm_ops *mm_ops)
+{
+	/*
+	 * Clear the existing PTE, and perform break-before-make if it was
+	 * valid. Depending on the system support, defer the TLB maintenance
+	 * for the same until the entire unmap walk is completed.
+	 */
+	stage2_unmap_clear_pte(ctx, mmu);
 	mm_ops->put_page(ctx->ptep);
 }