diff --git a/arch/s390/include/asm/hugetlb.h b/arch/s390/include/asm/hugetlb.h
index e33a5b587ee4..02db08429b7c 100644
--- a/arch/s390/include/asm/hugetlb.h
+++ b/arch/s390/include/asm/hugetlb.h
@@ -42,9 +42,9 @@ static inline void huge_pte_clear(struct mm_struct *mm, unsigned long addr,
 				  pte_t *ptep, unsigned long sz)
 {
 	if ((pte_val(ptep_get(ptep)) & _REGION_ENTRY_TYPE_MASK) == _REGION_ENTRY_TYPE_R3)
-		set_pte(ptep, __pte(_REGION3_ENTRY_EMPTY));
+		set_pud((pud_t *)ptep, __pud(_REGION3_ENTRY_EMPTY));
 	else
-		set_pte(ptep, __pte(_SEGMENT_ENTRY_EMPTY));
+		set_pmd((pmd_t *)ptep, __pmd(_SEGMENT_ENTRY_EMPTY));
 }
 
 #define __HAVE_ARCH_HUGE_PTEP_CLEAR_FLUSH
diff --git a/arch/s390/mm/hugetlbpage.c b/arch/s390/mm/hugetlbpage.c
index db35d8fe8609..f84aa9265430 100644
--- a/arch/s390/mm/hugetlbpage.c
+++ b/arch/s390/mm/hugetlbpage.c
@@ -147,10 +147,12 @@ void __set_huge_pte_at(struct mm_struct *mm, unsigned long addr,
 		if (likely(pte_present(pte)))
 			rste |= _REGION3_ENTRY_LARGE;
 		rste |= _REGION_ENTRY_TYPE_R3;
-	} else if (likely(pte_present(pte)))
-		rste |= _SEGMENT_ENTRY_LARGE;
-
-	set_pte(ptep, __pte(rste));
+		set_pud((pud_t *)ptep, __pud(rste));
+	} else {
+		if (likely(pte_present(pte)))
+			rste |= _SEGMENT_ENTRY_LARGE;
+		set_pmd((pmd_t *)ptep, __pmd(rste));
+	}
 }
 
 void set_huge_pte_at(struct mm_struct *mm, unsigned long addr,