diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 15f998e501..12f3c0b8e8 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -547,17 +547,14 @@ FlushTlbForAll (
   VOID
   )
 {
-  UINTN  Index;
-
   FlushTlbOnCurrentProcessor (NULL);
-
-  for (Index = 0; Index < gSmst->NumberOfCpus; Index++) {
-    if (Index != gSmst->CurrentlyExecutingCpu) {
-      // Force to start up AP in blocking mode,
-      SmmBlockingStartupThisAp (FlushTlbOnCurrentProcessor, Index, NULL);
-      // Do not check return status, because AP might not be present in some corner cases.
-    }
-  }
+  InternalSmmStartupAllAPs (
+    (EFI_AP_PROCEDURE2)FlushTlbOnCurrentProcessor,
+    0,
+    NULL,
+    NULL,
+    NULL
+    );
 }
 
 /**
@@ -799,72 +796,108 @@ PatchSmmSaveStateMap (
   UINTN  TileCodeSize;
   UINTN  TileDataSize;
   UINTN  TileSize;
+  UINTN  PageTableBase;
 
-  TileCodeSize = GetSmiHandlerSize ();
-  TileCodeSize = ALIGN_VALUE (TileCodeSize, SIZE_4KB);
-  TileDataSize = (SMRAM_SAVE_STATE_MAP_OFFSET - SMM_PSD_OFFSET) + sizeof (SMRAM_SAVE_STATE_MAP);
-  TileDataSize = ALIGN_VALUE (TileDataSize, SIZE_4KB);
-  TileSize     = TileDataSize + TileCodeSize - 1;
-  TileSize     = 2 * GetPowerOfTwo32 ((UINT32)TileSize);
+  TileCodeSize  = GetSmiHandlerSize ();
+  TileCodeSize  = ALIGN_VALUE (TileCodeSize, SIZE_4KB);
+  TileDataSize  = (SMRAM_SAVE_STATE_MAP_OFFSET - SMM_PSD_OFFSET) + sizeof (SMRAM_SAVE_STATE_MAP);
+  TileDataSize  = ALIGN_VALUE (TileDataSize, SIZE_4KB);
+  TileSize      = TileDataSize + TileCodeSize - 1;
+  TileSize      = 2 * GetPowerOfTwo32 ((UINT32)TileSize);
+  PageTableBase = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;
 
   DEBUG ((DEBUG_INFO, "PatchSmmSaveStateMap:\n"));
   for (Index = 0; Index < mMaxNumberOfCpus - 1; Index++) {
     //
     // Code
     //
-    SmmSetMemoryAttributes (
+    ConvertMemoryPageAttributes (
+      PageTableBase,
+      mPagingMode,
       mCpuHotPlugData.SmBase[Index] + SMM_HANDLER_OFFSET,
       TileCodeSize,
-      EFI_MEMORY_RO
+      EFI_MEMORY_RO,
+      TRUE,
+      NULL
       );
-    SmmClearMemoryAttributes (
+    ConvertMemoryPageAttributes (
+      PageTableBase,
+      mPagingMode,
       mCpuHotPlugData.SmBase[Index] + SMM_HANDLER_OFFSET,
       TileCodeSize,
-      EFI_MEMORY_XP
+      EFI_MEMORY_XP,
+      FALSE,
+      NULL
       );
 
     //
     // Data
     //
-    SmmClearMemoryAttributes (
+    ConvertMemoryPageAttributes (
+      PageTableBase,
+      mPagingMode,
       mCpuHotPlugData.SmBase[Index] + SMM_HANDLER_OFFSET + TileCodeSize,
       TileSize - TileCodeSize,
-      EFI_MEMORY_RO
+      EFI_MEMORY_RO,
+      FALSE,
+      NULL
       );
-    SmmSetMemoryAttributes (
+    ConvertMemoryPageAttributes (
+      PageTableBase,
+      mPagingMode,
       mCpuHotPlugData.SmBase[Index] + SMM_HANDLER_OFFSET + TileCodeSize,
       TileSize - TileCodeSize,
-      EFI_MEMORY_XP
+      EFI_MEMORY_XP,
+      TRUE,
+      NULL
       );
   }
 
   //
   // Code
   //
-  SmmSetMemoryAttributes (
+  ConvertMemoryPageAttributes (
+    PageTableBase,
+    mPagingMode,
     mCpuHotPlugData.SmBase[mMaxNumberOfCpus - 1] + SMM_HANDLER_OFFSET,
     TileCodeSize,
-    EFI_MEMORY_RO
+    EFI_MEMORY_RO,
+    TRUE,
+    NULL
     );
-  SmmClearMemoryAttributes (
+  ConvertMemoryPageAttributes (
+    PageTableBase,
+    mPagingMode,
     mCpuHotPlugData.SmBase[mMaxNumberOfCpus - 1] + SMM_HANDLER_OFFSET,
     TileCodeSize,
-    EFI_MEMORY_XP
+    EFI_MEMORY_XP,
+    FALSE,
+    NULL
     );
 
   //
   // Data
   //
-  SmmClearMemoryAttributes (
+  ConvertMemoryPageAttributes (
+    PageTableBase,
+    mPagingMode,
     mCpuHotPlugData.SmBase[mMaxNumberOfCpus - 1] + SMM_HANDLER_OFFSET + TileCodeSize,
     SIZE_32KB - TileCodeSize,
-    EFI_MEMORY_RO
+    EFI_MEMORY_RO,
+    FALSE,
+    NULL
     );
-  SmmSetMemoryAttributes (
+  ConvertMemoryPageAttributes (
+    PageTableBase,
+    mPagingMode,
     mCpuHotPlugData.SmBase[mMaxNumberOfCpus - 1] + SMM_HANDLER_OFFSET + TileCodeSize,
     SIZE_32KB - TileCodeSize,
-    EFI_MEMORY_XP
+    EFI_MEMORY_XP,
+    TRUE,
+    NULL
     );
+
+  FlushTlbForAll ();
 }
 
 /**