diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index 742a929..7c204e3 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -1675,16 +1675,15 @@
 }
 
 /* mempolicy aware migration callback */
-struct page *alloc_huge_page_vma(struct vm_area_struct *vma, unsigned long address)
+struct page *alloc_huge_page_vma(struct hstate *h, struct vm_area_struct *vma,
+		unsigned long address)
 {
 	struct mempolicy *mpol;
 	nodemask_t *nodemask;
 	struct page *page;
-	struct hstate *h;
 	gfp_t gfp_mask;
 	int node;
 
-	h = hstate_vma(vma);
 	gfp_mask = htlb_alloc_mask(h);
 	node = huge_node(vma, address, gfp_mask, &mpol, &nodemask);
 	page = alloc_huge_page_nodemask(h, node, nodemask);
diff --git a/mm/mempolicy.c b/mm/mempolicy.c
index 96823fa..d879f1d 100644
--- a/mm/mempolicy.c
+++ b/mm/mempolicy.c
@@ -1121,7 +1121,8 @@
 	}
 
 	if (PageHuge(page)) {
-		return alloc_huge_page_vma(vma, address);
+		return alloc_huge_page_vma(page_hstate(compound_head(page)),
+				vma, address);
 	} else if (thp_migration_supported() && PageTransHuge(page)) {
 		struct page *thp;