diff --git a/sdk/xtdk/amd64/mmtypes.h b/sdk/xtdk/amd64/mmtypes.h
index f472dff..8fccfaf 100644
--- a/sdk/xtdk/amd64/mmtypes.h
+++ b/sdk/xtdk/amd64/mmtypes.h
@@ -101,6 +101,7 @@ typedef struct _MMPAGEMAP_INFO
     ULONGLONG PpeBase;
     ULONGLONG PxeBase;
     ULONGLONG P5eBase;
+    ULONG VaBits;
 } MMPAGEMAP_INFO, *PMMPAGEMAP_INFO;
 
 /* A Page Table Entry on AMD64 system */
diff --git a/xtoskrnl/mm/amd64/init.c b/xtoskrnl/mm/amd64/init.c
index 0b12f80..31483d4 100644
--- a/xtoskrnl/mm/amd64/init.c
+++ b/xtoskrnl/mm/amd64/init.c
@@ -36,6 +36,9 @@ MmInitializePageMapSupport(VOID)
         MmpPageMapInfo.PpeBase = MM_PPE_LA57_BASE;
         MmpPageMapInfo.PxeBase = MM_PXE_LA57_BASE;
         MmpPageMapInfo.P5eBase = MM_P5E_LA57_BASE;
+
+        /* PML5 use 57-bit virtual addresses */
+        MmpPageMapInfo.VaBits = 57;
     }
     else
     {
@@ -51,6 +54,9 @@ MmInitializePageMapSupport(VOID)
         MmpPageMapInfo.PpeBase = MM_PPE_BASE;
         MmpPageMapInfo.PxeBase = MM_PXE_BASE;
         MmpPageMapInfo.P5eBase = 0x0;
+
+        /* PML use 48-bit virtual addresses */
+        MmpPageMapInfo.VaBits = 48;
     }
 }
 
diff --git a/xtoskrnl/mm/amd64/pmap.c b/xtoskrnl/mm/amd64/pmap.c
index a776720..bc29a75 100644
--- a/xtoskrnl/mm/amd64/pmap.c
+++ b/xtoskrnl/mm/amd64/pmap.c
@@ -61,7 +61,7 @@ MmpGetP5eAddress(PVOID Address)
 {
     ULONGLONG Offset;
 
-    Offset = (((ULONGLONG)Address >> MM_P5I_SHIFT) << MM_PTE_SHIFT);
+    Offset = ((((ULONGLONG)Address & (((ULONGLONG)1 << MmpPageMapInfo.VaBits) - 1)) >> MM_P5I_SHIFT) << MM_PTE_SHIFT);
     return (PMMP5E)((MmpPageMapInfo.P5eBase + Offset) * MmpPageMapInfo.Xpa);
 }
 
@@ -81,7 +81,7 @@ MmpGetPdeAddress(PVOID Address)
 {
     ULONGLONG Offset;
 
-    Offset = (((ULONGLONG)Address >> MM_PDI_SHIFT) << MM_PTE_SHIFT);
+    Offset = ((((ULONGLONG)Address & (((ULONGLONG)1 << MmpPageMapInfo.VaBits) - 1)) >> MM_PDI_SHIFT) << MM_PTE_SHIFT);
     return (PMMPDE)(MmpPageMapInfo.PdeBase + Offset);
 }
 
@@ -101,7 +101,7 @@ MmpGetPpeAddress(PVOID Address)
 {
     ULONGLONG Offset;
 
-    Offset = (((ULONGLONG)Address >> MM_PPI_SHIFT) << MM_PTE_SHIFT);
+    Offset = ((((ULONGLONG)Address & (((ULONGLONG)1 << MmpPageMapInfo.VaBits) - 1)) >> MM_PPI_SHIFT) << MM_PTE_SHIFT);
     return (PMMPPE)(MmpPageMapInfo.PpeBase + Offset);
 }
 
@@ -121,7 +121,7 @@ MmpGetPteAddress(PVOID Address)
 {
     ULONGLONG Offset;
 
-    Offset = (((ULONGLONG)Address >> MM_PTI_SHIFT) << MM_PTE_SHIFT);
+    Offset = ((((ULONGLONG)Address & (((ULONGLONG)1 << MmpPageMapInfo.VaBits) - 1)) >> MM_PTI_SHIFT) << MM_PTE_SHIFT);
     return (PMMPTE)(MmpPageMapInfo.PteBase + Offset);
 }
 
@@ -141,7 +141,7 @@ MmpGetPxeAddress(PVOID Address)
 {
     ULONGLONG Offset;
 
-    Offset = (((ULONGLONG)Address >> MM_PXI_SHIFT) << MM_PTE_SHIFT);
+    Offset = ((((ULONGLONG)Address & (((ULONGLONG)1 << MmpPageMapInfo.VaBits) - 1)) >> MM_PXI_SHIFT) << MM_PTE_SHIFT);
     return (PMMPXE)(MmpPageMapInfo.PxeBase + Offset);
 }