diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index 60d14a5a0e..78cc3e2b93 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -162,7 +162,7 @@ GetWakeupBuffer (
   @retval 0       Cannot find free memory below 4GB.
 **/
 UINTN
-GetModeTransitionBuffer (
+AllocateCodeBuffer (
   IN UINTN  BufferSize
   )
 {
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 4a73787ee4..d761bdc487 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -1056,7 +1056,7 @@ AllocateResetVector (
                                    (CpuMpData->WakeupBuffer +
                                     CpuMpData->AddressMap.RendezvousFunnelSize +
                                     CpuMpData->AddressMap.SwitchToRealSize);
-    CpuMpData->WakeupBufferHigh = GetModeTransitionBuffer (
+    CpuMpData->WakeupBufferHigh = AllocateCodeBuffer (
                                     CpuMpData->AddressMap.RendezvousFunnelSize +
                                     CpuMpData->AddressMap.SwitchToRealSize -
                                     CpuMpData->AddressMap.ModeTransitionOffset
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index f8c52426dd..59ab960897 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -442,7 +442,7 @@ GetWakeupBuffer (
   @retval 0       Cannot find free memory below 4GB.
 **/
 UINTN
-GetModeTransitionBuffer (
+AllocateCodeBuffer (
   IN UINTN  BufferSize
   );
 
diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
index efce574727..65400b95a2 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
@@ -299,14 +299,19 @@ GetWakeupBuffer (
   @retval 0       Cannot find free memory below 4GB.
 **/
 UINTN
-GetModeTransitionBuffer (
+AllocateCodeBuffer (
   IN UINTN  BufferSize
   )
 {
-  //
-  // PEI phase doesn't need to do such transition. So simply return 0.
-  //
-  return 0;
+  EFI_STATUS            Status;
+  EFI_PHYSICAL_ADDRESS  Address;
+
+  Status = PeiServicesAllocatePages (EfiBootServicesCode, EFI_SIZE_TO_PAGES (BufferSize), &Address);
+  if (EFI_ERROR (Status)) {
+    Address = 0;
+  }
+
+  return (UINTN)Address;
 }
 
 /**