struct page *alloc_huge_page_node(struct hstate *h, int nid);
 struct page *alloc_huge_page_nodemask(struct hstate *h, int preferred_nid,
                                nodemask_t *nmask);
-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);
 int huge_add_to_page_cache(struct page *page, struct address_space *mapping,
                        pgoff_t idx);
 
 #define alloc_huge_page(v, a, r) NULL
 #define alloc_huge_page_node(h, nid) NULL
 #define alloc_huge_page_nodemask(h, preferred_nid, nmask) NULL
-#define alloc_huge_page_vma(vma, address) NULL
+#define alloc_huge_page_vma(h, vma, address) NULL
 #define alloc_bootmem_huge_page(h) NULL
 #define hstate_file(f) NULL
 #define hstate_sizelog(s) NULL
 
 }
 
 /* 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);
 
        }
 
        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;