diff --git a/arch/x86/include/asm/tlbflush.h b/arch/x86/include/asm/tlbflush.h
index 85315d1d2d70..7d04aa2a5f86 100644
--- a/arch/x86/include/asm/tlbflush.h
+++ b/arch/x86/include/asm/tlbflush.h
@@ -296,6 +296,14 @@ unsigned long build_cr3_pcid(pgd_t *pgd, u16 pcid, bool noflush);
 u16 kern_pcid(u16 asid);
 u16 asi_pcid(struct asi *asi, u16 asid);
 
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+
+void __asi_prepare_tlb_flush(struct asi *asi, u64 *new_tlb_gen);
+void __asi_flush_tlb_range(u64 mm_context_id, u16 pcid_index, u64 new_tlb_gen,
+			   size_t start, size_t end, const cpumask_t *cpu_mask);
+
+#endif /* CONFIG_ADDRESS_SPACE_ISOLATION */
+
 #endif /* !MODULE */
 
 #endif /* _ASM_X86_TLBFLUSH_H */
diff --git a/arch/x86/mm/tlb.c b/arch/x86/mm/tlb.c
index 2a442335501f..fcd2c8e92f83 100644
--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -1302,21 +1302,10 @@ static bool is_asi_active_on_cpu(int cpu, void *info)
 	return per_cpu(asi_cpu_state.curr_asi, cpu);
 }
 
-void asi_flush_tlb_range(struct asi *asi, void *addr, size_t len)
+void __asi_prepare_tlb_flush(struct asi *asi, u64 *new_tlb_gen)
 {
-	size_t start = (size_t)addr;
-	size_t end = start + len;
-	struct flush_tlb_info *info;
-	u64 mm_context_id;
-	const cpumask_t *cpu_mask;
-	u64 new_tlb_gen = 0;
-
-	if (!static_cpu_has(X86_FEATURE_ASI))
-		return;
-
 	if (static_cpu_has(X86_FEATURE_PCID)) {
-		new_tlb_gen = atomic64_inc_return(asi->tlb_gen);
-
+		*new_tlb_gen = atomic64_inc_return(asi->tlb_gen);
 		/*
 		 * The increment of tlb_gen must happen before the curr_asi
 		 * reads in is_asi_active_on_cpu(). That ensures that if another
@@ -1326,8 +1315,35 @@ void asi_flush_tlb_range(struct asi *asi, void *addr, size_t len)
 		 */
 		smp_mb__after_atomic();
 	}
+}
+
+void __asi_flush_tlb_range(u64 mm_context_id, u16 pcid_index, u64 new_tlb_gen,
+			   size_t start, size_t end, const cpumask_t *cpu_mask)
+{
+	struct flush_tlb_info *info;
 
 	preempt_disable();
+	info = get_flush_tlb_info(NULL, start, end, 0, false, new_tlb_gen,
+				  mm_context_id, pcid_index);
+
+	on_each_cpu_cond_mask(is_asi_active_on_cpu, do_asi_tlb_flush, info,
+			      true, cpu_mask);
+	put_flush_tlb_info();
+	preempt_enable();
+}
+
+void asi_flush_tlb_range(struct asi *asi, void *addr, size_t len)
+{
+	size_t start = (size_t)addr;
+	size_t end = start + len;
+	u64 mm_context_id;
+	u64 new_tlb_gen = 0;
+	const cpumask_t *cpu_mask;
+
+	if (!static_cpu_has(X86_FEATURE_ASI))
+		return;
+
+	__asi_prepare_tlb_flush(asi, &new_tlb_gen);
 
 	if (asi == ASI_GLOBAL_NONSENSITIVE) {
 		mm_context_id = U64_MAX;
@@ -1337,14 +1353,8 @@ void asi_flush_tlb_range(struct asi *asi, void *addr, size_t len)
 		cpu_mask = mm_cpumask(asi->mm);
 	}
 
-	info = get_flush_tlb_info(NULL, start, end, 0, false, new_tlb_gen,
-				  mm_context_id, asi->pcid_index);
-
-	on_each_cpu_cond_mask(is_asi_active_on_cpu, do_asi_tlb_flush, info,
-			      true, cpu_mask);
-
-	put_flush_tlb_info();
-	preempt_enable();
+	__asi_flush_tlb_range(mm_context_id, asi->pcid_index, new_tlb_gen,
+			      start, end, cpu_mask);
 }
 
 #endif
diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 56511adc263e..7d38229ca85c 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -193,21 +193,33 @@ struct page {
 		/** @rcu_head: You can use this to free a page by RCU. */
 		struct rcu_head rcu_head;
 
-#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+#if defined(CONFIG_ADDRESS_SPACE_ISOLATION) && !defined(BUILD_VDSO32)
 		struct {
 			/* Links the pages_to_free_async list */
 			struct llist_node async_free_node;
 
 			unsigned long _asi_pad_1;
-			unsigned long _asi_pad_2;
+			u64 asi_tlb_gen;
 
-			/*
-			 * Upon allocation of a locally non-sensitive page, set
-			 * to the allocating mm. Must be set to the same mm when
-			 * the page is freed. May potentially be overwritten in
-			 * the meantime, as long as it is restored before free.
-			 */
-			struct mm_struct *asi_mm;
+			union {
+				/*
+				 * Upon allocation of a locally non-sensitive
+				 * page, set to the allocating mm. Must be set
+				 * to the same mm when the page is freed. May
+				 * potentially be overwritten in the meantime,
+				 * as long as it is restored before free.
+				 */
+				struct mm_struct *asi_mm;
+
+				/*
+				 * Set to the above mm's context ID if the page
+				 * is being freed asynchronously. Can't directly
+				 * use the mm_struct, unless we take additional
+				 * steps to avoid it from being freed while the
+				 * async work is pending.
+				 */
+				u64 asi_mm_ctx_id;
+			};
 		};
 #endif
 	};
diff --git a/mm/page_alloc.c b/mm/page_alloc.c
index 01784bff2a80..998ff6a56732 100644
--- a/mm/page_alloc.c
+++ b/mm/page_alloc.c
@@ -5182,20 +5182,41 @@ static void async_free_work_fn(struct work_struct *work)
 {
 	struct page *page, *tmp;
 	struct llist_node *pages_to_free;
-	void *va;
-	size_t len;
+	size_t addr;
 	uint order;
 
 	pages_to_free = llist_del_all(this_cpu_ptr(&pages_to_free_async));
 
-	/* A later patch will do a more optimized TLB flush. */
+	if (!pages_to_free)
+		return;
+
+	/* If we only have one page to free, then do a targeted TLB flush. */
+	if (!llist_next(pages_to_free)) {
+		page = llist_entry(pages_to_free, struct page, async_free_node);
+		addr = (size_t)page_to_virt(page);
+		order = page->private;
+
+		__asi_flush_tlb_range(page->asi_mm_ctx_id, 0, page->asi_tlb_gen,
+				      addr, addr + PAGE_SIZE * (1 << order),
+				      cpu_online_mask);
+		/* Need to clear, since it shares space with page->mapping. */
+		page->asi_tlb_gen = 0;
+
+		__free_the_page(page, order);
+		return;
+	}
+
+	/*
+	 * Otherwise, do a full flush. We could potentially try to optimize it
+	 * via taking a union of what needs to be flushed, but it may not be
+	 * worth the additional complexity.
+	 */
+	asi_flush_tlb_range(ASI_GLOBAL_NONSENSITIVE, 0, TLB_FLUSH_ALL);
 
 	llist_for_each_entry_safe(page, tmp, pages_to_free, async_free_node) {
-		va = page_to_virt(page);
 		order = page->private;
-		len = PAGE_SIZE * (1 << order);
-
-		asi_flush_tlb_range(ASI_GLOBAL_NONSENSITIVE, va, len);
+		/* Need to clear, since it shares space with page->mapping. */
+		page->asi_tlb_gen = 0;
 		__free_the_page(page, order);
 	}
 }
@@ -5291,6 +5312,11 @@ static bool asi_unmap_freed_pages(struct page *page, unsigned int order)
 	if (!async_flush_needed)
 		return true;
 
+	page->asi_mm_ctx_id = PageGlobalNonSensitive(page)
+			      ? U64_MAX : asi->mm->context.ctx_id;
+
+	__asi_prepare_tlb_flush(asi, &page->asi_tlb_gen);
+
 	page->private = order;
 	llist_add(&page->async_free_node, this_cpu_ptr(&pages_to_free_async));
 
