diff --git a/arch/x86/mm/pageattr_32.c b/arch/x86/mm/pageattr_32.c
index 5cb5c7101f413ec24a558c72d7ed89b69e927bd1..66688a630839ef6e6c6593644ca61d4952ffc692 100644
--- a/arch/x86/mm/pageattr_32.c
+++ b/arch/x86/mm/pageattr_32.c
@@ -172,7 +172,7 @@ static int __change_page_attr(struct page *page, pgprot_t prot)
 	BUG_ON(PageCompound(kpte_page));
 
 	if (pgprot_val(prot) != pgprot_val(PAGE_KERNEL)) {
-		if (!pte_huge(*kpte)) {
+		if (level == 3) {
 			set_pte_atomic(kpte, mk_pte(page, prot));
 		} else {
 			struct page *split;
@@ -190,7 +190,7 @@ static int __change_page_attr(struct page *page, pgprot_t prot)
 		}
 		page_private(kpte_page)++;
 	} else {
-		if (!pte_huge(*kpte)) {
+		if (level == 3) {
 			set_pte_atomic(kpte, mk_pte(page, PAGE_KERNEL));
 			BUG_ON(page_private(kpte_page) == 0);
 			page_private(kpte_page)--;
diff --git a/arch/x86/mm/pageattr_64.c b/arch/x86/mm/pageattr_64.c
index 3ccdb1401e673c219ac7f8aafb208da012b182a0..73dbbb4048e664d58f2bac71e49c15234cf5b877 100644
--- a/arch/x86/mm/pageattr_64.c
+++ b/arch/x86/mm/pageattr_64.c
@@ -153,7 +153,7 @@ __change_page_attr(unsigned long address, unsigned long pfn, pgprot_t prot,
 	BUG_ON(PageLRU(kpte_page));
 	BUG_ON(PageCompound(kpte_page));
 	if (pgprot_val(prot) != pgprot_val(ref_prot)) {
-		if (!pte_huge(*kpte)) {
+		if (level == 4) {
 			set_pte(kpte, pfn_pte(pfn, prot));
 		} else {
 			/*
@@ -172,7 +172,7 @@ __change_page_attr(unsigned long address, unsigned long pfn, pgprot_t prot,
 		}
 		page_private(kpte_page)++;
 	} else {
-		if (!pte_huge(*kpte)) {
+		if (level == 4) {
 			set_pte(kpte, pfn_pte(pfn, ref_prot));
 			BUG_ON(page_private(kpte_page) == 0);
 			page_private(kpte_page)--;