diff --git a/drivers/kvm/mmu.c b/drivers/kvm/mmu.c
index c3362ba33b0ac8c32bd2f6157f711adabfb3e7b5..1dcffc49eff8677de1156b367f3fe08819118a5f 100644
--- a/drivers/kvm/mmu.c
+++ b/drivers/kvm/mmu.c
@@ -860,9 +860,9 @@ static void mmu_unshadow(struct kvm *kvm, gfn_t gfn)
 	}
 }
 
-static void page_header_update_slot(struct kvm *kvm, void *pte, gpa_t gpa)
+static void page_header_update_slot(struct kvm *kvm, void *pte, gfn_t gfn)
 {
-	int slot = memslot_id(kvm, gfn_to_memslot(kvm, gpa >> PAGE_SHIFT));
+	int slot = memslot_id(kvm, gfn_to_memslot(kvm, gfn));
 	struct kvm_mmu_page *page_head = page_header(__pa(pte));
 
 	__set_bit(slot, &page_head->slot_bitmap);
@@ -928,7 +928,8 @@ static int nonpaging_map(struct kvm_vcpu *vcpu, gva_t v, hpa_t p)
 				return 0;
 			}
 			mark_page_dirty(vcpu->kvm, v >> PAGE_SHIFT);
-			page_header_update_slot(vcpu->kvm, table, v);
+			page_header_update_slot(vcpu->kvm, table,
+						v >> PAGE_SHIFT);
 			table[index] = p | PT_PRESENT_MASK | PT_WRITABLE_MASK |
 								PT_USER_MASK;
 			if (!was_rmapped)
diff --git a/drivers/kvm/paging_tmpl.h b/drivers/kvm/paging_tmpl.h
index 54a6ee802dc27dc0852ae0223439790d6db5a942..a3da98bfd6a62561dea5b5217fd2b3ff2da95094 100644
--- a/drivers/kvm/paging_tmpl.h
+++ b/drivers/kvm/paging_tmpl.h
@@ -259,8 +259,7 @@ static void FNAME(set_pte)(struct kvm_vcpu *vcpu, pt_element_t gpte,
 
 	pgprintk("%s: setting spte %llx\n", __FUNCTION__, spte);
 	set_shadow_pte(shadow_pte, spte);
-	page_header_update_slot(vcpu->kvm, shadow_pte,
-				(gpa_t)gfn << PAGE_SHIFT);
+	page_header_update_slot(vcpu->kvm, shadow_pte, gfn);
 	if (!was_rmapped) {
 		rmap_add(vcpu, shadow_pte, gfn);
 		if (!is_rmap_pte(*shadow_pte))