diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index bf82faf16c53..d9d9d72d8beb 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -2607,6 +2607,7 @@ int kvm_gmem_slot_set_shared(struct kvm_memory_slot *slot, gfn_t start,
 int kvm_gmem_slot_clear_shared(struct kvm_memory_slot *slot, gfn_t start,
 			       gfn_t end);
 bool kvm_gmem_slot_is_guest_shared(struct kvm_memory_slot *slot, gfn_t gfn);
+int kvm_gmem_slot_register_callback(struct kvm_memory_slot *slot, gfn_t gfn);
 void kvm_gmem_handle_folio_put(struct folio *folio);
 #else
 static inline int kvm_gmem_set_shared(struct kvm *kvm, gfn_t start, gfn_t end)
@@ -2638,6 +2639,11 @@ static inline bool kvm_gmem_slot_is_guest_shared(struct kvm_memory_slot *slot,
 	WARN_ON_ONCE(1);
 	return false;
 }
+static inline int kvm_gmem_slot_register_callback(struct kvm_memory_slot *slot, gfn_t gfn)
+{
+	WARN_ON_ONCE(1);
+	return -EINVAL;
+}
 #endif /* CONFIG_KVM_GMEM_SHARED_MEM */
 
 #endif
diff --git a/virt/kvm/guest_memfd.c b/virt/kvm/guest_memfd.c
index 3b4d724084a8..ce19bd6c2031 100644
--- a/virt/kvm/guest_memfd.c
+++ b/virt/kvm/guest_memfd.c
@@ -392,6 +392,27 @@ enum folio_shareability {
 	KVM_GMEM_NONE_SHARED	= 0b11, /* Not shared, transient state. */
 };
 
+/*
+ * Unregisters the __folio_put() callback from the folio.
+ *
+ * Restores a folio's refcount after all pending references have been released,
+ * and removes the folio type, thereby removing the callback. Now the folio can
+ * be freed normaly once all actual references have been dropped.
+ *
+ * Must be called with the folio locked and the offsets_lock write lock held.
+ */
+static void kvm_gmem_restore_pending_folio(struct folio *folio, struct inode *inode)
+{
+	rwlock_t *offsets_lock = &kvm_gmem_private(inode)->offsets_lock;
+
+	lockdep_assert_held_write(offsets_lock);
+	WARN_ON_ONCE(!folio_test_locked(folio));
+	WARN_ON_ONCE(!folio_test_guestmem(folio));
+
+	__folio_clear_guestmem(folio);
+	folio_ref_add(folio, folio_nr_pages(folio));
+}
+
 static int kvm_gmem_offset_set_shared(struct inode *inode, pgoff_t index)
 {
 	struct xarray *shared_offsets = &kvm_gmem_private(inode)->shared_offsets;
@@ -400,6 +421,24 @@ static int kvm_gmem_offset_set_shared(struct inode *inode, pgoff_t index)
 
 	lockdep_assert_held_write(offsets_lock);
 
+	/*
+	 * If the folio is NONE_SHARED, it indicates that it is transitioning to
+	 * private (GUEST_SHARED). Transition it to shared (ALL_SHARED)
+	 * immediately, and remove the callback.
+	 */
+	if (xa_to_value(xa_load(shared_offsets, index)) == KVM_GMEM_NONE_SHARED) {
+		struct folio *folio = filemap_lock_folio(inode->i_mapping, index);
+
+		if (WARN_ON_ONCE(IS_ERR(folio)))
+			return PTR_ERR(folio);
+
+		if (folio_test_guestmem(folio))
+			kvm_gmem_restore_pending_folio(folio, inode);
+
+		folio_unlock(folio);
+		folio_put(folio);
+	}
+
 	return xa_err(xa_store(shared_offsets, index, xval, GFP_KERNEL));
 }
 
@@ -503,9 +542,111 @@ static int kvm_gmem_offset_range_clear_shared(struct inode *inode,
 	return r;
 }
 
+/*
+ * Registers a callback to __folio_put(), so that gmem knows that the host does
+ * not have any references to the folio. The callback itself is registered by
+ * setting the folio type to guestmem.
+ *
+ * Returns 0 if a callback was registered or already has been registered, or
+ * -EAGAIN if the host has references, indicating a callback wasn't registered.
+ *
+ * Must be called with the folio locked and the offsets_lock write lock held.
+ */
+static int kvm_gmem_register_callback(struct folio *folio, struct inode *inode, pgoff_t index)
+{
+	struct xarray *shared_offsets = &kvm_gmem_private(inode)->shared_offsets;
+	rwlock_t *offsets_lock = &kvm_gmem_private(inode)->offsets_lock;
+	void *xval_guest = xa_mk_value(KVM_GMEM_GUEST_SHARED);
+	int refcount;
+	int r = 0;
+
+	lockdep_assert_held_write(offsets_lock);
+	WARN_ON_ONCE(!folio_test_locked(folio));
+
+	if (folio_test_guestmem(folio))
+		return 0;
+
+	if (folio_mapped(folio))
+		return -EAGAIN;
+
+	refcount = folio_ref_count(folio);
+	if (!folio_ref_freeze(folio, refcount))
+		return -EAGAIN;
+
+	/*
+	 * Register callback by setting the folio type and subtracting gmem's
+	 * references for it to trigger once outstanding references are dropped.
+	 */
+	if (refcount > 1) {
+		__folio_set_guestmem(folio);
+		refcount -= folio_nr_pages(folio);
+	} else {
+		/* No outstanding references, transition it to guest shared. */
+		r = WARN_ON_ONCE(xa_err(xa_store(shared_offsets, index, xval_guest, GFP_KERNEL)));
+	}
+
+	folio_ref_unfreeze(folio, refcount);
+	return r;
+}
+
+int kvm_gmem_slot_register_callback(struct kvm_memory_slot *slot, gfn_t gfn)
+{
+	unsigned long pgoff = slot->gmem.pgoff + gfn - slot->base_gfn;
+	struct inode *inode = file_inode(READ_ONCE(slot->gmem.file));
+	rwlock_t *offsets_lock = &kvm_gmem_private(inode)->offsets_lock;
+	struct folio *folio;
+	int r;
+
+	write_lock(offsets_lock);
+
+	folio = filemap_lock_folio(inode->i_mapping, pgoff);
+	if (WARN_ON_ONCE(IS_ERR(folio))) {
+		write_unlock(offsets_lock);
+		return PTR_ERR(folio);
+	}
+
+	r = kvm_gmem_register_callback(folio, inode, pgoff);
+
+	folio_unlock(folio);
+	folio_put(folio);
+	write_unlock(offsets_lock);
+
+	return r;
+}
+EXPORT_SYMBOL_GPL(kvm_gmem_slot_register_callback);
+
+/*
+ * Callback function for __folio_put(), i.e., called once all references by the
+ * host to the folio have been dropped. This allows gmem to transition the state
+ * of the folio to shared with the guest, and allows the hypervisor to continue
+ * transitioning its state to private, since the host cannot attempt to access
+ * it anymore.
+ */
 void kvm_gmem_handle_folio_put(struct folio *folio)
 {
-	WARN_ONCE(1, "A placeholder that shouldn't trigger. Work in progress.");
+	struct address_space *mapping;
+	struct xarray *shared_offsets;
+	rwlock_t *offsets_lock;
+	struct inode *inode;
+	pgoff_t index;
+	void *xval;
+
+	mapping = folio->mapping;
+	if (WARN_ON_ONCE(!mapping))
+		return;
+
+	inode = mapping->host;
+	index = folio->index;
+	shared_offsets = &kvm_gmem_private(inode)->shared_offsets;
+	offsets_lock = &kvm_gmem_private(inode)->offsets_lock;
+	xval = xa_mk_value(KVM_GMEM_GUEST_SHARED);
+
+	write_lock(offsets_lock);
+	folio_lock(folio);
+	kvm_gmem_restore_pending_folio(folio, inode);
+	folio_unlock(folio);
+	WARN_ON_ONCE(xa_err(xa_store(shared_offsets, index, xval, GFP_KERNEL)));
+	write_unlock(offsets_lock);
 }
 EXPORT_SYMBOL_GPL(kvm_gmem_handle_folio_put);
 
