diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 9097741..b6d4151 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -1590,15 +1590,17 @@
 			 int offset, int len)
 {
 	int r;
+	struct kvm_memory_slot *memslot;
 	unsigned long addr;
 
-	addr = gfn_to_hva(kvm, gfn);
+	memslot = gfn_to_memslot(kvm, gfn);
+	addr = gfn_to_hva_memslot(memslot, gfn);
 	if (kvm_is_error_hva(addr))
 		return -EFAULT;
 	r = __copy_to_user((void __user *)addr + offset, data, len);
 	if (r)
 		return -EFAULT;
-	mark_page_dirty(kvm, gfn);
+	mark_page_dirty_in_slot(kvm, memslot, gfn);
 	return 0;
 }
 EXPORT_SYMBOL_GPL(kvm_write_guest_page);