chore: import upstream snapshot with attribution

test/cpp/fluid/memory/CMakeLists.txt

@@ -0,0 +1,232 @@
+cc_test(
+  memory_stats_test
+  SRCS memory_stats_test.cc
+  DEPS)
+cc_test(
+  stats_test
+  SRCS stats_test.cc
+  DEPS)
+
+cc_test(
+  naive_best_fit_allocator_test
+  SRCS naive_best_fit_allocator_test.cc
+  DEPS phi common)
+cc_test(
+  buffered_allocator_test
+  SRCS buffered_allocator_test.cc
+  DEPS phi common)
+
+if(WITH_GPU)
+  nv_test(
+    thread_local_allocator_test
+    SRCS thread_local_allocator_test.cc
+    DEPS phi common)
+endif()
+if(WITH_ROCM)
+  hip_test(
+    thread_local_allocator_test
+    SRCS thread_local_allocator_test.cc
+    DEPS phi common)
+endif()
+
+if(WITH_GPU)
+  nv_test(
+    best_fit_allocator_test
+    SRCS best_fit_allocator_test.cc best_fit_allocator_test.cu
+    DEPS phi common)
+elseif(WITH_ROCM)
+  hip_test(
+    best_fit_allocator_test
+    SRCS best_fit_allocator_test.cc best_fit_allocator_test.cu
+    DEPS phi common)
+else()
+  if(WIN32)
+    cc_test(
+      best_fit_allocator_test
+      SRCS best_fit_allocator_test.cc
+      DEPS type_info common)
+  else()
+    cc_test(
+      best_fit_allocator_test
+      SRCS best_fit_allocator_test.cc
+      DEPS phi common)
+  endif()
+endif()
+
+cc_test(
+  test_aligned_allocator
+  SRCS test_aligned_allocator.cc
+  DEPS phi common)
+
+if(WITH_GPU)
+  if(WIN32)
+    message(STATUS "Skip allocator_visitor_test on Windows")
+  else()
+    nv_test(
+      allocator_visitor_test
+      SRCS allocator_visitor_test.cc
+      DEPS phi common)
+  endif()
+endif()
+
+if(WIN32)
+  cc_test(
+    retry_allocator_test
+    SRCS retry_allocator_test.cc
+    DEPS type_info common)
+else()
+  cc_test(
+    retry_allocator_test
+    SRCS retry_allocator_test.cc
+    DEPS phi common)
+endif()
+if(TEST retry_allocator_test)
+  set_tests_properties(retry_allocator_test PROPERTIES LABELS
+                                                       "RUN_TYPE=EXCLUSIVE")
+endif()
+
+cc_test(
+  allocator_facade_abs_flags_test
+  SRCS allocator_facade_abs_flags_test.cc
+  DEPS phi common)
+
+if(WITH_GPU)
+  if(WIN32)
+    message(STATUS "Skip vmm_allocator_tail_test on Windows")
+  else()
+    nv_test(
+      vmm_allocator_tail_test
+      SRCS vmm_allocator_tail_test.cc
+      DEPS phi common)
+    nv_test(
+      cuda_virtual_mem_allocator_v2_test
+      SRCS cuda_virtual_mem_allocator_v2_test.cu
+      DEPS phi common)
+  endif()
+endif()
+
+cc_test(
+  allocator_facade_frac_flags_test
+  SRCS allocator_facade_frac_flags_test.cc
+  DEPS phi common)
+
+if(WITH_GPU)
+  nv_test(
+    malloc_test
+    SRCS malloc_test.cu
+    DEPS phi)
+  nv_test(
+    stream_safe_cuda_alloc_test
+    SRCS stream_safe_cuda_alloc_test.cu
+    DEPS phi common)
+  nv_test(
+    cuda_managed_memory_test
+    SRCS cuda_managed_memory_test.cu
+    DEPS phi common)
+  nv_test(
+    cuda_malloc_async_allocator_test
+    SRCS cuda_malloc_async_allocator_test.cu
+    DEPS phi common)
+  if(WITH_TESTING AND TEST stream_safe_cuda_alloc_test)
+    set_tests_properties(
+      stream_safe_cuda_alloc_test
+      PROPERTIES ENVIRONMENT "FLAGS_use_stream_safe_cuda_allocator=true; \
+        FLAGS_allocator_strategy=auto_growth")
+  endif()
+  if(WITH_TESTING AND TEST cuda_malloc_async_allocator_test)
+    set_tests_properties(
+      cuda_malloc_async_allocator_test
+      PROPERTIES ENVIRONMENT "FLAGS_use_cuda_malloc_async_allocator=true")
+  endif()
+endif()
+
+if(WITH_ROCM)
+  hip_test(
+    malloc_test
+    SRCS malloc_test.cu
+    DEPS phi)
+  hip_test(
+    cuda_managed_memory_test
+    SRCS cuda_managed_memory_test.cu
+    DEPS phi common)
+endif()
+
+if(WITH_TESTING AND TEST cuda_managed_memory_test)
+  set_tests_properties(
+    cuda_managed_memory_test
+    PROPERTIES
+      ENVIRONMENT
+      "FLAGS_use_cuda_managed_memory=true;FLAGS_use_cuda_malloc_async_allocator=false;FLAGS_allocator_strategy=auto_growth"
+      TIMEOUT
+      50)
+endif()
+
+if(WITH_GPU AND WITH_TESTING)
+  nv_test(
+    get_base_ptr_test
+    SRCS get_base_ptr_test.cu
+    DEPS phi common)
+  set_tests_properties(
+    get_base_ptr_test
+    PROPERTIES ENVIRONMENT "FLAGS_allocator_strategy=auto_growth;
+                                    FLAGS_use_stream_safe_cuda_allocator=true;")
+endif()
+
+if(WIN32)
+  cc_test(
+    auto_growth_best_fit_allocator_facade_test
+    SRCS auto_growth_best_fit_allocator_facade_test.cc
+    DEPS type_info common)
+else()
+  cc_test(
+    auto_growth_best_fit_allocator_facade_test
+    SRCS auto_growth_best_fit_allocator_facade_test.cc
+    DEPS phi common)
+endif()
+cc_test(
+  auto_growth_best_fit_allocator_test
+  SRCS auto_growth_best_fit_allocator_test.cc
+  DEPS phi common)
+
+if(NOT WIN32)
+  cc_test(
+    mmap_allocator_test
+    SRCS mmap_allocator_test.cc
+    DEPS phi common)
+endif()
+
+cc_test(
+  system_allocator_test
+  SRCS system_allocator_test.cc
+  DEPS phi common)
+
+if(WITH_GPU)
+  if(WIN32)
+    message(STATUS "Skip vmm_auto_growth_best_fit_allocator_test on Windows")
+  else()
+    nv_test(
+      vmm_auto_growth_best_fit_allocator_test
+      SRCS vmm_auto_growth_best_fit_allocator_test.cu
+      DEPS phi common)
+    nv_test(
+      vmm_auto_growth_best_fit_allocator_v2_test
+      SRCS vmm_auto_growth_best_fit_allocator_v2_test.cu
+      DEPS phi common)
+    if(WITH_TESTING AND TEST vmm_auto_growth_best_fit_allocator_test)
+      set_tests_properties(
+        vmm_auto_growth_best_fit_allocator_test
+        PROPERTIES ENVIRONMENT "FLAGS_use_virtual_memory_auto_growth=true")
+    endif()
+  endif()
+endif()
+
+if(WITH_GPU)
+  if(WIN32)
+    message(STATUS "Skip multi_scale_allocator_test on Windows")
+  else()
+    nv_test(
+      multi_scale_allocator_test
+      SRCS multi_scale_allocator_test.cc
+      DEPS phi common)
+  endif()
+endif()

test/cpp/fluid/memory/allocator_facade_abs_flags_test.cc

@@ -0,0 +1,104 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/core/memory/allocation/allocator_facade.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+COMMON_DECLARE_double(fraction_of_gpu_memory_to_use);
+COMMON_DECLARE_double(fraction_of_cuda_pinned_memory_to_use);
+COMMON_DECLARE_uint64(initial_gpu_memory_in_mb);
+COMMON_DECLARE_uint64(reallocate_gpu_memory_in_mb);
+PD_DECLARE_int64(gpu_allocator_retry_time);
+#endif
+COMMON_DECLARE_string(allocator_strategy);
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+//! Run allocate test cases for different places
+void AllocateTestCases() {
+  auto &instance = AllocatorFacade::Instance();
+  phi::Place place;
+  size_t size = 1024;
+
+  {
+    place = phi::CPUPlace();
+    size = 1024;
+    auto cpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(cpu_allocation, nullptr);
+    ASSERT_NE(cpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(cpu_allocation->place(), place);
+    ASSERT_EQ(cpu_allocation->size(), size);
+  }
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  {
+    place = phi::GPUPlace(0);
+    size = 1024;
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(), size);
+  }
+
+  {
+    // Allocate 2GB gpu memory
+    place = phi::GPUPlace(0);
+    size = 2 * static_cast<size_t>(1 << 30);
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(), size);
+  }
+
+  {
+    place = phi::GPUPinnedPlace();
+    size = (1 << 20);
+    auto cuda_pinned_allocation =
+        instance.Alloc(phi::GPUPinnedPlace(), 1 << 20);
+    ASSERT_NE(cuda_pinned_allocation, nullptr);
+    ASSERT_NE(cuda_pinned_allocation->ptr(), nullptr);
+    ASSERT_EQ(cuda_pinned_allocation->place(), place);
+    ASSERT_GE(cuda_pinned_allocation->size(), size);
+  }
+#endif
+}
+
+TEST(Allocator, SpecifyGpuMemory) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  // Set to 0.0 to test FLAGS_initial_gpu_memory_in_mb and
+  // FLAGS_reallocate_gpu_memory_in_mb
+  FLAGS_fraction_of_gpu_memory_to_use = 0.0;
+  // 512 MB
+  FLAGS_initial_gpu_memory_in_mb = 512;
+  // 4 MB
+  FLAGS_reallocate_gpu_memory_in_mb = 4;
+  FLAGS_gpu_allocator_retry_time = 500;
+  FLAGS_fraction_of_cuda_pinned_memory_to_use = 0.5;
+#endif
+
+  FLAGS_allocator_strategy = "naive_best_fit";
+
+  AllocateTestCases();
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/allocator_facade_frac_flags_test.cc

@@ -0,0 +1,97 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/core/memory/allocation/allocator_facade.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+COMMON_DECLARE_double(fraction_of_gpu_memory_to_use);
+COMMON_DECLARE_double(fraction_of_cuda_pinned_memory_to_use);
+COMMON_DECLARE_uint64(initial_gpu_memory_in_mb);
+COMMON_DECLARE_uint64(reallocate_gpu_memory_in_mb);
+PD_DECLARE_int64(gpu_allocator_retry_time);
+#endif
+COMMON_DECLARE_string(allocator_strategy);
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+//! Run allocate test cases for different places
+void AllocateTestCases() {
+  auto &instance = AllocatorFacade::Instance();
+  phi::Place place;
+  size_t size = 1024;
+
+  {
+    place = phi::CPUPlace();
+    size = 1024;
+    auto cpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(cpu_allocation, nullptr);
+    ASSERT_NE(cpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(cpu_allocation->place(), place);
+    ASSERT_EQ(cpu_allocation->size(), size);
+  }
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  {
+    place = phi::GPUPlace(0);
+    size = 1024;
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(), size);
+  }
+
+  {
+    // Allocate 2GB gpu memory
+    place = phi::GPUPlace(0);
+    size = 2 * static_cast<size_t>(1 << 30);
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(), size);
+  }
+
+  {
+    place = phi::GPUPinnedPlace();
+    size = (1 << 20);
+    auto cuda_pinned_allocation =
+        instance.Alloc(phi::GPUPinnedPlace(), 1 << 20);
+    ASSERT_NE(cuda_pinned_allocation, nullptr);
+    ASSERT_NE(cuda_pinned_allocation->ptr(), nullptr);
+    ASSERT_EQ(cuda_pinned_allocation->place(), place);
+    ASSERT_GE(cuda_pinned_allocation->size(), size);
+  }
+#endif
+}
+
+TEST(Allocator, Allocator) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  FLAGS_fraction_of_gpu_memory_to_use = 0.01;
+  FLAGS_gpu_allocator_retry_time = 500;
+  FLAGS_fraction_of_cuda_pinned_memory_to_use = 0.5;
+#endif
+  FLAGS_allocator_strategy = "naive_best_fit";
+
+  AllocateTestCases();
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/allocator_visitor_test.cc

@@ -0,0 +1,213 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.h"
+#include "paddle/phi/core/memory/allocation/cpu_allocator.h"
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.h"
+#include "paddle/phi/core/memory/allocation/retry_allocator.h"
+#include "paddle/phi/core/memory/allocation/virtual_memory_auto_growth_best_fit_allocator.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+#include "glog/logging.h"
+#include "gtest/gtest.h"
+
+PD_DECLARE_uint64(vmm_small_pool_size_in_mb);
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(VirtualMemoryAutoGrowthBestFitAllocator, TestAllocatorVisitor) {
+  FLAGS_v = 1;
+  auto vmm_cuda_allocator =
+      std::make_shared<CUDAVirtualMemAllocator>(phi::GPUPlace());
+  auto vma_allocator =
+      std::make_shared<VirtualMemoryAutoGrowthBestFitAllocator>(
+          vmm_cuda_allocator, platform::GpuMinChunkSize(), phi::GPUPlace());
+  memory::AllocatorVisitor visitor;
+  vma_allocator->Accept(&visitor);
+}
+
+TEST(VirtualMemoryAutoGrowthBestFitAllocator, TestAllBlocksInfoVisitor) {
+  auto vmm_cuda_allocator =
+      std::make_shared<CUDAVirtualMemAllocator>(phi::GPUPlace());
+  auto vma_allocator =
+      std::make_shared<VirtualMemoryAutoGrowthBestFitAllocator>(
+          vmm_cuda_allocator, platform::GpuMinChunkSize(), phi::GPUPlace());
+
+  auto allocation = vma_allocator->Allocate(platform::GpuMinChunkSize());
+  AllBlocksInfoVisitor visitor;
+  vma_allocator->Accept(&visitor);
+  auto all_blocks_info = std::move(visitor).GetAllBlocksInfo();
+
+  ASSERT_EQ(all_blocks_info.size(), 1UL);
+  ASSERT_GE(all_blocks_info[0].size(), 1UL);
+  uintptr_t expected_addr = std::get<1>(all_blocks_info[0][0]);
+  size_t total_size = 0;
+  bool found_allocated_block = false;
+  for (const auto& block_info : all_blocks_info[0]) {
+    const size_t block_size = std::get<0>(block_info);
+    const uintptr_t block_addr = std::get<1>(block_info);
+    EXPECT_GT(block_size, 0UL);
+    EXPECT_EQ(block_addr, expected_addr);
+    expected_addr += block_size;
+    total_size += block_size;
+    if (std::get<1>(block_info) ==
+        reinterpret_cast<uintptr_t>(allocation->ptr())) {
+      EXPECT_EQ(std::get<0>(block_info), platform::GpuMinChunkSize());
+      EXPECT_FALSE(std::get<2>(block_info));
+      found_allocated_block = true;
+    }
+  }
+  EXPECT_TRUE(found_allocated_block);
+
+  allocation.reset();
+  AllBlocksInfoVisitor after_free_visitor;
+  vma_allocator->Accept(&after_free_visitor);
+  auto after_free_info = std::move(after_free_visitor).GetAllBlocksInfo();
+
+  ASSERT_EQ(after_free_info.size(), 1UL);
+  ASSERT_EQ(after_free_info[0].size(), 1UL);
+  EXPECT_EQ(std::get<0>(after_free_info[0][0]), total_size);
+  EXPECT_EQ(std::get<1>(after_free_info[0][0]),
+            std::get<1>(all_blocks_info[0][0]));
+  EXPECT_TRUE(std::get<2>(after_free_info[0][0]));
+}
+
+TEST(AutoGrowthBestFitAllocator, TestAllBlocksInfoVisitor) {
+  constexpr size_t kAlignment = 256;
+  constexpr size_t kChunkSize = 4096;
+  auto cpu_allocator = std::make_shared<CPUAllocator>();
+  auto ag_allocator = std::make_shared<AutoGrowthBestFitAllocator>(
+      cpu_allocator, kAlignment, kChunkSize);
+
+  auto first = ag_allocator->Allocate(1024);
+  auto second = ag_allocator->Allocate(2048);
+  AllBlocksInfoVisitor visitor;
+  ag_allocator->Accept(&visitor);
+  auto all_blocks_info = std::move(visitor).GetAllBlocksInfo();
+
+  ASSERT_EQ(all_blocks_info.size(), 1UL);
+  ASSERT_EQ(all_blocks_info[0].size(), 3UL);
+  uintptr_t expected_addr = std::get<1>(all_blocks_info[0][0]);
+  size_t total_size = 0;
+  for (const auto& block_info : all_blocks_info[0]) {
+    const size_t block_size = std::get<0>(block_info);
+    const uintptr_t block_addr = std::get<1>(block_info);
+    EXPECT_GT(block_size, 0UL);
+    EXPECT_EQ(block_addr, expected_addr);
+    expected_addr += block_size;
+    total_size += block_size;
+  }
+  ASSERT_EQ(total_size, kChunkSize);
+
+  EXPECT_EQ(std::get<0>(all_blocks_info[0][0]), 1024UL);
+  EXPECT_TRUE(std::get<2>(all_blocks_info[0][0]));
+  EXPECT_EQ(std::get<0>(all_blocks_info[0][1]), 2048UL);
+  EXPECT_EQ(std::get<1>(all_blocks_info[0][1]),
+            reinterpret_cast<uintptr_t>(second->ptr()));
+  EXPECT_FALSE(std::get<2>(all_blocks_info[0][1]));
+  EXPECT_EQ(std::get<0>(all_blocks_info[0][2]), 1024UL);
+  EXPECT_EQ(std::get<1>(all_blocks_info[0][2]),
+            reinterpret_cast<uintptr_t>(first->ptr()));
+  EXPECT_FALSE(std::get<2>(all_blocks_info[0][2]));
+
+  first.reset();
+  second.reset();
+  AllBlocksInfoVisitor after_free_visitor;
+  ag_allocator->Accept(&after_free_visitor);
+  auto after_free_info = std::move(after_free_visitor).GetAllBlocksInfo();
+
+  ASSERT_EQ(after_free_info.size(), 1UL);
+  ASSERT_EQ(after_free_info[0].size(), 1UL);
+  EXPECT_EQ(std::get<0>(after_free_info[0][0]), kChunkSize);
+  EXPECT_EQ(std::get<1>(after_free_info[0][0]),
+            std::get<1>(all_blocks_info[0][0]));
+  EXPECT_TRUE(std::get<2>(after_free_info[0][0]));
+}
+
+TEST(MultiScalePoolAllocator, TestPoolFilterLargeOnly) {
+  FLAGS_vmm_small_pool_size_in_mb = 20;
+
+  auto vmm_cuda_allocator_small =
+      std::make_shared<CUDAVirtualMemAllocator>(phi::GPUPlace(0));
+  auto vmm_cuda_allocator_large =
+      std::make_shared<CUDAVirtualMemAllocator>(phi::GPUPlace(0));
+  auto small_alloc = std::make_shared<VirtualMemoryAutoGrowthBestFitAllocator>(
+      vmm_cuda_allocator_small, platform::GpuMinChunkSize(), GPUPlace(0));
+  auto large_alloc = std::make_shared<VirtualMemoryAutoGrowthBestFitAllocator>(
+      vmm_cuda_allocator_large, platform::GpuMinChunkSize(), GPUPlace(0));
+  auto multi_scale =
+      std::make_shared<VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator>(
+          small_alloc, large_alloc, platform::GpuMinChunkSize(), GPUPlace(0));
+
+  // Allocate in small pool (< 20MB)
+  auto small_allocation = multi_scale->Allocate(1 << 20);  // 1MB
+  // Allocate in large pool (>= 20MB)
+  auto large_allocation = multi_scale->Allocate(30 << 20);  // 30MB
+
+  // Query all pools
+  AllBlocksInfoVisitor all_visitor;
+  multi_scale->Accept(&all_visitor);
+  auto all_info = std::move(all_visitor).GetAllBlocksInfo();
+  ASSERT_EQ(all_info.size(), 2UL);  // small + large
+
+  // Query large pool only
+  AllBlocksInfoVisitor large_visitor(
+      AllBlocksInfoVisitor::PoolFilter::kLargeOnly);
+  multi_scale->Accept(&large_visitor);
+  auto large_info = std::move(large_visitor).GetAllBlocksInfo();
+  ASSERT_EQ(large_info.size(), 1UL);  // only large
+
+  // Verify large pool contains allocated block
+  bool found_large = false;
+  for (const auto& block : large_info[0]) {
+    if (std::get<1>(block) ==
+        reinterpret_cast<uintptr_t>(large_allocation->ptr())) {
+      found_large = true;
+      EXPECT_FALSE(std::get<2>(block));  // not free
+    }
+  }
+  EXPECT_TRUE(found_large);
+
+  // Query small pool only
+  AllBlocksInfoVisitor small_visitor(
+      AllBlocksInfoVisitor::PoolFilter::kSmallOnly);
+  multi_scale->Accept(&small_visitor);
+  auto small_info = std::move(small_visitor).GetAllBlocksInfo();
+  ASSERT_EQ(small_info.size(), 1UL);  // only small
+
+  // Verify small pool contains allocated block
+  bool found_small = false;
+  for (const auto& block : small_info[0]) {
+    if (std::get<1>(block) ==
+        reinterpret_cast<uintptr_t>(small_allocation->ptr())) {
+      found_small = true;
+      EXPECT_FALSE(std::get<2>(block));  // not free
+    }
+  }
+  EXPECT_TRUE(found_small);
+
+  // Verify large pool does NOT contain small allocation
+  for (const auto& block : large_info[0]) {
+    EXPECT_NE(std::get<1>(block),
+              reinterpret_cast<uintptr_t>(small_allocation->ptr()));
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/auto_growth_best_fit_allocator_facade_test.cc

@@ -0,0 +1,149 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include <condition_variable>  // NOLINT
+#include <mutex>               // NOLINT
+#include <random>
+#include <thread>  // NOLINT
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/core/memory/allocation/allocator_facade.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+COMMON_DECLARE_double(fraction_of_gpu_memory_to_use);
+COMMON_DECLARE_double(fraction_of_cuda_pinned_memory_to_use);
+PD_DECLARE_int64(gpu_allocator_retry_time);
+#endif
+
+COMMON_DECLARE_string(allocator_strategy);
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+static inline size_t AlignTo(size_t size, size_t alignment) {
+  auto remaining = size % alignment;
+  return remaining == 0 ? size : size + alignment - remaining;
+}
+
+TEST(allocator, allocator) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  FLAGS_fraction_of_gpu_memory_to_use = 0.01;
+  FLAGS_gpu_allocator_retry_time = 500;
+  FLAGS_fraction_of_cuda_pinned_memory_to_use = 0.5;
+#endif
+
+  FLAGS_allocator_strategy = "auto_growth";
+
+  auto &instance = AllocatorFacade::Instance();
+  size_t size = 1024;
+  phi::Place place;
+
+  {
+    place = phi::CPUPlace();
+    size = 1024;
+    auto cpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(cpu_allocation, nullptr);
+    ASSERT_NE(cpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(cpu_allocation->place(), place);
+    ASSERT_EQ(cpu_allocation->size(), AlignedSize(size, 1024));
+  }
+
+#if (defined PADDLE_WITH_CUDA || defined PADDLE_WITH_HIP)
+  {
+    place = phi::GPUPlace(0);
+    size = 1024;
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(),
+              AlignedSize(size, platform::GpuMinChunkSize()));
+  }
+
+  {
+    // Allocate 2GB gpu memory
+    place = phi::GPUPlace(0);
+    size = 2 * static_cast<size_t>(1 << 30);
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(),
+              AlignedSize(size, platform::GpuMinChunkSize()));
+  }
+
+  {
+    place = phi::GPUPinnedPlace();
+    size = (1 << 20);
+    auto cuda_pinned_allocation =
+        instance.Alloc(phi::GPUPinnedPlace(), 1 << 20);
+    ASSERT_NE(cuda_pinned_allocation, nullptr);
+    ASSERT_NE(cuda_pinned_allocation->ptr(), nullptr);
+    ASSERT_EQ(cuda_pinned_allocation->place(), place);
+    ASSERT_GE(cuda_pinned_allocation->size(), AlignedSize(size, 1 << 20));
+  }
+#endif
+}
+
+TEST(multithread_allocate, test_segfault) {
+  FLAGS_allocator_strategy = "auto_growth";
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  std::mutex mtx;
+  std::condition_variable cv;
+  bool flag = false;
+
+  auto alloc_func = [&](int dev_id, unsigned int seed) {
+    auto &instance = AllocatorFacade::Instance();
+
+    std::mt19937 gen(seed);
+    std::uniform_int_distribution<size_t> dist(1 << 20, 1 << 25);
+
+    {
+      std::unique_lock<std::mutex> lock(mtx);
+      cv.wait(lock, [&] { return flag; });
+    }
+
+    for (int i = 0; i < 50; i++) {
+      size_t size = dist(gen);
+      for (int j = 0; j < 10; j++) {
+        instance.Alloc(phi::GPUPlace(dev_id), size);
+      }
+    }
+  };
+
+  std::vector<std::thread> ths;
+  for (size_t i = 0; i < 50; ++i) {
+    std::random_device rd;
+    ths.emplace_back(alloc_func, 0, rd());
+  }
+
+  {
+    std::lock_guard<std::mutex> guard(mtx);
+    flag = true;
+  }
+  cv.notify_all();
+
+  for (auto &th : ths) {
+    th.join();
+  }
+#endif
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/auto_growth_best_fit_allocator_test.cc

@@ -0,0 +1,176 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.h"
+
+#include <cstdlib>
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/allocation/aligned_allocator.h"
+
+PD_DECLARE_bool(free_idle_chunk);
+PD_DECLARE_bool(free_when_no_cache_hit);
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class RecordedAllocator : public Allocator {
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override {
+    allocated_size_ += size;
+    return new Allocation(malloc(size), size, phi::CPUPlace());  // NOLINT
+  }
+
+  void FreeImpl(phi::Allocation *allocation) override {
+    allocated_size_ -= allocation->size();
+    free(allocation->ptr());  // NOLINT
+    delete allocation;
+  }
+
+ public:
+  size_t AllocatedSize() const { return allocated_size_; }
+
+ private:
+  size_t allocated_size_{0};
+};
+
+static void TestFreeIdleChunk(bool free_idle_chunk,
+                              bool free_when_no_cache_hit) {
+  FLAGS_free_idle_chunk = free_idle_chunk;
+  FLAGS_free_when_no_cache_hit = free_when_no_cache_hit;
+  auto recorded_allocator = std::make_shared<RecordedAllocator>();
+
+  size_t alignment = 4096;
+  size_t memory_size = 8192;
+  auto underlying_allocator =
+      std::make_shared<AlignedAllocator>(recorded_allocator, alignment);
+  auto ag_allocator = std::make_shared<AutoGrowthBestFitAllocator>(
+      underlying_allocator, alignment);
+
+  for (size_t i = 0; i < 10; ++i) {
+    auto allocation = ag_allocator->Allocate(memory_size);
+    ASSERT_EQ(recorded_allocator->AllocatedSize(), memory_size + alignment);
+    allocation.reset();
+    if (free_idle_chunk) {
+      ASSERT_EQ(recorded_allocator->AllocatedSize(), 0UL);
+    } else {
+      ASSERT_EQ(recorded_allocator->AllocatedSize(), memory_size + alignment);
+    }
+    ag_allocator->Release(phi::CPUPlace());
+  }
+}
+
+class LimitedResourceAllocator : public Allocator {
+ public:
+  explicit LimitedResourceAllocator(size_t capacity) : capacity_(capacity) {}
+
+  size_t AllocatedSize() const { return allocated_size_; }
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override {
+    if (allocated_size_ + size > capacity_) {
+      throw BadAlloc("", __FILE__, __LINE__);
+    }
+
+    allocated_size_ += size;
+    return new Allocation(malloc(size), size, phi::CPUPlace());  // NOLINT
+  }
+
+  void FreeImpl(phi::Allocation *allocation) override {
+    allocated_size_ -= allocation->size();
+    free(allocation->ptr());  // NOLINT
+    delete allocation;
+  }
+
+ private:
+  size_t allocated_size_{0};
+  const size_t capacity_;
+};
+
+void TestFreeWhenNoCacheHit(bool free_when_no_cache_hit) {
+  FLAGS_free_idle_chunk = false;
+  FLAGS_free_when_no_cache_hit = free_when_no_cache_hit;
+  size_t alignment = 256;
+  size_t base_memory_size = 4096;
+
+  /*
+   * Suppose that we have 3 memory allocation request, that is:
+   *  - allocate x1, and then free x1
+   *  - allocate x2, and then free x2
+   *  - allocate x3, and then free x3
+   *
+   * where:
+   *  - x1 + alignment < x2
+   *  - x2 + alignment < x3
+   *  - x1 + x2 <= memory_capacity < x1 + x2 + x3
+   *
+   * In this unittest, we obtain memory_capacity by
+   * ((x1 + x2) + (x1 + x2 + x3) / 2 = x1 + x2 + x3 / 2.
+   *
+   * In this case, when FLAGS_free_when_no_cache_hit is true,
+   * the cached memory size when each allocation request ends
+   * would be: x1 + alignment, x2 + alignment, x3 + alignment.
+   *
+   * When FLAGS_free_when_no_cache_hit is false, the cached
+   * memory size when each allocation request ends would be:
+   * x1 + alignment, x1 + x2 + 2 * alignment, x3 + alignment.
+   */
+  std::vector<size_t> allocate_size = {base_memory_size,
+                                       base_memory_size + alignment * 2,
+                                       base_memory_size + alignment * 4};
+  size_t memory_capacity =
+      allocate_size[0] + allocate_size[1] + allocate_size[2] / 2;
+
+  auto underlying_allocator =
+      std::make_shared<LimitedResourceAllocator>(memory_capacity);
+  auto aligned_allocator =
+      std::make_shared<AlignedAllocator>(underlying_allocator, alignment);
+  auto ag_allocator = std::make_shared<AutoGrowthBestFitAllocator>(
+      aligned_allocator, alignment);
+
+  ag_allocator->Allocate(allocate_size[0]);
+  ASSERT_EQ(underlying_allocator->AllocatedSize(),
+            allocate_size[0] + alignment);
+
+  ag_allocator->Allocate(allocate_size[1]);
+  if (free_when_no_cache_hit) {
+    ASSERT_EQ(underlying_allocator->AllocatedSize(),
+              allocate_size[1] + alignment);
+  } else {
+    ASSERT_EQ(underlying_allocator->AllocatedSize(),
+              allocate_size[0] + allocate_size[1] + 2 * alignment);
+  }
+
+  ag_allocator->Allocate(allocate_size[2]);
+  ASSERT_EQ(underlying_allocator->AllocatedSize(),
+            allocate_size[2] + alignment);
+}
+
+TEST(test_auto_growth_allocator, test_free_idle_chunk) {
+  for (auto free_idle_chunk : {false, true}) {
+    for (auto free_when_no_cache_hit : {false, true}) {
+      TestFreeIdleChunk(free_idle_chunk, free_when_no_cache_hit);
+    }
+  }
+}
+
+TEST(test_auto_growth_allocator, test_free_when_no_cache_hit) {
+  TestFreeWhenNoCacheHit(false);
+  TestFreeWhenNoCacheHit(true);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/best_fit_allocator_test.cc

@@ -0,0 +1,140 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/best_fit_allocator.h"
+
+#include <random>
+#include <thread>  // NOLINT
+
+#include "gtest/gtest-message.h"
+#include "gtest/gtest-test-part.h"
+#include "gtest/gtest.h"
+#include "gtest/gtest_pred_impl.h"
+#include "paddle/phi/core/memory/allocation/cpu_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class StubAllocation : public Allocation {
+ public:
+  explicit StubAllocation(size_t size)
+      : Allocation(nullptr, size, phi::CPUPlace()) {}
+};
+
+TEST(BestFitAllocator, test_allocation) {
+  // NOTE(zhiqiu): On windows with msvc compiler, unsigned long (UL) is 32bits,
+  // so 4UL * 1024 * 1024 * 1024 becomes 0.
+  // We need to use 4ULL (unsigned long long) here.
+  StubAllocation stub(4ULL * 1024 * 1024 * 1024);
+  BestFitAllocator allocator(&stub);
+  { auto allocation = allocator.Allocate(64); }
+
+  {
+    auto allocation = allocator.Allocate(80);
+
+    {
+      auto best_fit_allocation =
+          dynamic_cast<BestFitAllocation*>(allocation.get());
+      ASSERT_NE(best_fit_allocation, nullptr);
+      ASSERT_FALSE(best_fit_allocation->ChunkIterator()->is_free);
+      ASSERT_EQ(best_fit_allocation->ChunkIterator()->offset_, 0UL);
+      ASSERT_EQ(allocation->size(), 80UL);
+      ASSERT_EQ(allocation->ptr(), nullptr);
+    }
+
+    auto allocation2 = allocator.Allocate(60);
+    auto allocation3 = allocator.Allocate(90);
+    allocation2.reset();
+    allocation2 = allocator.Allocate(30);
+
+    {
+      auto best_fit_allocation =
+          dynamic_cast<BestFitAllocation*>(allocation2.get());
+      ASSERT_EQ(best_fit_allocation->ChunkIterator()->offset_, 80UL);
+    }
+    allocation2.reset();
+    allocation2 = allocator.Allocate(60);
+
+    {
+      auto best_fit_allocation =
+          dynamic_cast<BestFitAllocation*>(allocation2.get());
+      ASSERT_EQ(best_fit_allocation->ChunkIterator()->offset_, 80UL);
+    }
+
+    allocation.reset();
+    allocation2.reset();
+
+    allocation = allocator.Allocate(80 + 60);
+    {
+      auto best_fit_allocation =
+          dynamic_cast<BestFitAllocation*>(allocation.get());
+      ASSERT_EQ(best_fit_allocation->ChunkIterator()->offset_, 0UL);
+    }
+
+    allocation.reset();
+
+    allocation = allocator.Allocate(80);
+    allocation2 = allocator.Allocate(60);
+    allocation = nullptr;
+    allocation2 = nullptr;
+    allocation3 = nullptr;
+
+    ASSERT_EQ(allocator.NumFreeChunks(), 1U);
+  }
+}
+
+TEST(BestFitAllocator, test_concurrent_cpu_allocation) {
+  CPUAllocator allocator;
+  auto global_allocation = allocator.Allocate(256UL * 1024 * 1024);
+
+  BestFitAllocator best_fit_allocator(global_allocation.get());
+
+  auto th_main = [&](std::random_device::result_type seed) {
+    std::default_random_engine engine(seed);
+    std::uniform_int_distribution<size_t> dist(1U, 1024U);
+
+    for (size_t i = 0; i < 128; ++i) {
+      size_t allocate_size = dist(engine);
+
+      auto allocation =
+          best_fit_allocator.Allocate(sizeof(size_t) * allocate_size);
+
+      size_t* data = reinterpret_cast<size_t*>(allocation->ptr());
+
+      for (size_t j = 0; j < allocate_size; ++j) {
+        data[j] = j;
+      }
+      std::this_thread::yield();
+
+      for (size_t j = 0; j < allocate_size; ++j) {
+        ASSERT_EQ(data[j], j);
+      }
+    }
+  };
+  {
+    std::vector<std::thread> threads;
+    for (size_t i = 0; i < 1024; ++i) {
+      std::random_device dev;
+      threads.emplace_back(th_main, dev());
+    }
+    for (auto& th : threads) {
+      th.join();
+    }
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/best_fit_allocator_test.cu

@@ -0,0 +1,97 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <memory>
+#include <random>
+#include <thread>  // NOLINT
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/allocation/allocator_facade.h"
+#include "paddle/phi/core/memory/allocation/best_fit_allocator.h"
+#include "paddle/phi/core/memory/allocation/cuda_allocator.h"
+#include "paddle/phi/core/memory/memcpy.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+struct ForEachFill {
+  size_t* ptr_;
+
+  explicit ForEachFill(size_t* ptr) : ptr_(ptr) {}
+
+  __device__ void operator()(size_t i) { ptr_[i] = i; }
+};
+
+TEST(BestFitAllocator, concurrent_cuda) {
+  CUDAAllocator allocator(phi::GPUPlace(0));
+  // 256 MB
+  auto cuda_allocation = allocator.Allocate(256U * 1024 * 1024);
+  BestFitAllocator concurrent_allocator(cuda_allocation.get());
+
+  phi::GPUPlace gpu(0);
+  phi::GPUContext dev_ctx(gpu);
+  dev_ctx.SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
+                           .GetAllocator(gpu, dev_ctx.stream())
+                           .get());
+  dev_ctx.PartialInitWithAllocator();
+
+  auto th_main = [&](std::random_device::result_type seed) {
+    std::default_random_engine engine(seed);
+    std::uniform_int_distribution<size_t> dist(1U, 1024U);
+    std::array<size_t, 1024> buf;
+
+    for (size_t i = 0; i < 128; ++i) {
+      size_t allocate_size = dist(engine);
+
+      auto allocation =
+          concurrent_allocator.Allocate(sizeof(size_t) * allocate_size);
+
+      size_t* data = reinterpret_cast<size_t*>(allocation->ptr());
+
+      ForEachFill fill(data);
+      phi::funcs::ForRange<phi::GPUContext> for_range(dev_ctx, allocate_size);
+      for_range(fill);
+
+      memory::Copy(phi::CPUPlace(),
+                   buf.data(),
+                   gpu,
+                   data,
+                   sizeof(size_t) * allocate_size,
+                   dev_ctx.stream());
+
+      dev_ctx.Wait();
+      for (size_t j = 0; j < allocate_size; ++j) {
+        ASSERT_EQ(buf[j], j);
+      }
+      allocation = nullptr;
+    }
+  };
+
+  {
+    std::vector<std::thread> threads;
+    for (size_t i = 0; i < 1024; ++i) {
+      std::random_device dev;
+      threads.emplace_back(th_main, dev());
+    }
+    for (auto& th : threads) {
+      th.join();
+    }
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/buffered_allocator_test.cc

@@ -0,0 +1,129 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/buffered_allocator.h"
+
+#include <gtest/gtest.h>
+
+#include <utility>
+
+#include "paddle/phi/core/memory/allocation/best_fit_allocator.h"
+#include "paddle/phi/core/memory/allocation/cpu_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class StubAllocation : public Allocation {
+ public:
+  using Allocation::Allocation;
+};
+
+class StubAllocator : public Allocator {
+ public:
+  void ResetCounter() {
+    construct_count_ = 0;
+    destruct_count_ = 0;
+  }
+
+  size_t GetAllocCount() const { return construct_count_; }
+
+  size_t GetFreeCount() const { return destruct_count_; }
+
+ protected:
+  void FreeImpl(phi::Allocation *allocation) override {
+    auto *alloc = dynamic_cast<StubAllocation *>(allocation);
+    PADDLE_ENFORCE_NOT_NULL(
+        alloc,
+        common::errors::InvalidArgument(
+            "The input allocation is not type of StubAllocation."));
+    if (alloc->ptr()) delete[] static_cast<uint8_t *>(alloc->ptr());
+    ++destruct_count_;
+    delete allocation;
+  }
+
+  phi::Allocation *AllocateImpl(size_t size) override {
+    ++construct_count_;
+    if (size == 0) {
+      return new StubAllocation(nullptr, 0, phi::CPUPlace());
+    } else {
+      return new StubAllocation(new uint8_t[size], size, phi::CPUPlace());
+    }
+  }
+
+ private:
+  size_t construct_count_ = 0;
+  size_t destruct_count_ = 0;
+};
+
+constexpr size_t kZero = 0;
+constexpr size_t kOne = 1;
+constexpr size_t kTwo = 2;
+
+TEST(buffered_allocator, lazy_free) {
+  std::unique_ptr<StubAllocator> stub_allocator(new StubAllocator());
+  auto *underlying_allocator = stub_allocator.get();
+  std::unique_ptr<BufferedAllocator> allocator(
+      new BufferedAllocator(std::move(stub_allocator)));
+
+  {
+    underlying_allocator->ResetCounter();
+    auto x = allocator->Allocate(1025);
+    ASSERT_EQ(underlying_allocator->GetAllocCount(), kOne);
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+    x = nullptr;
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+  }
+
+  {
+    underlying_allocator->ResetCounter();
+    auto x = allocator->Allocate(900);
+    ASSERT_EQ(underlying_allocator->GetAllocCount(), kZero);
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+    auto y = allocator->Allocate(2048);
+    ASSERT_EQ(underlying_allocator->GetAllocCount(), kOne);
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+    x = nullptr;
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+    y = nullptr;
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+  }
+
+  {
+    underlying_allocator->ResetCounter();
+    allocator->ClearCache();
+    ASSERT_EQ(underlying_allocator->GetAllocCount(), kZero);
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kTwo);
+  }
+}
+
+TEST(buffered_allocator, garbage_collection) {
+  std::unique_ptr<CPUAllocator> cpu_allocator(new CPUAllocator());
+  auto chunk = cpu_allocator->Allocate(2048);
+  std::unique_ptr<Allocator> allocator(new BestFitAllocator(chunk.get()));
+
+  auto buffered_allocator =
+      std::make_unique<BufferedAllocator>(std::move(allocator));
+  auto x1 = buffered_allocator->Allocate(1600);
+  auto x2 = buffered_allocator->Allocate(400);
+  x1 = nullptr;
+  x2 = nullptr;
+  auto x3 = buffered_allocator->Allocate(1600);
+  ASSERT_NE(x3, nullptr);
+  ASSERT_NE(x3->ptr(), nullptr);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/cuda_malloc_async_allocator_test.cu

@@ -0,0 +1,381 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <thread>  // NOLINT
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/allocation/allocator_facade.h"
+#include "paddle/phi/core/memory/memory.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#include "paddle/phi/core/platform/device_context.h"
+#include "paddle/phi/core/stream.h"
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include "paddle/phi/core/platform/cuda_graph_with_memory_pool.h"
+#endif
+
+namespace paddle {
+namespace memory {
+
+// y += (x + 1)
+__global__ void add_kernel(int *x, int *y, int n) {
+  int thread_num = gridDim.x * blockDim.x;
+  int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int i = thread_id; i < n; i += thread_num) {
+    y[i] += x[i] + 1;
+  }
+}
+
+void CheckMemLeak(const phi::GPUPlace &place) {
+  uint64_t cuda_malloc_size =
+      platform::RecordedGpuMallocSize(place.GetDeviceId());
+  ASSERT_EQ(cuda_malloc_size, 0)
+      << "Found " << cuda_malloc_size << " bytes memory that not released yet,"
+      << " there may be a memory leak problem";
+}
+
+#if (defined(PADDLE_WITH_CUDA) && (CUDA_VERSION >= 12200))
+
+TEST(CUDAMallocAsyncInterfaceTest, AllocInterfaceTest) {
+  phi::GPUPlace place = phi::GPUPlace();
+  size_t alloc_size = 256;
+
+  std::shared_ptr<Allocation> allocation_implicit_stream =
+      AllocShared(place, alloc_size);
+  EXPECT_GE(allocation_implicit_stream->size(), alloc_size);
+
+  void *address = allocation_implicit_stream->ptr();
+  allocation_implicit_stream.reset();
+  gpuStream_t default_stream =
+      dynamic_cast<phi::GPUContext *>(
+          phi::DeviceContextPool::Instance().Get(place))
+          ->stream();
+  allocation::AllocationPtr allocation_unique =
+      Alloc(place,
+            alloc_size,
+            phi::Stream(reinterpret_cast<phi::StreamId>(default_stream)));
+  EXPECT_GE(allocation_unique->size(), alloc_size);
+  EXPECT_EQ(allocation_unique->ptr(), address);
+  allocation_unique.reset();
+
+  Release(place);
+  CheckMemLeak(place);
+}
+
+TEST(CUDAMallocAsyncInterfaceTest, GetAllocatorInterfaceTest) {
+  phi::GPUPlace place = phi::GPUPlace();
+  size_t alloc_size = 256;
+
+  allocation::AllocationPtr allocation_implicit_stream =
+      Alloc(place, alloc_size);
+  EXPECT_GE(allocation_implicit_stream->size(), alloc_size);
+  void *address = allocation_implicit_stream->ptr();
+  allocation_implicit_stream.reset();
+
+  auto &instance = allocation::AllocatorFacade::Instance();
+  const std::shared_ptr<Allocator> &allocator = instance.GetAllocator(place);
+
+  allocation::AllocationPtr allocation_from_allocator =
+      allocator->Allocate(alloc_size);
+  EXPECT_GE(allocation_from_allocator->size(), alloc_size);
+  EXPECT_EQ(allocation_from_allocator->ptr(), address);
+  allocation_from_allocator.reset();
+
+  Release(place);
+  CheckMemLeak(place);
+}
+
+TEST(CUDAMallocAsyncInterfaceTest, GetAllocatorWithDefaultStreamTest) {
+  auto &instance = allocation::AllocatorFacade::Instance();
+  phi::GPUPlace place = phi::GPUPlace();
+  const std::shared_ptr<Allocator> allocator_implicit_stream =
+      instance.GetAllocator(place);
+  const std::shared_ptr<Allocator> allocator_default_stream =
+      instance.GetAllocator(place,
+                            static_cast<phi::GPUContext *>(
+                                phi::DeviceContextPool::Instance().Get(place))
+                                ->stream());
+  EXPECT_EQ(allocator_implicit_stream.get(), allocator_default_stream.get());
+}
+
+TEST(CUDAMallocAsyncInterfaceTest, ZeroSizeRecordStreamTest) {
+  phi::GPUPlace place = phi::GPUPlace();
+  std::shared_ptr<Allocation> zero_size_allocation = AllocShared(place, 0);
+  EXPECT_EQ(zero_size_allocation->ptr(), nullptr);
+
+  gpuStream_t stream;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream));
+
+  EXPECT_NO_THROW(RecordStream(zero_size_allocation, stream));
+
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(stream));
+}
+
+TEST(CUDAMallocAsyncInterfaceTest, GetStreamInterfaceTest) {
+  phi::GPUPlace place = phi::GPUPlace();
+  size_t alloc_size = 256;
+
+  gpuStream_t default_stream =
+      dynamic_cast<phi::GPUContext *>(
+          phi::DeviceContextPool::Instance().Get(place))
+          ->stream();
+  std::shared_ptr<Allocation> allocation_implicit_stream =
+      AllocShared(place, alloc_size);
+  EXPECT_EQ(GetStream(allocation_implicit_stream), default_stream);
+
+  gpuStream_t new_stream;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&new_stream));
+
+  std::shared_ptr<Allocation> allocation_new_stream =
+      AllocShared(place,
+                  alloc_size,
+                  phi::Stream(reinterpret_cast<phi::StreamId>(new_stream)));
+  EXPECT_EQ(GetStream(allocation_new_stream), new_stream);
+
+  allocation_implicit_stream.reset();
+  allocation_new_stream.reset();
+  // we should release the block before destroy the stream
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(new_stream));
+
+  Release(place);
+  CheckMemLeak(place);
+}
+
+TEST(CUDAMallocAsyncRetryTest, RetryTest) {
+  phi::GPUPlace place = phi::GPUPlace();
+  gpuStream_t stream1, stream2;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream1));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream2));
+  size_t available_size = platform::GpuAvailableMemToAlloc();
+  // alloc_size < available_size < 2 * alloc_size,
+  // so the second alloc will fail and retry
+  size_t alloc_size = available_size / 4 * 3;
+
+  allocation::AllocationPtr allocation1 = Alloc(
+      place, alloc_size, phi::Stream(reinterpret_cast<phi::StreamId>(stream1)));
+  allocation::AllocationPtr allocation2;
+
+  std::thread th([&allocation2, &place, &stream2, alloc_size]() {
+    std::this_thread::sleep_for(std::chrono::seconds(1));
+    allocation2 = Alloc(place,
+                        alloc_size,
+                        phi::Stream(reinterpret_cast<phi::StreamId>(stream2)));
+  });
+  allocation1.reset();  // free but not release
+  th.join();
+  EXPECT_GE(allocation2->size(), alloc_size);
+  allocation2.reset();
+
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
+
+  Release(place, stream1);
+  Release(place, stream2);
+  CheckMemLeak(place);
+}
+
+class CUDAMallocAsyncTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    place_ = phi::GPUPlace();
+    stream_num_ = 64;
+    grid_num_ = 1;
+    block_num_ = 32;
+    data_num_ = 131072;
+    workspace_size_ = data_num_ * sizeof(int);
+
+    for (size_t i = 0; i < stream_num_; ++i) {
+      gpuStream_t stream;
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream));
+
+      std::shared_ptr<phi::Allocation> workspace_allocation =
+          AllocShared(place_,
+                      workspace_size_,
+                      phi::Stream(reinterpret_cast<phi::StreamId>(stream)));
+      std::shared_ptr<phi::Allocation> result_allocation =
+          AllocShared(place_,
+                      workspace_size_,
+                      phi::Stream(reinterpret_cast<phi::StreamId>(stream)));
+      std::shared_ptr<phi::Allocation> host_result_allocation =
+          AllocShared(phi::CPUPlace(), workspace_size_);
+
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaMemset(
+          workspace_allocation->ptr(), 0, workspace_allocation->size()));
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaMemset(result_allocation->ptr(), 0, result_allocation->size()));
+
+      streams_.emplace_back(stream);
+      workspaces_.emplace_back(workspace_allocation);
+      results_.emplace_back(result_allocation);
+      host_results_.emplace_back(host_result_allocation);
+    }
+  }
+
+  void SingleStreamRun(size_t idx) {
+    int *y = reinterpret_cast<int *>(results_[idx]->ptr());
+    int neighbouring_idx = idx > 0 ? idx - 1 : idx;
+
+    add_kernel<<<grid_num_, block_num_, 0, streams_[idx]>>>(
+        reinterpret_cast<int *>(workspaces_[idx]->ptr()), y, data_num_);
+    add_kernel<<<grid_num_, block_num_, 0, streams_[idx]>>>(
+        reinterpret_cast<int *>(workspaces_[neighbouring_idx]->ptr()),
+        y,
+        data_num_);
+    RecordStream(workspaces_[neighbouring_idx], streams_[idx]);
+  }
+
+  void MultiStreamRun() {
+    // Must run in reverse order, or the workspace_[i - 1] will be released
+    // before streams_[i]'s kernel launch
+    for (int i = stream_num_ - 1; i >= 0; --i) {
+      SingleStreamRun(i);
+      workspaces_[i].reset();  // fast GC
+    }
+  }
+
+  void MultiThreadMultiStreamRun() {
+    std::vector<std::thread> threads;
+    for (size_t i = 0; i < stream_num_; ++i) {
+      threads.emplace_back(&CUDAMallocAsyncTest::SingleStreamRun, this, i);
+    }
+    for (size_t i = 0; i < stream_num_; ++i) {
+      threads[i].join();
+    }
+    workspaces_.clear();
+  }
+
+  void CUDAGraphRun() {
+    testing_cuda_graph_ = true;
+    platform::BeginCUDAGraphCapture(phi::GPUPlace(),
+                                    cudaStreamCaptureModeGlobal);
+
+    std::shared_ptr<Allocation> data_allocation =
+        AllocShared(phi::GPUPlace(), workspace_size_);
+    std::shared_ptr<Allocation> result_allocation =
+        AllocShared(phi::GPUPlace(), workspace_size_);
+
+    int *data = static_cast<int *>(data_allocation->ptr());
+    int *result = static_cast<int *>(result_allocation->ptr());
+
+    gpuStream_t main_stream = GetStream(data_allocation);
+    gpuStream_t other_stream;
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&other_stream));
+
+    add_kernel<<<grid_num_, block_num_, 0, main_stream>>>(
+        data, result, data_num_);
+    RecordStream(data_allocation, other_stream);
+
+    std::unique_ptr<phi::backends::gpu::CUDAGraph> cuda_graph =
+        platform::EndCUDAGraphCapture();
+
+    int replay_times = 10;
+    for (int i = 0; i < replay_times; ++i) {
+      cuda_graph->Replay();
+    }
+
+    std::shared_ptr<Allocation> host_result_allocation =
+        AllocShared(phi::CPUPlace(), workspace_size_);
+    Copy(host_result_allocation->place(),
+         host_result_allocation->ptr(),
+         result_allocation->place(),
+         result_allocation->ptr(),
+         workspace_size_,
+         main_stream);
+    cudaStreamSynchronize(main_stream);
+
+    int *host_result = static_cast<int *>(host_result_allocation->ptr());
+    for (int i = 0; i < data_num_; ++i) {
+      EXPECT_EQ(host_result[i], replay_times);
+    }
+
+    data_allocation.reset();
+    result_allocation.reset();
+    cuda_graph.release();
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(other_stream));
+  }
+
+  void CheckResult() {
+    for (size_t i = 0; i < stream_num_; ++i) {
+      Copy(host_results_[i]->place(),
+           host_results_[i]->ptr(),
+           results_[i]->place(),
+           results_[i]->ptr(),
+           workspace_size_,
+           streams_[i]);
+    }
+    cudaDeviceSynchronize();
+
+    size_t thread_num = grid_num_ * block_num_;
+    for (size_t i = 0; i < stream_num_; ++i) {
+      int *result = static_cast<int *>(host_results_[i]->ptr());
+      for (size_t j = 0; j < data_num_; ++j) {
+        EXPECT_EQ(result[j], 2);
+      }
+    }
+  }
+
+  void TearDown() override {
+    workspaces_.clear();
+    results_.clear();
+    host_results_.clear();
+    for (gpuStream_t stream : streams_) {
+      Release(place_, stream);
+    }
+
+    for (size_t i = 0; i < stream_num_; ++i) {
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamSynchronize(streams_[i]));
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(streams_[i]));
+    }
+
+    // Memory release for CUDA Graph memory pool is forbidden
+    if (!testing_cuda_graph_) {
+      CheckMemLeak(place_);
+    }
+  }
+
+  bool testing_cuda_graph_{0};
+  size_t stream_num_;
+  size_t grid_num_;
+  size_t block_num_;
+  size_t data_num_;
+  size_t workspace_size_;
+  phi::GPUPlace place_;
+  std::vector<gpuStream_t> streams_;
+  std::vector<std::shared_ptr<phi::Allocation>> workspaces_;
+  std::vector<std::shared_ptr<phi::Allocation>> results_;
+  std::vector<std::shared_ptr<phi::Allocation>> host_results_;
+};
+
+TEST_F(CUDAMallocAsyncTest, CUDAMutilStreamTest) {
+  MultiStreamRun();
+  CheckResult();
+}
+
+TEST_F(CUDAMallocAsyncTest, CUDAMutilThreadMutilStreamTest) {
+  MultiThreadMultiStreamRun();
+  CheckResult();
+}
+
+TEST_F(CUDAMallocAsyncTest, CUDAGraphTest) {
+  MultiStreamRun();
+  CUDAGraphRun();
+  CheckResult();
+}
+#endif
+
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/cuda_managed_memory_test.cu

@@ -0,0 +1,137 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda_runtime.h>
+#endif
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#endif
+
+#include "gtest/gtest.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/malloc.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+namespace paddle {
+namespace memory {
+
+__global__ void write_kernel(int* data, uint64_t n, uint64_t step) {
+  int thread_num = gridDim.x * blockDim.x;
+  int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
+  for (uint64_t i = thread_id; i * step < n; i += thread_num) {
+    *(data + i * step) = 1;
+  }
+}
+
+__global__ void sum_kernel(int* data, uint64_t n, uint64_t step, int* sum) {
+  int thread_num = gridDim.x * blockDim.x;
+  int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
+  for (uint64_t i = thread_id; i * step < n; i += thread_num) {
+    atomicAdd(sum, *(data + i * step));
+  }
+}
+
+TEST(ManagedMemoryTest, H2DTest) {
+  if (!platform::IsGPUManagedMemorySupported(0)) {
+    return;
+  }
+
+  uint64_t n_data = 1024;
+  uint64_t step = 1;
+  allocation::AllocationPtr allocation =
+      Alloc(phi::GPUPlace(0), n_data * sizeof(int));
+  int* data = static_cast<int*>(allocation->ptr());
+
+  memset(data, 0, n_data * sizeof(int));          // located on host memory
+  write_kernel<<<1, 1024>>>(data, n_data, step);  // trans to device memory
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceSynchronize());
+#endif
+
+  int sum = 0;
+  for (uint64_t i = 0; i < n_data; ++i) {
+    sum += *(data + i);
+  }
+  EXPECT_EQ(sum, n_data / step);
+  allocation = nullptr;
+}
+
+TEST(ManagedMemoryTest, D2HTest) {
+  if (!platform::IsGPUManagedMemorySupported(0)) {
+    return;
+  }
+
+  uint64_t n_data = 1024;
+  uint64_t step = 1;
+  AllocationPtr allocation = Alloc(phi::GPUPlace(0), n_data * sizeof(int));
+  int* data = static_cast<int*>(allocation->ptr());
+
+  write_kernel<<<1, 1024>>>(data, n_data, step);  // located on device memory
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceSynchronize());
+#endif
+
+  memset(data, 0, n_data * sizeof(int));  // trans to host memory
+
+  int sum = 0;
+  for (uint64_t i = 0; i < n_data; ++i) {
+    sum += *(data + i);
+  }
+  EXPECT_EQ(sum, 0);
+}
+
+TEST(ManagedMemoryTest, OversubscribeGPUMemoryTest) {
+  if (!platform::IsGPUManagedMemoryOversubscriptionSupported(0)) {
+    return;
+  }
+
+  uint64_t available_mem = platform::GpuAvailableMemToAlloc();
+  uint64_t n_data = available_mem * 2 / sizeof(int) +
+                    1;  // requires more than 2 * available_mem bytes
+  uint64_t step = std::max(n_data / 1024, static_cast<uint64_t>(1));
+  AllocationPtr data_allocation = Alloc(phi::GPUPlace(0), n_data * sizeof(int));
+  AllocationPtr sum_allocation = Alloc(phi::GPUPlace(0), sizeof(int));
+  int* data = static_cast<int*>(data_allocation->ptr());
+  int* sum = static_cast<int*>(sum_allocation->ptr());
+  (*sum) = 0;
+
+  write_kernel<<<1, 1024>>>(data, n_data, step);
+  sum_kernel<<<1, 1024>>>(data, n_data, step, sum);
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceSynchronize());
+#endif
+
+  EXPECT_EQ(*sum, (n_data + step - 1) / step);
+}
+
+TEST(ManagedMemoryTest, OOMExceptionTest) {
+  if (!platform::IsGPUManagedMemorySupported(0)) {
+    return;
+  }
+  EXPECT_THROW(Alloc(phi::GPUPlace(0), size_t(1) << 60),
+               memory::allocation::BadAlloc);
+}
+
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/cuda_virtual_mem_allocator_v2_test.cu

@@ -0,0 +1,362 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "paddle/phi/core/enforce.h"
+#define private public
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator_v2.h"
+#undef private
+#include "paddle/phi/core/memory/allocation/vmm_backing_map.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(VMMBackingMap, TracksMappedAndUnmappedRanges) {
+  VMMBackingMap map;
+  const VMMDevicePtr base = 0x10000000;
+  const size_t page_size = 2UL << 20;
+  map.Configure(base, page_size * 4, page_size, 0);
+
+  EXPECT_TRUE(map.IsRangeUnmapped(base, page_size * 4));
+  EXPECT_FALSE(map.IsRangeMapped(base, page_size));
+  EXPECT_EQ(map.total_mapped_bytes(), 0UL);
+
+  const VMMAllocHandle first_handle = static_cast<VMMAllocHandle>(0x101);
+  const VMMAllocHandle second_handle = static_cast<VMMAllocHandle>(0x102);
+  auto first_meta =
+      std::make_shared<VMMHandleMeta>(base, page_size, first_handle, 0);
+  auto second_meta = std::make_shared<VMMHandleMeta>(
+      base + page_size, page_size, second_handle, 0);
+  map.MarkMapped(base, first_meta, page_size);
+  map.MarkMapped(base + page_size, second_meta, page_size);
+  map.MarkMapped(base, first_meta, page_size);
+  EXPECT_EQ(map.total_mapped_bytes(), page_size * 2);
+  EXPECT_FALSE(map.IsRangeReleasable(base, page_size * 4));
+  EXPECT_FALSE(map.IsRangeReleasable(base - page_size, page_size));
+
+  EXPECT_TRUE(map.IsRangeMapped(base, page_size * 2));
+  EXPECT_FALSE(map.IsRangeMapped(base, page_size * 3));
+  EXPECT_FALSE(map.IsRangeUnmapped(base, page_size));
+  EXPECT_TRUE(map.IsRangeUnmapped(base + page_size * 2, page_size * 2));
+  EXPECT_EQ(map.total_mapped_bytes(), page_size * 2);
+
+  map.MarkUnmapped(base, page_size);
+  map.MarkUnmapped(base, page_size);
+
+  EXPECT_FALSE(map.IsRangeMapped(base, page_size * 2));
+  EXPECT_TRUE(map.IsRangeUnmapped(base, page_size));
+  EXPECT_TRUE(map.IsRangeMapped(base + page_size, page_size));
+  EXPECT_EQ(map.total_mapped_bytes(), page_size);
+
+  map.MarkReleased(base + page_size, second_handle, page_size);
+  map.MarkReleased(base + page_size, second_handle, page_size);
+  EXPECT_TRUE(map.IsRangeUnmapped(base, page_size * 4));
+  EXPECT_EQ(map.total_mapped_bytes(), 0UL);
+  const VMMAllocHandle third_handle = static_cast<VMMAllocHandle>(0x103);
+  map.MarkMapped(base, third_handle, page_size);
+  EXPECT_TRUE(map.IsRangeMapped(base, page_size));
+}
+
+TEST(VMMBackingMap, RejectsMappedPageHandleOverwrite) {
+  VMMBackingMap map;
+  const VMMDevicePtr base = 0x18000000;
+  const size_t page_size = 2UL << 20;
+  map.Configure(base, page_size, page_size, 0);
+
+  const VMMAllocHandle first_handle = static_cast<VMMAllocHandle>(0x181);
+  const VMMAllocHandle second_handle = static_cast<VMMAllocHandle>(0x182);
+  map.MarkMapped(base, first_handle, page_size);
+  EXPECT_THROW(map.MarkMapped(base, second_handle, page_size),
+               common::enforce::EnforceNotMet);
+
+  map.MarkUnmapped(base, page_size);
+  auto meta = std::make_shared<VMMHandleMeta>(base, page_size, first_handle, 0);
+  map.MarkMapped(base, meta, page_size);
+  auto other_meta =
+      std::make_shared<VMMHandleMeta>(base, page_size, second_handle, 0);
+  EXPECT_THROW(map.MarkMapped(base, other_meta, page_size),
+               common::enforce::EnforceNotMet);
+}
+
+TEST(VMMBackingMap, RejectsInvalidConfiguration) {
+  VMMBackingMap map;
+  const VMMDevicePtr base = 0x1a000000;
+  const size_t page_size = 2UL << 20;
+
+  EXPECT_THROW(map.Configure(base, page_size, 0, 0),
+               common::enforce::EnforceNotMet);
+  EXPECT_THROW(map.Configure(base, page_size + 1, page_size, 0),
+               common::enforce::EnforceNotMet);
+}
+
+TEST(VMMBackingMap, ReconfigureKeepsOriginalLayout) {
+  VMMBackingMap map;
+  const VMMDevicePtr base = 0x1c000000;
+  const size_t page_size = 2UL << 20;
+
+  map.Configure(base, page_size * 2, page_size, 0);
+  EXPECT_TRUE(map.IsConfigured());
+  map.Configure(base, page_size * 2, page_size, 0);
+  map.Configure(base + page_size, page_size * 2, page_size, 1);
+
+  EXPECT_TRUE(map.IsRangeUnmapped(base, page_size * 2));
+  EXPECT_FALSE(map.IsRangeUnmapped(base + page_size * 2, page_size));
+}
+
+TEST(VMMBackingMap, UnconfiguredAndInvalidRangesReturnFalse) {
+  VMMBackingMap map;
+  const VMMDevicePtr base = 0x1e000000;
+  const size_t page_size = 2UL << 20;
+  auto meta = std::make_shared<VMMHandleMeta>(
+      base, page_size, static_cast<VMMAllocHandle>(0x1e1), 0);
+  HandleLayout layout{meta};
+
+  EXPECT_FALSE(map.IsRangeMapped(base, page_size));
+  EXPECT_FALSE(map.IsRangeUnmapped(base, page_size));
+  EXPECT_FALSE(map.IsRangeReleasable(base, page_size));
+  EXPECT_FALSE(map.ValidateLayout(layout, "unconfigured"));
+  map.MarkMapped(base, meta, page_size);
+  map.MarkUnmapped(base, page_size);
+  map.MarkReleased(base, meta->handle(), page_size);
+
+  map.Configure(base, page_size * 2, page_size, 0);
+  EXPECT_FALSE(map.IsRangeMapped(base, page_size / 2));
+  EXPECT_FALSE(map.IsRangeUnmapped(base - page_size, page_size));
+  EXPECT_FALSE(map.IsRangeReleasable(base + page_size * 2, page_size));
+}
+
+TEST(VMMBackingMap, ValidateLayoutDetectsMissingAndMismatchedPages) {
+  VMMBackingMap map;
+  const VMMDevicePtr base = 0x22000000;
+  const size_t page_size = 2UL << 20;
+  map.Configure(base, page_size * 2, page_size, 0);
+
+  auto first = std::make_shared<VMMHandleMeta>(
+      base, page_size, static_cast<VMMAllocHandle>(0x221), 0);
+  auto second = std::make_shared<VMMHandleMeta>(
+      base + page_size, page_size, static_cast<VMMAllocHandle>(0x222), 0);
+  EXPECT_FALSE(map.ValidateLayout(HandleLayout{first}, "missing mapped page"));
+
+  map.MarkMapped(base, first, page_size);
+  map.MarkMapped(base + page_size, second, page_size);
+  auto wrong = std::make_shared<VMMHandleMeta>(
+      base + page_size, page_size, static_cast<VMMAllocHandle>(0x223), 0);
+  EXPECT_FALSE(
+      map.ValidateLayout(HandleLayout{first, wrong}, "handle mismatch"));
+}
+
+TEST(VMMBackingMap, MarkReleasedAllowsMismatchAndNullMetaIsNotReleasable) {
+  VMMBackingMap map;
+  const VMMDevicePtr base = 0x24000000;
+  const size_t page_size = 2UL << 20;
+  map.Configure(base, page_size * 2, page_size, 0);
+
+  map.MarkMapped(base, static_cast<VMMAllocHandle>(0), page_size);
+  EXPECT_TRUE(map.IsRangeMapped(base, page_size));
+  EXPECT_FALSE(map.IsRangeReleasable(base, page_size));
+  map.MarkUnmapped(base + page_size, page_size);
+
+  auto meta = std::make_shared<VMMHandleMeta>(
+      base + page_size, page_size, static_cast<VMMAllocHandle>(0x241), 0);
+  map.MarkMapped(base + page_size, meta, page_size);
+  map.MarkReleased(
+      base + page_size, static_cast<VMMAllocHandle>(0x242), page_size);
+  EXPECT_TRUE(map.IsRangeUnmapped(base + page_size, page_size));
+}
+
+TEST(VMMBackingMap, PageCanUseBackingRejectsInvalidPageStates) {
+  VMMBackingMap map;
+  VMMBackingMap::Page page;
+
+  EXPECT_FALSE(map.PageCanUseBackingLocked(nullptr, "null page"));
+  EXPECT_FALSE(map.PageCanUseBackingLocked(&page, "unmapped page"));
+  page.mapped = true;
+  EXPECT_FALSE(map.PageCanUseBackingLocked(&page, "missing meta"));
+}
+
+TEST(BlockV2, UnmappedSubBlockTrimAndMerge) {
+  auto* base = reinterpret_cast<void*>(0x20000000);
+  BlockV2 block =
+      BlockV2::MakeUnmappedFreeBlock(base, 8UL << 20, PoolType::kSmall);
+
+  BlockV2 sub_block = block.MakeUnmappedFreeSubBlock(2UL << 20, 4UL << 20);
+  EXPECT_TRUE(sub_block.IsUnmappedFree());
+  EXPECT_EQ(sub_block.ptr(), reinterpret_cast<uint8_t*>(base) + (2UL << 20));
+  EXPECT_EQ(sub_block.size(), 4UL << 20);
+  EXPECT_EQ(sub_block.pool_type_, PoolType::kSmall);
+
+  sub_block.TrimToSuffix(1UL << 20, 3UL << 20);
+  EXPECT_EQ(sub_block.ptr(), reinterpret_cast<uint8_t*>(base) + (3UL << 20));
+  EXPECT_EQ(sub_block.size(), 3UL << 20);
+
+  BlockV2 next = BlockV2::MakeUnmappedFreeBlock(
+      reinterpret_cast<uint8_t*>(base) + (6UL << 20),
+      2UL << 20,
+      PoolType::kSmall);
+  sub_block.MergeAdjacentUnmappedFreeBlock(next);
+  EXPECT_EQ(sub_block.size(), 5UL << 20);
+}
+
+TEST(CUDAVirtualMemAllocatorV2, AppendWithBlockReturnsMappedFreeBlock) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kLarge);
+
+  auto allocation_with_block =
+      allocator.AppendWithBlock(allocator.handle_size() * 2);
+  ASSERT_NE(allocation_with_block.allocation, nullptr);
+
+  const auto& block = allocation_with_block.block;
+  ASSERT_EQ(block.size_, allocation_with_block.allocation->size());
+  EXPECT_EQ(block.ptr_, allocation_with_block.allocation->ptr());
+  EXPECT_TRUE(block.IsMappedFree());
+}
+
+TEST(CUDAVirtualMemAllocatorV2,
+     SmallPoolAppendWithBlockReturnsMappedFreeBlock) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kSmall);
+
+  auto allocation_with_block =
+      allocator.AppendWithBlock(allocator.handle_size());
+  ASSERT_NE(allocation_with_block.allocation, nullptr);
+
+  const auto& block = allocation_with_block.block;
+  EXPECT_EQ(allocator.pool_type(), PoolType::kSmall);
+  EXPECT_EQ(block.ptr_, allocation_with_block.allocation->ptr());
+  EXPECT_EQ(block.size_, allocator.handle_size());
+  EXPECT_TRUE(block.IsMappedFree());
+}
+
+TEST(CUDAVirtualMemAllocatorV2, CollectAllocationHandleLayoutTracksLifecycle) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kLarge);
+
+  HandleLayout layout;
+  EXPECT_FALSE(allocator.CollectAllocationHandleLayout(
+      reinterpret_cast<void*>(0x1234), &layout));
+
+  auto allocation_with_block =
+      allocator.AppendWithBlock(allocator.handle_size() * 2);
+  ASSERT_NE(allocation_with_block.allocation, nullptr);
+  void* ptr = allocation_with_block.allocation->ptr();
+
+  ASSERT_TRUE(allocator.CollectAllocationHandleLayout(ptr, &layout));
+  ASSERT_EQ(layout.size(), 2UL);
+  EXPECT_EQ(layout.front()->base(), reinterpret_cast<VMMDevicePtr>(ptr));
+  EXPECT_EQ(layout.front()->size(), allocator.handle_size());
+
+  layout.clear();
+  allocation_with_block.allocation.reset();
+  EXPECT_FALSE(allocator.CollectAllocationHandleLayout(ptr, &layout));
+}
+
+TEST(CUDAVirtualMemAllocatorV2, AllocateImplReturnsTrackedAllocation) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kLarge);
+
+  auto allocation = allocator.Allocate(allocator.handle_size());
+  ASSERT_NE(allocation, nullptr);
+  EXPECT_NE(allocation->ptr(), nullptr);
+  EXPECT_EQ(allocation->size(), allocator.handle_size());
+
+  HandleLayout layout;
+  EXPECT_TRUE(
+      allocator.CollectAllocationHandleLayout(allocation->ptr(), &layout));
+  EXPECT_EQ(layout.size(), 1UL);
+}
+
+TEST(CUDAVirtualMemAllocatorV2, RollbackCreatedHandlesReleasesLayout) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kLarge);
+
+  allocator.InitOnce();
+  HandleLayout layout;
+  layout.push_back(nullptr);
+  allocator.RollbackCreatedHandles(layout);
+
+  layout = allocator.CreateMappedHandleLayout(
+      allocator.virtual_mem_base(), allocator.handle_size(), "test rollback");
+  ASSERT_EQ(layout.size(), 1UL);
+  allocator.RollbackCreatedHandles(layout);
+}
+
+TEST(CUDAVirtualMemAllocatorV2, RequireHandleLayoutRejectsUnknownPointer) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kLarge);
+
+  EXPECT_THROW(allocator.RequireHandleLayout(reinterpret_cast<void*>(0x1234)),
+               common::enforce::EnforceNotMet);
+}
+
+TEST(CUDAVirtualMemAllocatorV2, PlaceAtVARejectsUnalignedAddress) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kLarge);
+
+  auto allocation_with_block =
+      allocator.AppendWithBlock(allocator.handle_size());
+  ASSERT_NE(allocation_with_block.allocation, nullptr);
+  auto ptr =
+      reinterpret_cast<VMMDevicePtr>(allocation_with_block.allocation->ptr());
+
+  EXPECT_THROW(allocator.PlaceAtVAWithBlock(ptr + 1, allocator.handle_size()),
+               common::enforce::EnforceNotMet);
+}
+
+TEST(CUDAVirtualMemAllocatorV2, PlaceAtVARejectsOutOfRangeAddress) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kLarge);
+
+  allocator.InitOnce();
+  const VMMDevicePtr base = allocator.virtual_mem_base();
+  const size_t handle_size = allocator.handle_size();
+
+  EXPECT_THROW(allocator.PlaceAtVAWithBlock(base - handle_size, handle_size),
+               common::enforce::EnforceNotMet);
+  EXPECT_THROW(allocator.PlaceAtVAWithBlock(base + allocator.virtual_mem_size(),
+                                            handle_size),
+               common::enforce::EnforceNotMet);
+  EXPECT_THROW(
+      allocator.PlaceAtVAWithBlock(
+          base + allocator.virtual_mem_size() - handle_size, handle_size * 2),
+      common::enforce::EnforceNotMet);
+}
+
+TEST(CUDAVirtualMemAllocatorV2, FreeRemovesHandleRegistration) {
+  CUDAVirtualMemAllocatorV2 allocator(
+      phi::GPUPlace(), 2UL << 20, PoolType::kLarge);
+
+  auto allocation_with_block =
+      allocator.AppendWithBlock(allocator.handle_size());
+  ASSERT_NE(allocation_with_block.allocation, nullptr);
+  void* ptr = allocation_with_block.allocation->ptr();
+
+  allocation_with_block.allocation.reset();
+
+  auto reused = allocator.PlaceAtVAWithBlock(
+      reinterpret_cast<VMMDevicePtr>(ptr), allocator.handle_size());
+  ASSERT_NE(reused.allocation, nullptr);
+  EXPECT_EQ(reused.allocation->ptr(), ptr);
+  EXPECT_TRUE(allocator.IsRangeReleasable(reinterpret_cast<VMMDevicePtr>(ptr),
+                                          allocator.handle_size()));
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/get_base_ptr_test.cu

@@ -0,0 +1,119 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <random>
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/malloc.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class CUDAAllocatoionBasePtrTest : public ::testing::Test {
+ public:
+  void SetUp() override {
+    place_ = phi::GPUPlace();
+    alloc_times_ = 100;
+    batch_size_ = 10;
+    max_alloc_size_ = platform::GpuMaxAllocSize() / alloc_times_;
+    random_engine_ = std::default_random_engine(time(NULL));
+    dis_ = std::uniform_int_distribution<int>(0, max_alloc_size_);
+  }
+
+  void OneByOneAllocTest() {
+    for (size_t i = 0; i < alloc_times_; ++i) {
+      size_t size = dis_(random_engine_);
+      auto allocation = AllocShared(place_, size);
+
+      void* base_ptr = GetBasePtr(allocation);
+      void* system_ptr =
+          platform::GetGpuBasePtr(allocation->ptr(), place_.GetDeviceId());
+      EXPECT_EQ(base_ptr, system_ptr);
+    }
+
+    Release(place_);
+  }
+
+  void BatchByBatchAllocTest() {
+    std::vector<std::shared_ptr<phi::Allocation>> allocations;
+    allocations.reserve(batch_size_);
+    size_t batch_num = alloc_times_ / batch_size_;
+
+    for (size_t i = 0; i < batch_num; ++i) {
+      for (size_t j = 0; j < batch_size_; ++j) {
+        size_t size = dis_(random_engine_);
+        auto allocation = AllocShared(place_, size);
+
+        void* base_ptr = GetBasePtr(allocation);
+        void* system_ptr =
+            platform::GetGpuBasePtr(allocation->ptr(), place_.GetDeviceId());
+        EXPECT_EQ(base_ptr, system_ptr);
+
+        allocations.emplace_back(allocation);
+      }
+      allocations.clear();
+    }
+
+    Release(place_);
+  }
+
+  void ContinuousAllocTest() {
+    std::vector<std::shared_ptr<phi::Allocation>> allocations;
+    allocations.reserve(alloc_times_);
+
+    for (size_t i = 0; i < alloc_times_; ++i) {
+      size_t size = dis_(random_engine_);
+      auto allocation = AllocShared(place_, size);
+
+      void* base_ptr = GetBasePtr(allocation);
+      void* system_ptr =
+          platform::GetGpuBasePtr(allocation->ptr(), place_.GetDeviceId());
+      EXPECT_EQ(base_ptr, system_ptr);
+
+      allocations.emplace_back(allocation);
+    }
+
+    allocations.clear();
+    Release(place_);
+  }
+
+  void ZeroSizeAllocTest() {
+    auto allocation = AllocShared(place_, 0);
+    void* base_ptr = GetBasePtr(allocation);
+    void* system_ptr =
+        platform::GetGpuBasePtr(allocation->ptr(), place_.GetDeviceId());
+    EXPECT_EQ(base_ptr, system_ptr);
+  }
+
+ private:
+  phi::GPUPlace place_;
+  size_t max_alloc_size_;
+  size_t alloc_times_;
+  size_t batch_size_;
+  std::default_random_engine random_engine_;
+  std::uniform_int_distribution<int> dis_;
+};
+
+TEST_F(CUDAAllocatoionBasePtrTest, base_ptr_test) {
+  OneByOneAllocTest();
+  BatchByBatchAllocTest();
+  ContinuousAllocTest();
+  ZeroSizeAllocTest();
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/malloc_test.cu

@@ -0,0 +1,277 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <thread>  // NOLINT
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/allocation/allocator_facade.h"
+#include "paddle/phi/core/memory/malloc.h"
+#include "paddle/phi/core/platform/device_context.h"
+#include "paddle/phi/core/stream.h"
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#endif
+
+namespace paddle {
+namespace memory {
+
+const int NUM_STREAMS = 8;
+const int N = 2;
+const float DELTA = 1e-1;
+
+__global__ void kernel(float *x, int n) {
+  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = tid; i < n; i += blockDim.x * gridDim.x) {
+    x[i] = 3.14159 * i;
+  }
+}
+
+void CheckKernelOutput(const AllocationPtr &x, int n) {
+  auto host_x = std::unique_ptr<float[]>(new float[n]);
+  for (int i = 0; i < n; ++i) {
+#ifdef PADDLE_WITH_HIP
+    EXPECT_TRUE(hipSuccess == hipMemcpy(host_x.get(),
+                                        (x->ptr()),
+                                        n * sizeof(float),
+                                        hipMemcpyDeviceToHost));
+#else
+    EXPECT_TRUE(cudaSuccess == cudaMemcpy(host_x.get(),
+                                          (x->ptr()),
+                                          n * sizeof(float),
+                                          cudaMemcpyDeviceToHost));
+#endif
+    EXPECT_GE(host_x[i] + DELTA, 3.14159f * i);
+    EXPECT_LE(host_x[i] - DELTA, 3.14159f * i);
+  }
+}
+
+void MultiStreamCompute(const AllocationPtr &first_data,
+                        const AllocationPtr &second_data,
+                        phi::GPUContext *ctx) {
+  // multi-streams
+  EXPECT_GE(first_data->size(), N * sizeof(float));
+
+#ifdef PADDLE_WITH_HIP
+  hipLaunchKernelGGL((kernel),
+                     dim3(1),
+                     dim3(64),
+                     0,
+                     ctx->stream(),
+                     reinterpret_cast<float *>(first_data->ptr()),
+                     N);
+#else
+  kernel<<<1, 64, 0, ctx->stream()>>>(
+      reinterpret_cast<float *>(first_data->ptr()), N);
+#endif
+
+  EXPECT_GE(second_data->size(), N * sizeof(float));
+  // allocate and compute on same stream again
+
+#ifdef PADDLE_WITH_HIP
+  hipLaunchKernelGGL((kernel),
+                     dim3(1),
+                     dim3(64),
+                     0,
+                     ctx->stream(),
+                     reinterpret_cast<float *>(second_data->ptr()),
+                     N);
+#else
+  kernel<<<1, 64, 0, ctx->stream()>>>(
+      reinterpret_cast<float *>(second_data->ptr()), N);
+#endif
+}
+
+TEST(Malloc, GPUContextMultiStream) {
+  auto place = phi::GPUPlace(0);
+  platform::SetDeviceId(0);
+
+  AllocationPtr main_stream_alloc_ptr = Alloc(place, N * sizeof(float));
+  EXPECT_GE(main_stream_alloc_ptr->size(), N * sizeof(float));
+
+  AllocationPtr first_data[NUM_STREAMS], second_data[NUM_STREAMS];
+  std::vector<phi::GPUContext *> dev_ctx;
+
+// default stream
+#ifdef PADDLE_WITH_HIP
+  hipLaunchKernelGGL((kernel),
+                     dim3(1),
+                     dim3(64),
+                     0,
+                     0,
+                     reinterpret_cast<float *>(main_stream_alloc_ptr->ptr()),
+                     N);
+#else
+  kernel<<<1, 64>>>(reinterpret_cast<float *>(main_stream_alloc_ptr->ptr()), N);
+#endif
+  main_stream_alloc_ptr.reset();
+
+  for (int i = 0; i < NUM_STREAMS; ++i) {
+    auto ctx = new phi::GPUContext(place);
+    ctx->SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
+                          .GetAllocator(place, ctx->stream())
+                          .get());
+    ctx->SetHostAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetAllocator(phi::CPUPlace())
+            .get());
+    ctx->SetZeroAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetZeroAllocator(place)
+            .get());
+    ctx->SetPinnedAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetAllocator(phi::GPUPinnedPlace())
+            .get());
+    ctx->PartialInitWithAllocator();
+    dev_ctx.emplace_back(ctx);
+    first_data[i] =
+        Alloc(ctx->GetPlace(),
+              N * sizeof(float),
+              phi::Stream(reinterpret_cast<phi::StreamId>(ctx->stream())));
+    second_data[i] =
+        Alloc(ctx->GetPlace(),
+              N * sizeof(float),
+              phi::Stream(reinterpret_cast<phi::StreamId>(ctx->stream())));
+    MultiStreamCompute(first_data[i], second_data[i], ctx);
+  }
+
+#ifdef PADDLE_WITH_HIP
+  EXPECT_TRUE(hipSuccess == hipDeviceSynchronize());
+#else
+  EXPECT_TRUE(cudaSuccess == cudaDeviceSynchronize());
+#endif
+
+  for (int i = 0; i < NUM_STREAMS; ++i) {
+    CheckKernelOutput(first_data[i], N);
+    CheckKernelOutput(second_data[i], N);
+  }
+
+  // For cudaMallocAsyncAllocator, cudaFreeAsync is executed on _malloc_stream,
+  // which is the stream passed at Alloc(). Therefore, the stream must be
+  // postponed until the the memory is freed. Otherwise, the stream would be
+  // destroyed before the cudaFreeAsync is called.
+  for (int i = 0; i < NUM_STREAMS; i++) {
+    first_data[i].release();
+    second_data[i].release();
+    delete dev_ctx[i];
+  }
+}
+
+TEST(Malloc, GPUContextMultiThreadMultiStream) {
+  auto place = phi::GPUPlace(0);
+  platform::SetDeviceId(0);
+
+  AllocationPtr main_stream_alloc_ptr = Alloc(place, N * sizeof(float));
+  EXPECT_GE(main_stream_alloc_ptr->size(), N * sizeof(float));
+
+  AllocationPtr first_data[NUM_STREAMS], second_data[NUM_STREAMS];
+  std::vector<phi::GPUContext *> dev_ctx;
+
+// default stream
+#ifdef PADDLE_WITH_HIP
+  hipLaunchKernelGGL((kernel),
+                     dim3(1),
+                     dim3(64),
+                     0,
+                     0,
+                     reinterpret_cast<float *>(main_stream_alloc_ptr->ptr()),
+                     N);
+#else
+  kernel<<<1, 64>>>(reinterpret_cast<float *>(main_stream_alloc_ptr->ptr()), N);
+#endif
+  main_stream_alloc_ptr.reset();
+  std::vector<std::thread> threads;
+
+  for (int i = 0; i < NUM_STREAMS; ++i) {
+    auto ctx = new phi::GPUContext(place);
+    ctx->SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
+                          .GetAllocator(place, ctx->stream())
+                          .get());
+    ctx->SetHostAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetAllocator(phi::CPUPlace())
+            .get());
+    ctx->SetZeroAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetZeroAllocator(place)
+            .get());
+    ctx->SetHostZeroAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetZeroAllocator(phi::CPUPlace())
+            .get());
+    ctx->SetPinnedAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetAllocator(phi::GPUPinnedPlace())
+            .get());
+    ctx->PartialInitWithAllocator();
+    dev_ctx.emplace_back(ctx);
+    first_data[i] =
+        Alloc(ctx->GetPlace(),
+              N * sizeof(float),
+              phi::Stream(reinterpret_cast<phi::StreamId>(ctx->stream())));
+    second_data[i] =
+        Alloc(ctx->GetPlace(),
+              N * sizeof(float),
+              phi::Stream(reinterpret_cast<phi::StreamId>(ctx->stream())));
+    threads.emplace_back(MultiStreamCompute,
+                         std::ref(first_data[i]),
+                         std::ref(second_data[i]),
+                         ctx);
+  }
+
+  for (int i = 0; i < NUM_STREAMS; ++i) {
+    threads[i].join();
+  }
+
+#ifdef PADDLE_WITH_HIP
+  EXPECT_TRUE(hipSuccess == hipDeviceSynchronize());
+#else
+  EXPECT_TRUE(cudaSuccess == cudaDeviceSynchronize());
+#endif
+
+  for (int i = 0; i < NUM_STREAMS; ++i) {
+    CheckKernelOutput(first_data[i], N);
+    CheckKernelOutput(second_data[i], N);
+  }
+
+  // There are dependencies on the pointer deconstructing. Manually
+  // release the pointers would resolve the conflict.
+  for (int i = 0; i < NUM_STREAMS; i++) {
+    first_data[i].release();
+    second_data[i].release();
+    delete dev_ctx[i];
+  }
+}
+
+TEST(Malloc, AllocZero) {
+  auto place = phi::GPUPlace(0);
+  AllocationPtr allocation_ptr = Alloc(place, 0);
+  EXPECT_GE(allocation_ptr->size(), 0);
+}
+
+TEST(Malloc, AllocWithStream) {
+  size_t size = 1024;
+  AllocationPtr allocation = Alloc(phi::GPUPlace(), size, phi::Stream(0));
+  EXPECT_EQ(allocation->size(), 1024);
+}
+
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/memory_stats_test.cc

@@ -0,0 +1,65 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <algorithm>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/memory.h"
+
+namespace paddle {
+namespace memory {
+
+TEST(stat_allocator_test, host_memory_stat_test) {
+  std::vector<int64_t> alloc_sizes{
+      5278, 9593, 8492, 5041, 3351, 4232, 3706, 5963, 5896, 5057, 7527,
+      6235, 0,    7810, 940,  1239, 1945, 789,  2891, 7553, 8046, 2685,
+      1332, 6547, 5238, 5345, 1133, 5475, 9137, 3111, 8478, 6350, 9395,
+      4,    1185, 2186, 357,  9774, 6743, 6136, 7073, 7674, 5640, 3935,
+      528,  6699, 9821, 8717, 2264, 4708, 9936, 3566, 1373, 6955, 3694,
+      221,  309,  3617, 3793, 3334, 7281, 1302};
+
+  int64_t max_alloc_size = 0;
+  for (int64_t size : alloc_sizes) {
+    AllocationPtr allocation = Alloc(phi::CPUPlace(), size);
+    int64_t alloc_size = static_cast<int64_t>(allocation->size());
+    max_alloc_size = std::max(max_alloc_size, alloc_size);
+    EXPECT_EQ(HostMemoryStatCurrentValue("Allocated", 0), alloc_size);
+  }
+  EXPECT_EQ(HostMemoryStatPeakValue("Allocated", 0), max_alloc_size);
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+TEST(stat_allocator_test, device_memory_stat_test) {
+  std::vector<int64_t> alloc_sizes{
+      5278, 9593, 8492, 5041, 3351, 4232, 3706, 5963, 5896, 5057, 7527,
+      6235, 0,    7810, 940,  1239, 1945, 789,  2891, 7553, 8046, 2685,
+      1332, 6547, 5238, 5345, 1133, 5475, 9137, 3111, 8478, 6350, 9395,
+      4,    1185, 2186, 357,  9774, 6743, 6136, 7073, 7674, 5640, 3935,
+      528,  6699, 9821, 8717, 2264, 4708, 9936, 3566, 1373, 6955, 3694,
+      221,  309,  3617, 3793, 3334, 7281, 1302};
+
+  int64_t max_alloc_size = 0;
+  for (int64_t size : alloc_sizes) {
+    AllocationPtr allocation = Alloc(phi::GPUPlace(), size);
+    int64_t alloc_size = static_cast<int64_t>(allocation->size());
+    max_alloc_size = std::max(max_alloc_size, alloc_size);
+    EXPECT_EQ(DeviceMemoryStatCurrentValue("Allocated", 0), alloc_size);
+  }
+  EXPECT_EQ(DeviceMemoryStatPeakValue("Allocated", 0), max_alloc_size);
+}
+#endif
+
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/mmap_allocator_test.cc

@@ -0,0 +1,53 @@
+// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef _WIN32
+
+#include "paddle/phi/core/memory/allocation/mmap_allocator.h"
+
+#include "gtest/gtest.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(MemoryMapAllocation, test_allocation_base) {
+  size_t data_size = 4UL * 1024;
+
+  // 1. allocate writer holder
+  auto mmap_writer_holder = AllocateMemoryMapWriterAllocation(data_size);
+  std::string ipc_name = mmap_writer_holder->ipc_name();
+  // 2. write data
+  auto* writer_ptr = static_cast<int32_t*>(mmap_writer_holder->ptr());
+  for (int32_t i = 0; i < 1024; ++i) {
+    writer_ptr[i] = i;
+  }
+  // 3. create child process
+  pid_t fpid = fork();
+  if (fpid == 0) {
+    // 4. rebuild reader holder
+    auto mmap_reader_holder =
+        RebuildMemoryMapReaderAllocation(ipc_name, data_size);
+    auto* reader_ptr = static_cast<int32_t*>(mmap_reader_holder->ptr());
+    for (int32_t i = 0; i < 1024; ++i) {
+      ASSERT_EQ(reader_ptr[i], i);
+    }
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

test/cpp/fluid/memory/multi_scale_allocator_test.cc

@@ -0,0 +1,145 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.h"
+#include "paddle/phi/core/memory/allocation/retry_allocator.h"
+#include "paddle/phi/core/memory/allocation/virtual_memory_auto_growth_best_fit_allocator.h"
+#include "paddle/phi/core/memory/memory.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+#include "glog/logging.h"
+#include "gtest/gtest.h"
+PD_DECLARE_uint64(vmm_small_pool_pre_alloc_in_mb);
+PD_DECLARE_uint64(vmm_large_pool_pre_alloc_in_mb);
+PD_DECLARE_uint64(vmm_pre_alloc_in_mb);
+PD_DECLARE_uint64(vmm_small_pool_size_in_mb);
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Test fixture
+class VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocatorTest
+    : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    auto vmm_cuda_allocator_small =
+        std::make_shared<CUDAVirtualMemAllocator>(phi::GPUPlace(0));
+    auto vmm_cuda_allocator_large =
+        std::make_shared<CUDAVirtualMemAllocator>(phi::GPUPlace(0));
+    // Create mock underlying allocators
+    auto underlying_small =
+        std::make_shared<VirtualMemoryAutoGrowthBestFitAllocator>(
+            vmm_cuda_allocator_small, platform::GpuMinChunkSize(), GPUPlace(0));
+    auto underlying_large =
+        std::make_shared<VirtualMemoryAutoGrowthBestFitAllocator>(
+            vmm_cuda_allocator_large, platform::GpuMinChunkSize(), GPUPlace(0));
+
+    // Create the multi-scale pool allocator
+    multi_scale_allocator_ =
+        std::make_shared<VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator>(
+            underlying_small,
+            underlying_large,
+            platform::GpuMinChunkSize(),
+            GPUPlace(0));
+
+    small_allocator_ = underlying_small;
+    large_allocator_ = underlying_large;
+  }
+
+  size_t mb = (1 << 20);
+  std::shared_ptr<VirtualMemoryAutoGrowthBestFitAllocator> small_allocator_;
+  std::shared_ptr<VirtualMemoryAutoGrowthBestFitAllocator> large_allocator_;
+  std::shared_ptr<VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator>
+      multi_scale_allocator_;
+};
+
+// Test case for small pool pre-allocation only
+TEST_F(VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocatorTest,
+       PreAllocSmallPoolOnly) {
+  // Set flags for small pool pre-allocation
+  FLAGS_vmm_small_pool_pre_alloc_in_mb = 10;  // 10 MB
+  FLAGS_vmm_large_pool_pre_alloc_in_mb = 0;   // No large pool pre-allocation
+
+  multi_scale_allocator_->PreAlloc();
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 12 * mb);
+}
+
+// Test case for large pool pre-allocation only
+TEST_F(VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocatorTest,
+       PreAllocLargePoolOnly) {
+  // Set flags for large pool pre-allocation
+  FLAGS_vmm_small_pool_pre_alloc_in_mb = 0;   // No small pool pre-allocation
+  FLAGS_vmm_large_pool_pre_alloc_in_mb = 20;  // 20 MB
+
+  multi_scale_allocator_->PreAlloc();
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 22 * mb);
+}
+
+// Test case for both pools pre-allocation
+TEST_F(VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocatorTest,
+       PreAllocBothPools) {
+  // Set flags for both pools pre-allocation
+  FLAGS_v = 4;
+  FLAGS_vmm_small_pool_pre_alloc_in_mb = 5;   // 5 MB
+  FLAGS_vmm_large_pool_pre_alloc_in_mb = 15;  // 15 MB
+
+  multi_scale_allocator_->PreAlloc();
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 22 * mb);
+}
+
+// Test case for no pre-allocation
+TEST_F(VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocatorTest,
+       PreAllocNone) {
+  // Set flags for no pre-allocation
+  FLAGS_vmm_small_pool_pre_alloc_in_mb = 0;  // No pre-allocation
+  FLAGS_vmm_large_pool_pre_alloc_in_mb = 0;  // No pre-allocation
+
+  multi_scale_allocator_->PreAlloc();
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 0 * mb);
+}
+
+TEST_F(VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocatorTest,
+       PreAllocWithZeroSize) {
+  FLAGS_v = 4;
+  FLAGS_vmm_pre_alloc_in_mb = 0;
+  small_allocator_->PreAlloc();
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 0 * mb);
+}
+
+TEST_F(VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocatorTest,
+       PreAllocWithPositiveSize) {
+  FLAGS_vmm_pre_alloc_in_mb = 10;  // 10MB
+  large_allocator_->PreAlloc();
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 12 * mb);
+}
+
+TEST_F(VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocatorTest,
+       MultiScaleAlloc) {
+  FLAGS_vmm_small_pool_size_in_mb = 20;
+
+  auto allocation_small = multi_scale_allocator_->Allocate(10 * mb);
+  auto allocation_large = multi_scale_allocator_->Allocate(30 * mb);
+  auto safe = multi_scale_allocator_->IsAllocThreadSafe();
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 44 * mb);
+  EXPECT_EQ(safe, true);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/naive_best_fit_allocator_test.cc

@@ -0,0 +1,66 @@
+// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/naive_best_fit_allocator.h"
+
+#include "gtest/gtest.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(NaiveBestFitAllocatorTest, CpuAlloc) {
+  NaiveBestFitAllocator alloc{phi::CPUPlace()};
+  {
+    size_t size = (1 << 20);
+    auto allocation = alloc.Allocate(size);
+  }
+  alloc.Release(phi::CPUPlace());
+
+  size_t size = (1 << 20);
+  auto allocation = alloc.Allocate(size);
+  alloc.Release(phi::CPUPlace());
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+TEST(NaiveBestFitAllocatorTest, GpuAlloc) {
+  NaiveBestFitAllocator alloc{phi::GPUPlace(0)};
+  {
+    size_t size = (1 << 20);
+    auto allocation = alloc.Allocate(size);
+  }
+  alloc.Release(phi::GPUPlace(0));
+
+  size_t size = (1 << 20);
+  auto allocation = alloc.Allocate(size);
+  alloc.Release(phi::GPUPlace(0));
+}
+
+TEST(NaiveBestFitAllocatorTest, CudaPinnedAlloc) {
+  NaiveBestFitAllocator alloc{phi::GPUPinnedPlace()};
+  {
+    size_t size = (1 << 20);
+    auto allocation = alloc.Allocate(size);
+  }
+  alloc.Release(phi::GPUPinnedPlace());
+
+  size_t size = (1 << 20);
+  auto allocation = alloc.Allocate(size);
+  alloc.Release(phi::GPUPinnedPlace());
+}
+#endif
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/pinned_memory_test.cu

@@ -0,0 +1,202 @@
+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+#include <gtest/gtest.h>
+
+#include <unordered_map>
+
+#include "paddle/phi/backends/cpu/cpu_info.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/memory_block.h"
+#include "paddle/phi/core/memory/memcpy.h"
+#include "paddle/phi/core/memory/memory.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+// This unit test is an example comparing the performance between using pinned
+// memory and not. In general, using pinned memory will be faster.
+template <typename T>
+__global__ void Kernel(T* output, int dim) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (tid < dim) {
+    output[tid] = output[tid] * output[tid] / 100;
+  }
+}
+
+template <typename Place>
+float test_pinned_memory() {
+  Place cpu_place;
+  phi::GPUPlace cuda_place;
+
+  const int data_size = 4096;
+  const int iteration = 10;
+
+  // create event start and end
+  gpuEvent_t start_e, stop_e, copying_e;
+  float elapsedTime = 0;
+
+#ifdef PADDLE_WITH_HIP
+  hipEventCreate(&start_e);
+  hipEventCreate(&stop_e);
+  hipEventCreate(&copying_e);
+#else
+  cudaEventCreate(&start_e);
+  cudaEventCreate(&stop_e);
+  cudaEventCreate(&copying_e);
+#endif
+
+  // create computation stream, data copying stream
+  gpuStream_t computation_stream, copying_stream;
+#ifdef PADDLE_WITH_HIP
+  hipStreamCreate(&computation_stream);
+  hipStreamCreate(&copying_stream);
+#else
+  cudaStreamCreate(&computation_stream);
+  cudaStreamCreate(&copying_stream);
+#endif
+
+  // create record event, pinned memory, gpu memory
+  std::vector<gpuEvent_t> record_event(iteration);
+  std::vector<float*> input_pinned_mem(iteration);
+  std::vector<float*> gpu_mem(iteration);
+  std::vector<float*> output_pinned_mem(iteration);
+
+  // initial data
+  for (int j = 0; j < iteration; ++j) {
+#ifdef PADDLE_WITH_HIP
+    hipEventCreateWithFlags(&record_event[j], hipEventDisableTiming);
+    hipEventCreate(&(record_event[j]));
+#else
+    cudaEventCreateWithFlags(&record_event[j], cudaEventDisableTiming);
+    cudaEventCreate(&(record_event[j]));
+#endif
+    input_pinned_mem[j] = static_cast<float*>(
+        paddle::memory::Alloc(cpu_place, data_size * sizeof(float)));
+    output_pinned_mem[j] = static_cast<float*>(
+        paddle::memory::Alloc(cpu_place, data_size * sizeof(float)));
+    gpu_mem[j] = static_cast<float*>(
+        paddle::memory::Alloc(cuda_place, data_size * sizeof(float)));
+
+    for (int k = 0; k < data_size; ++k) {
+      input_pinned_mem[j][k] = k;
+    }
+  }
+
+#ifdef PADDLE_WITH_HIP
+  hipEventRecord(start_e, computation_stream);
+#else
+  cudaEventRecord(start_e, computation_stream);
+#endif
+
+  // computation
+  for (int m = 0; m < 30; ++m) {
+    for (int i = 0; i < iteration; ++i) {
+      // cpu -> GPU on computation stream.
+      // note: this operation is async for pinned memory.
+      paddle::memory::Copy(cuda_place,
+                           gpu_mem[i],
+                           cpu_place,
+                           input_pinned_mem[i],
+                           data_size * sizeof(float),
+                           computation_stream);
+
+      // call kernel on computation stream.
+      Kernel<<<4, 1024, 0, computation_stream>>>(gpu_mem[i], data_size);
+
+#ifdef PADDLE_WITH_HIP
+      // record event_computation on computation stream
+      hipEventRecord(record_event[i], computation_stream);
+
+      // wait event_computation on copy stream.
+      // note: this operation is async.
+      hipStreamWaitEvent(copying_stream, record_event[i], 0);
+#else
+      // record event_computation on computation stream
+      cudaEventRecord(record_event[i], computation_stream);
+
+      // wait event_computation on copy stream.
+      // note: this operation is async.
+      cudaStreamWaitEvent(copying_stream, record_event[i], 0);
+#endif
+      // copy data GPU->CPU, on copy stream.
+      // note: this operation is async for pinned memory.
+      paddle::memory::Copy(cpu_place,
+                           output_pinned_mem[i],
+                           cuda_place,
+                           gpu_mem[i],
+                           data_size * sizeof(float),
+                           copying_stream);
+    }
+  }
+
+#ifdef PADDLE_WITH_HIP
+  hipEventRecord(copying_e, copying_stream);
+  hipStreamWaitEvent(computation_stream, copying_e, 0);
+
+  hipEventRecord(stop_e, computation_stream);
+
+  hipEventSynchronize(start_e);
+  hipEventSynchronize(stop_e);
+  hipEventElapsedTime(&elapsedTime, start_e, stop_e);
+#else
+  cudaEventRecord(copying_e, copying_stream);
+  cudaStreamWaitEvent(computation_stream, copying_e, 0);
+
+  cudaEventRecord(stop_e, computation_stream);
+
+  cudaEventSynchronize(start_e);
+  cudaEventSynchronize(stop_e);
+  cudaEventElapsedTime(&elapsedTime, start_e, stop_e);
+#endif
+
+  // std::cout << cpu_place << " "
+  //          << "time consume:" << elapsedTime / 30 << std::endl;
+
+  for (int l = 0; l < iteration; ++l) {
+    for (int k = 0; k < data_size; ++k) {
+      float temp = input_pinned_mem[l][k];
+      temp = temp * temp / 100;
+      EXPECT_FLOAT_EQ(temp, output_pinned_mem[l][k]);
+    }
+  }
+
+// destroy resource
+#ifdef PADDLE_WITH_HIP
+  hipEventDestroy(copying_e);
+  hipEventDestroy(start_e);
+  hipEventDestroy(stop_e);
+#else
+  cudaEventDestroy(copying_e);
+  cudaEventDestroy(start_e);
+  cudaEventDestroy(stop_e);
+#endif
+  for (int j = 0; j < 10; ++j) {
+#ifdef PADDLE_WITH_HIP
+    hipEventDestroy((record_event[j]));
+#else
+    cudaEventDestroy((record_event[j]));
+#endif
+    paddle::memory::Free(cpu_place, input_pinned_mem[j]);
+    paddle::memory::Free(cpu_place, output_pinned_mem[j]);
+    paddle::memory::Free(cuda_place, gpu_mem[j]);
+  }
+  return elapsedTime / 30;
+}
+
+TEST(CPUANDCUDAPinned, CPUAllocatorAndCUDAPinnedAllocator) {
+  // Generally speaking, operation on pinned_memory is faster than that on
+  // unpinned-memory, but if this unit test fails frequently, please close this
+  // test for the time being.
+  float time1 = test_pinned_memory<phi::CPUPlace>();
+  float time2 = test_pinned_memory<phi::GPUPinnedPlace>();
+  EXPECT_GT(time1, time2);
+}

test/cpp/fluid/memory/retry_allocator_test.cc

@@ -0,0 +1,139 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/retry_allocator.h"
+
+#include <thread>  // NOLINT
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/allocation/best_fit_allocator.h"
+#include "paddle/phi/core/memory/allocation/cpu_allocator.h"
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include "paddle/phi/core/memory/allocation/cuda_allocator.h"
+#endif
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(RetryAllocator, RetryAllocator) {
+  CPUAllocator cpu_allocator;
+
+  size_t size = (1 << 20);
+  auto cpu_allocation = cpu_allocator.Allocate(size);
+
+  size_t thread_num = 4;
+  size_t sleep_time = 40;
+  size_t extra_time = 20;
+
+  // Reserve to perform more tests in the future
+  std::vector<std::shared_ptr<Allocator>> allocators;
+  {
+    std::unique_ptr<BestFitAllocator> best_fit_allocator(
+        new BestFitAllocator(cpu_allocation.get()));
+    allocators.push_back(std::make_shared<RetryAllocator>(
+        std::move(best_fit_allocator),
+        phi::CPUPlace(),
+        (thread_num - 1) * (sleep_time + extra_time)));
+  }
+
+  for (auto &allocator : allocators) {
+    std::vector<std::thread> threads(thread_num);
+    std::vector<void *> addresses(threads.size(), nullptr);
+
+    std::mutex mutex;
+    std::condition_variable cv;
+    bool flag = false;
+
+    for (size_t i = 0; i < threads.size(); ++i) {
+      threads[i] = std::thread([&, i]() {
+        {
+          std::unique_lock<std::mutex> lock(mutex);
+          cv.wait(lock, [&] { return flag; });
+        }
+
+        auto ret = allocator->Allocate(size - 1);
+        addresses[i] = ret->ptr();
+        std::this_thread::sleep_for(std::chrono::milliseconds(sleep_time));
+      });
+    }
+
+    {
+      std::lock_guard<std::mutex> lock(mutex);
+      flag = true;
+      cv.notify_all();
+    }
+
+    for (auto &th : threads) {
+      th.join();
+    }
+
+    void *val = cpu_allocation->ptr();
+    bool is_all_equal = std::all_of(addresses.begin(),
+                                    addresses.end(),
+                                    [val](void *p) { return p == val; });
+    ASSERT_TRUE(is_all_equal);
+    allocator->Release(phi::CPUPlace());
+  }
+}
+
+class DummyAllocator : public Allocator {
+ public:
+  bool IsAllocThreadSafe() const override { return true; }
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override {
+    PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+        "Here is a test exception, always BadAlloc."));
+  }
+
+  void FreeImpl(phi::Allocation *) override {}
+};
+
+TEST(RetryAllocator, RetryAllocatorLastAllocFailure) {
+  size_t retry_ms = 10;
+  {
+    RetryAllocator allocator(
+        std::make_shared<DummyAllocator>(), phi::CPUPlace(), retry_ms);
+    try {
+      auto allocation = allocator.Allocate(100);
+      ASSERT_TRUE(false);
+      allocation.reset();
+    } catch (BadAlloc &ex) {
+      ASSERT_TRUE(std::string(ex.what()).find("always BadAlloc") !=
+                  std::string::npos);
+    }
+  }
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  {
+    phi::GPUPlace p(0);
+    RetryAllocator allocator(std::make_shared<CUDAAllocator>(p), p, retry_ms);
+    size_t allocate_size = (static_cast<size_t>(1) << 40);  // Very large number
+    try {
+      auto allocation = allocator.Allocate(allocate_size);
+      ASSERT_TRUE(false);
+      allocation.reset();
+      allocator.Release(p);
+    } catch (BadAlloc &ex) {
+      ASSERT_TRUE(std::string(ex.what()).find("Cannot allocate") !=
+                  std::string::npos);
+    }
+  }
+#endif
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/stats_test.cc

@@ -0,0 +1,201 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/stats.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+namespace paddle {
+namespace memory {
+
+class StatsTest : public ::testing::Test {
+ protected:
+  void SetStatType(const std::string& stat_type) { stat_type_ = stat_type; }
+
+  void SetFunc(
+      std::function<void(const std::string, int, int64_t)> update_func,
+      std::function<int64_t(const std::string, int)> current_value_func,
+      std::function<int64_t(const std::string, int)> peak_value_func,
+      std::function<void(const std::string, int)> reset_peak_value_func) {
+    update_func_ = update_func;
+    current_value_func_ = current_value_func;
+    peak_value_func_ = peak_value_func;
+    reset_peak_value_func_ = reset_peak_value_func;
+  }
+
+  void RunTests() {
+    MultiThreadReadWriteTest();
+    PeakValueTest();
+    ResetPeakValueTest();
+  }
+
+ private:
+  void MultiThreadReadWriteTest() {
+    size_t thread_num = 3;
+    size_t data_num = 10;
+
+    std::condition_variable cv;
+    std::mutex mutex;
+    std::vector<std::thread> threads;
+    size_t ready_thread_num = 0;
+
+    for (size_t i = 0; i < thread_num; ++i) {
+      threads.emplace_back([&]() {
+        for (size_t data = 0; data < data_num; ++data) {
+          update_func_(stat_type_, 0, static_cast<int64_t>(data));
+        }
+        /* lock guard*/ {
+          std::lock_guard<std::mutex> lock_guard{mutex};
+          ++ready_thread_num;
+          cv.notify_one();
+        }
+        // Sleep here to not exit before the main thread checking stat
+        // results, because the thread-local stat data will be destroyed when
+        // the thread exit
+        std::this_thread::sleep_for(std::chrono::seconds(1));
+      });
+    }
+
+    std::unique_lock<std::mutex> unique_lock(mutex);
+    cv.wait(unique_lock, [&ready_thread_num, thread_num]() {
+      return ready_thread_num == thread_num;
+    });
+
+    EXPECT_EQ(current_value_func_(stat_type_, 0),
+              int64_t((thread_num * data_num * (data_num - 1)) >> 1));
+
+    for (size_t i = 0; i < thread_num; ++i) {
+      threads[i].join();
+    }
+  }
+
+  void PeakValueTest() {
+    int64_t peak_value = ((int64_t)1) << 63;
+    int64_t sum = 0;
+    for (int64_t data : datas_) {
+      update_func_(stat_type_, 0, data);
+      sum += data;
+      peak_value = std::max(peak_value, sum);
+    }
+    EXPECT_EQ(peak_value_func_(stat_type_, 0), peak_value);
+  }
+
+  void ResetPeakValueTest() {
+    for (int64_t data : datas_) {
+      update_func_(stat_type_, 0, data);
+
+      EXPECT_GE(peak_value_func_(stat_type_, 0),
+                current_value_func_(stat_type_, 0));
+      reset_peak_value_func_(stat_type_, 0);
+      EXPECT_EQ(peak_value_func_(stat_type_, 0),
+                current_value_func_(stat_type_, 0));
+    }
+  }
+
+  std::string stat_type_;
+  std::vector<int64_t> datas_{
+      543149808935355, 634698327471328, 706215795436611, 577939367795333,
+      419479490054362, 21975227714595,  812939817942250, 984428837942082,
+      537304104446806, 685008544452453, 563352858161268, 690143831596330,
+      964829938186077, 476984078018245, 804403365180177, -57918691189304,
+      947611269236893, 752188963801927, 710946451346683, -49226452527666,
+      -59049377393968, 14128239868858,  463298869064035, 71954818131880,
+      894087341752481, 971337322257029, 202325222441382, 128423535063606,
+      -89146949094815, 756429151957759, 444400180007032, 937040878834954,
+      303916192293233, 16628488962638,  544031750807065, 392396591234910,
+      686663859558365, 941126625484539, 120719755546781, 938838399629825,
+      364952832531949, 237865770815218, -64409925441421, 130095171433100,
+      140906146174023, 635514857321759, -65954585142544, 505369882354612,
+      939334896592688, 591590196329715, 424834428510773, 316569328289240,
+      44932622352645,  464924685290752, 396541464249293, 937169087747437,
+      437992536948503, 44395833829426,  968496835801562, 80493658180301,
+      836093264717766, 3339912102452,   -32067753603273, 77353521424986,
+      290980283590981, 496135147814915, 335112580987207, 571094882208895,
+      776581672377954, -83075504255716, -93690563747742, 115144063088100,
+      422629490055299, 917988755593299, 283511671626409, 715179006446336,
+      760708617230450, 183628659314298, 899792829140365, 214949068928854,
+      848851506468080, 791041814082114, 801591030978388, 526551272394511,
+      781034506085043, 279998089943681, 907197980150568, 974365521595836,
+      282127262539024, 272870474932399, 346617645597508, 409964014011113,
+      746465732805300, -74049761897414, -65640372433924, 852009039806484,
+      305079802044257, -48409757869238, 266031781660228, 327287322379820};
+
+  std::function<void(const std::string, int, int64_t)> update_func_;
+  std::function<int64_t(const std::string, int)> current_value_func_;
+  std::function<int64_t(const std::string, int)> peak_value_func_;
+  std::function<void(const std::string, int)> reset_peak_value_func_;
+};
+
+TEST_F(StatsTest, DeviceAllocatedTest) {
+  SetStatType("Allocated");
+  SetFunc(DeviceMemoryStatUpdate,
+          DeviceMemoryStatCurrentValue,
+          DeviceMemoryStatPeakValue,
+          DeviceMemoryStatResetPeakValue);
+  RunTests();
+}
+
+TEST_F(StatsTest, DeviceReservedMacroTest) {
+  SetStatType("Reserved");
+  SetFunc(
+      [](const std::string stat_type, int id, int64_t increment) {
+        return DEVICE_MEMORY_STAT_UPDATE(Reserved, id, increment);
+      },
+      [](const std::string stat_type, int id) {
+        return DEVICE_MEMORY_STAT_CURRENT_VALUE(Reserved, id);
+      },
+      [](const std::string stat_type, int id) {
+        return DEVICE_MEMORY_STAT_PEAK_VALUE(Reserved, id);
+      },
+      [](const std::string stat_type, int id) {
+        return DEVICE_MEMORY_STAT_RESET_PEAK_VALUE(Reserved, id);
+      });
+  RunTests();
+}
+
+TEST_F(StatsTest, HostAllocatedMacroTest) {
+  SetStatType("Allocated");
+  SetFunc(
+      [](const std::string stat_type, int id, int64_t increment) {
+        return HOST_MEMORY_STAT_UPDATE(Allocated, id, increment);
+      },
+      [](const std::string stat_type, int id) {
+        return HOST_MEMORY_STAT_CURRENT_VALUE(Allocated, id);
+      },
+      [](const std::string stat_type, int id) {
+        return HOST_MEMORY_STAT_PEAK_VALUE(Allocated, id);
+      },
+      [](const std::string stat_type, int id) {
+        return HOST_MEMORY_STAT_RESET_PEAK_VALUE(Allocated, id);
+      });
+  RunTests();
+}
+
+TEST_F(StatsTest, HostReservedTest) {
+  SetStatType("Reserved");
+  SetFunc(HostMemoryStatUpdate,
+          HostMemoryStatCurrentValue,
+          HostMemoryStatPeakValue,
+          HostMemoryStatResetPeakValue);
+  RunTests();
+}
+
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/stream_safe_cuda_alloc_test.cu

@@ -0,0 +1,449 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <thread>  // NOLINT
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/allocation/allocator_facade.h"
+#include "paddle/phi/core/memory/memory.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#include "paddle/phi/core/platform/device_context.h"
+#include "paddle/phi/core/stream.h"
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include "paddle/phi/core/platform/cuda_graph_with_memory_pool.h"
+#endif
+
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#endif
+
+#define RETURN_IF_NOT_ENABLED                            \
+  {                                                      \
+    if (!memory::allocation::AllocatorFacade::Instance() \
+             .IsStreamSafeCUDAAllocatorUsed()) {         \
+      return;                                            \
+    }                                                    \
+  }
+
+namespace paddle {
+namespace memory {
+
+// y += (x + 1)
+__global__ void add_kernel(int *x, int *y, int n) {
+  int thread_num = gridDim.x * blockDim.x;
+  int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int i = thread_id; i < n; i += thread_num) {
+    y[i] += x[i] + 1;
+  }
+}
+
+void CheckMemLeak(const phi::GPUPlace &place) {
+  uint64_t cuda_malloc_size =
+      platform::RecordedGpuMallocSize(place.GetDeviceId());
+  ASSERT_EQ(cuda_malloc_size, 0)
+      << "Found " << cuda_malloc_size << " bytes memory that not released yet,"
+      << " there may be a memory leak problem";
+}
+
+TEST(StreamSafeCUDAAllocInterfaceTest, AllocInterfaceTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  phi::GPUPlace place = phi::GPUPlace();
+  size_t alloc_size = 256;
+
+  std::shared_ptr<Allocation> allocation_implicit_stream =
+      AllocShared(place, alloc_size);
+  EXPECT_GE(allocation_implicit_stream->size(), alloc_size);
+
+  void *address = allocation_implicit_stream->ptr();
+  allocation_implicit_stream.reset();
+
+  gpuStream_t default_stream =
+      dynamic_cast<phi::GPUContext *>(
+          phi::DeviceContextPool::Instance().Get(place))
+          ->stream();
+  allocation::AllocationPtr allocation_unique =
+      Alloc(place,
+            alloc_size,
+            phi::Stream(reinterpret_cast<phi::StreamId>(default_stream)));
+  EXPECT_GE(allocation_unique->size(), alloc_size);
+  EXPECT_EQ(allocation_unique->ptr(), address);
+  allocation_unique.reset();
+
+  Release(place);
+  CheckMemLeak(place);
+}
+
+TEST(StreamSafeCUDAAllocInterfaceTest, GetAllocatorInterfaceTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  phi::GPUPlace place = phi::GPUPlace();
+  size_t alloc_size = 256;
+
+  allocation::AllocationPtr allocation_implicit_stream =
+      Alloc(place, alloc_size);
+  EXPECT_GE(allocation_implicit_stream->size(), alloc_size);
+  void *address = allocation_implicit_stream->ptr();
+  allocation_implicit_stream.reset();
+
+  auto &instance = allocation::AllocatorFacade::Instance();
+  const std::shared_ptr<Allocator> &allocator = instance.GetAllocator(place);
+
+  allocation::AllocationPtr allocation_from_allocator =
+      allocator->Allocate(alloc_size);
+  EXPECT_GE(allocation_from_allocator->size(), alloc_size);
+  EXPECT_EQ(allocation_from_allocator->ptr(), address);
+  allocation_from_allocator.reset();
+
+  Release(place);
+  CheckMemLeak(place);
+}
+
+TEST(StreamSafeCUDAAllocInterfaceTest, GetAllocatorWithDefaultStreamTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  auto &instance = allocation::AllocatorFacade::Instance();
+  phi::GPUPlace place = phi::GPUPlace();
+  const std::shared_ptr<Allocator> allocator_implicit_stream =
+      instance.GetAllocator(place);
+  const std::shared_ptr<Allocator> allocator_default_stream =
+      instance.GetAllocator(place,
+                            static_cast<phi::GPUContext *>(
+                                phi::DeviceContextPool::Instance().Get(place))
+                                ->stream());
+  EXPECT_EQ(allocator_implicit_stream.get(), allocator_default_stream.get());
+}
+
+TEST(StreamSafeCUDAAllocInterfaceTest, ZeroSizeRecordStreamTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  phi::GPUPlace place = phi::GPUPlace();
+  std::shared_ptr<Allocation> zero_size_allocation = AllocShared(place, 0);
+  EXPECT_EQ(zero_size_allocation->ptr(), nullptr);
+
+  gpuStream_t stream;
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamCreate(&stream));
+#endif
+
+  EXPECT_NO_THROW(RecordStream(zero_size_allocation, stream));
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(stream));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamDestroy(stream));
+#endif
+}
+
+TEST(StreamSafeCUDAAllocInterfaceTest, GetStreamInterfaceTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  phi::GPUPlace place = phi::GPUPlace();
+  size_t alloc_size = 256;
+
+  gpuStream_t default_stream =
+      dynamic_cast<phi::GPUContext *>(
+          phi::DeviceContextPool::Instance().Get(place))
+          ->stream();
+  std::shared_ptr<Allocation> allocation_implicit_stream =
+      AllocShared(place, alloc_size);
+  EXPECT_EQ(GetStream(allocation_implicit_stream), default_stream);
+
+  gpuStream_t new_stream;
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&new_stream));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamCreate(&new_stream));
+#endif
+
+  std::shared_ptr<Allocation> allocation_new_stream =
+      AllocShared(place,
+                  alloc_size,
+                  phi::Stream(reinterpret_cast<phi::StreamId>(new_stream)));
+  EXPECT_EQ(GetStream(allocation_new_stream), new_stream);
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(new_stream));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamDestroy(new_stream));
+#endif
+
+  allocation_implicit_stream.reset();
+  allocation_new_stream.reset();
+  Release(place);
+  CheckMemLeak(place);
+}
+
+TEST(StreamSafeCUDAAllocRetryTest, RetryTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  phi::GPUPlace place = phi::GPUPlace();
+  gpuStream_t stream1, stream2;
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream1));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream2));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamCreate(&stream1));
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamCreate(&stream2));
+#endif
+  size_t available_size = platform::GpuAvailableMemToAlloc();
+  // alloc_size < available_size < 2 * alloc_size,
+  // so the second alloc will fail and retry
+  size_t alloc_size = available_size / 4 * 3;
+
+  allocation::AllocationPtr allocation1 = Alloc(
+      place, alloc_size, phi::Stream(reinterpret_cast<phi::StreamId>(stream1)));
+  allocation::AllocationPtr allocation2;
+
+  std::thread th([&allocation2, &place, &stream2, alloc_size]() {
+    std::this_thread::sleep_for(std::chrono::seconds(1));
+    allocation2 = Alloc(place,
+                        alloc_size,
+                        phi::Stream(reinterpret_cast<phi::StreamId>(stream2)));
+  });
+  allocation1.reset();  // free but not release
+  th.join();
+  EXPECT_GE(allocation2->size(), alloc_size);
+  allocation2.reset();
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceSynchronize());
+#endif
+
+  Release(place, stream1);
+  Release(place, stream2);
+  CheckMemLeak(place);
+}
+
+class StreamSafeCUDAAllocTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    place_ = phi::GPUPlace();
+    stream_num_ = 64;
+    grid_num_ = 1;
+    block_num_ = 32;
+    data_num_ = 131072;
+    workspace_size_ = data_num_ * sizeof(int);
+
+    for (size_t i = 0; i < stream_num_; ++i) {
+      gpuStream_t stream;
+#ifdef PADDLE_WITH_CUDA
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream));
+#else
+      PADDLE_ENFORCE_GPU_SUCCESS(hipStreamCreate(&stream));
+#endif
+
+      std::shared_ptr<phi::Allocation> workspace_allocation =
+          AllocShared(place_,
+                      workspace_size_,
+                      phi::Stream(reinterpret_cast<phi::StreamId>(stream)));
+      std::shared_ptr<phi::Allocation> result_allocation =
+          AllocShared(place_,
+                      workspace_size_,
+                      phi::Stream(reinterpret_cast<phi::StreamId>(stream)));
+      std::shared_ptr<phi::Allocation> host_result_allocation =
+          AllocShared(phi::CPUPlace(), workspace_size_);
+
+#ifdef PADDLE_WITH_CUDA
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaMemset(
+          workspace_allocation->ptr(), 0, workspace_allocation->size()));
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaMemset(result_allocation->ptr(), 0, result_allocation->size()));
+#else
+      PADDLE_ENFORCE_GPU_SUCCESS(hipMemset(
+          workspace_allocation->ptr(), 0, workspace_allocation->size()));
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          hipMemset(result_allocation->ptr(), 0, result_allocation->size()));
+#endif
+
+      streams_.emplace_back(stream);
+      workspaces_.emplace_back(workspace_allocation);
+      results_.emplace_back(result_allocation);
+      host_results_.emplace_back(host_result_allocation);
+    }
+  }
+
+  void SingleStreamRun(size_t idx) {
+    int *y = reinterpret_cast<int *>(results_[idx]->ptr());
+    int neighbouring_idx = idx > 0 ? idx - 1 : idx;
+
+    add_kernel<<<grid_num_, block_num_, 0, streams_[idx]>>>(
+        reinterpret_cast<int *>(workspaces_[idx]->ptr()), y, data_num_);
+    add_kernel<<<grid_num_, block_num_, 0, streams_[idx]>>>(
+        reinterpret_cast<int *>(workspaces_[neighbouring_idx]->ptr()),
+        y,
+        data_num_);
+    RecordStream(workspaces_[neighbouring_idx], streams_[idx]);
+  }
+
+  void MultiStreamRun() {
+    // Must run in reverse order, or the workspace_[i - 1] will be released
+    // before streams_[i]'s kernel launch
+    for (int i = stream_num_ - 1; i >= 0; --i) {
+      SingleStreamRun(i);
+      workspaces_[i].reset();  // fast GC
+    }
+  }
+
+  void MultiThreadMultiStreamRun() {
+    std::vector<std::thread> threads;
+    for (size_t i = 0; i < stream_num_; ++i) {
+      threads.emplace_back(&StreamSafeCUDAAllocTest::SingleStreamRun, this, i);
+    }
+    for (size_t i = 0; i < stream_num_; ++i) {
+      threads[i].join();
+    }
+    workspaces_.clear();
+  }
+
+  void CUDAGraphRun() {
+    testing_cuda_graph_ = true;
+    platform::BeginCUDAGraphCapture(phi::GPUPlace(),
+                                    cudaStreamCaptureModeGlobal);
+
+    std::shared_ptr<Allocation> data_allocation =
+        AllocShared(phi::GPUPlace(), workspace_size_);
+    std::shared_ptr<Allocation> result_allocation =
+        AllocShared(phi::GPUPlace(), workspace_size_);
+
+    int *data = static_cast<int *>(data_allocation->ptr());
+    int *result = static_cast<int *>(result_allocation->ptr());
+
+    gpuStream_t main_stream = GetStream(data_allocation);
+    gpuStream_t other_stream;
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&other_stream));
+
+    add_kernel<<<grid_num_, block_num_, 0, main_stream>>>(
+        data, result, data_num_);
+    RecordStream(data_allocation, other_stream);
+
+    std::unique_ptr<phi::backends::gpu::CUDAGraph> cuda_graph =
+        platform::EndCUDAGraphCapture();
+
+    int replay_times = 10;
+    for (int i = 0; i < replay_times; ++i) {
+      cuda_graph->Replay();
+    }
+
+    std::shared_ptr<Allocation> host_result_allocation =
+        AllocShared(phi::CPUPlace(), workspace_size_);
+    Copy(host_result_allocation->place(),
+         host_result_allocation->ptr(),
+         result_allocation->place(),
+         result_allocation->ptr(),
+         workspace_size_,
+         main_stream);
+    cudaStreamSynchronize(main_stream);
+
+    int *host_result = static_cast<int *>(host_result_allocation->ptr());
+    for (int i = 0; i < data_num_; ++i) {
+      EXPECT_EQ(host_result[i], replay_times);
+    }
+
+    data_allocation.reset();
+    result_allocation.reset();
+    cuda_graph.release();
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(other_stream));
+  }
+
+  void CheckResult() {
+    for (size_t i = 0; i < stream_num_; ++i) {
+      Copy(host_results_[i]->place(),
+           host_results_[i]->ptr(),
+           results_[i]->place(),
+           results_[i]->ptr(),
+           workspace_size_,
+           streams_[i]);
+    }
+    cudaDeviceSynchronize();
+
+    size_t thread_num = grid_num_ * block_num_;
+    for (size_t i = 0; i < stream_num_; ++i) {
+      int *result = static_cast<int *>(host_results_[i]->ptr());
+      for (size_t j = 0; j < data_num_; ++j) {
+        EXPECT_EQ(result[j], 2);
+      }
+    }
+  }
+
+  void TearDown() override {
+    workspaces_.clear();
+    results_.clear();
+    host_results_.clear();
+    for (gpuStream_t stream : streams_) {
+      Release(place_, stream);
+    }
+
+    for (size_t i = 0; i < stream_num_; ++i) {
+#ifdef PADDLE_WITH_CUDA
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(streams_[i]));
+#else
+      PADDLE_ENFORCE_GPU_SUCCESS(hipStreamDestroy(streams_[i]));
+#endif
+    }
+
+    // Memory release for CUDA Graph memory pool is forbidden
+    if (!testing_cuda_graph_) {
+      CheckMemLeak(place_);
+    }
+  }
+
+  bool testing_cuda_graph_{0};
+  size_t stream_num_;
+  size_t grid_num_;
+  size_t block_num_;
+  size_t data_num_;
+  size_t workspace_size_;
+  phi::GPUPlace place_;
+  std::vector<gpuStream_t> streams_;
+  std::vector<std::shared_ptr<phi::Allocation>> workspaces_;
+  std::vector<std::shared_ptr<phi::Allocation>> results_;
+  std::vector<std::shared_ptr<phi::Allocation>> host_results_;
+};
+
+TEST_F(StreamSafeCUDAAllocTest, CUDAMutilStreamTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  MultiStreamRun();
+  CheckResult();
+}
+
+TEST_F(StreamSafeCUDAAllocTest, CUDAMutilThreadMutilStreamTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  MultiThreadMultiStreamRun();
+  CheckResult();
+}
+
+#if (defined(PADDLE_WITH_CUDA) && (CUDA_VERSION >= 11000))
+TEST_F(StreamSafeCUDAAllocTest, CUDAGraphTest) {
+  RETURN_IF_NOT_ENABLED;
+
+  MultiStreamRun();
+  CUDAGraphRun();
+  CheckResult();
+}
+#endif
+
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/system_allocator_test.cc

@@ -0,0 +1,84 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/core/memory/allocation/system_allocator.h"
+
+#include <gtest/gtest.h>
+#include <memory>
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/platform/device/device_wrapper.h"
+
+COMMON_DECLARE_bool(use_pinned_memory);
+
+void TestAllocator(paddle::memory::detail::SystemAllocator* a, size_t size) {
+  bool freed = false;
+  {
+    size_t index;  // NOLINT
+    void* p = a->Alloc(&index, size);
+    if (size > 0) {
+      EXPECT_NE(p, nullptr);
+    } else {
+      EXPECT_EQ(p, nullptr);
+    }
+
+    int* i = static_cast<int*>(p);
+    std::shared_ptr<int> ptr(i, [&](void* p) {
+      freed = true;
+      a->Free(p, size, index);
+    });
+  }
+  EXPECT_TRUE(freed);
+}
+
+TEST(CPUAllocator, NoLockMem) {
+  FLAGS_use_pinned_memory = false;
+  paddle::memory::detail::CPUAllocator a;
+  TestAllocator(&a, 2048);
+  TestAllocator(&a, 0);
+}
+
+TEST(CPUAllocator, LockMem) {
+  FLAGS_use_pinned_memory = true;
+  paddle::memory::detail::CPUAllocator a;
+  TestAllocator(&a, 2048);
+  TestAllocator(&a, 0);
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+TEST(GPUAllocator, Alloc) {
+  paddle::memory::detail::GPUAllocator a(0);
+  TestAllocator(&a, 2048);
+  TestAllocator(&a, 0);
+}
+
+TEST(CUDAPinnedAllocator, Alloc) {
+  paddle::memory::detail::CUDAPinnedAllocator a;
+  TestAllocator(&a, 2048);
+  TestAllocator(&a, 0);
+}
+
+TEST(GPUAllocator, AllocFailure) {
+  paddle::memory::detail::GPUAllocator allocator(0);
+  size_t index;
+  size_t alloc_size = (static_cast<size_t>(1) << 40);  // Very large number
+  try {
+    allocator.Alloc(&index, alloc_size);
+    ASSERT_TRUE(false);
+  } catch (paddle::memory::allocation::BadAlloc&) {
+    PADDLE_ENFORCE_GPU_SUCCESS(paddle::platform::GpuGetLastError());
+  }
+}
+#endif

test/cpp/fluid/memory/test_aligned_allocator.cc

@@ -0,0 +1,81 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "gtest/gtest.h"
+#include "paddle/phi/core/memory/allocation/aligned_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(aligned, aligned_size) {
+  ASSERT_EQ(AlignedSize(1024, 1024), 1024UL);
+  ASSERT_EQ(AlignedSize(1023, 1024), 1024UL);
+  ASSERT_EQ(AlignedSize(1025, 1024), 2048UL);
+}
+
+struct StubAllocator : public Allocator {
+ public:
+  StubAllocator() = default;
+
+  size_t AllocNum() const { return alloc_num_; }
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override {
+    ++alloc_num_;
+    return new Allocation(new uint8_t[size], size, phi::CPUPlace());
+  }
+
+  void FreeImpl(phi::Allocation *allocation) override {
+    delete[] static_cast<uint8_t *>(allocation->ptr());
+    delete allocation;
+    --alloc_num_;
+  }
+
+ private:
+  size_t alloc_num_{0};
+};
+
+bool IsAligned(const AllocationPtr &alloc, size_t alignment) {
+  return reinterpret_cast<uintptr_t>(alloc->ptr()) % alignment == 0;
+}
+
+TEST(aligned_allocator, aligned_allocator) {
+  size_t alignment = 1024;
+  auto allocator = std::make_shared<StubAllocator>();
+  auto aligned_allocator =
+      std::make_shared<AlignedAllocator>(allocator, alignment);
+
+  auto alloc1 = aligned_allocator->Allocate(1345);
+  ASSERT_EQ(allocator->AllocNum(), 1UL);
+  ASSERT_TRUE(IsAligned(alloc1, alignment));
+  alloc1.reset();
+  ASSERT_EQ(allocator->AllocNum(), 0UL);
+
+  {
+    auto alloc2 = aligned_allocator->Allocate(200);
+    ASSERT_TRUE(IsAligned(alloc2, alignment));
+    ASSERT_EQ(allocator->AllocNum(), 1UL);
+
+    auto alloc3 = aligned_allocator->Allocate(3021);
+    ASSERT_TRUE(IsAligned(alloc3, alignment));
+    ASSERT_EQ(allocator->AllocNum(), 2UL);
+  }
+
+  ASSERT_EQ(allocator->AllocNum(), 0UL);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/thread_local_allocator_test.cc

@@ -0,0 +1,92 @@
+// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/thread_local_allocator.h"
+
+#include <condition_variable>  // NOLINT
+#include <thread>              // NOLINT
+
+#include "gtest/gtest.h"
+#include "paddle/common/flags.h"
+#include "paddle/phi/core/memory/malloc.h"
+
+COMMON_DECLARE_double(fraction_of_gpu_memory_to_use);
+COMMON_DECLARE_string(allocator_strategy);
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(ThreadLocalAllocator, cross_scope_release) {
+  FLAGS_fraction_of_gpu_memory_to_use = 0.1;
+  FLAGS_allocator_strategy = "thread_local";
+
+  const size_t thread_num = 5;
+  const std::vector<int> devices = platform::GetSelectedDevices();
+
+  std::vector<std::vector<void *>> allocator_addresses(devices.size());
+  std::vector<std::vector<AllocationPtr>> thread_allocations(devices.size());
+
+  for (size_t i = 0; i < devices.size(); ++i) {
+    allocator_addresses[i].resize(thread_num);
+    thread_allocations[i].resize(thread_num);
+  }
+
+  std::vector<std::thread> threads(thread_num);
+  std::mutex mutex;
+  std::condition_variable cv;
+  bool flag = false;
+
+  for (size_t i = 0; i < threads.size(); ++i) {
+    threads[i] = std::thread([&, i]() {
+      {
+        std::unique_lock<std::mutex> lock(mutex);
+        cv.wait(lock, [&] { return flag; });
+      }
+      for (size_t j = 0; j < devices.size(); ++j) {
+        thread_allocations[j][i] = memory::Alloc(phi::GPUPlace(devices[j]), 10);
+        auto tl_allocator_impl =
+            ThreadLocalCUDAAllocatorPool::Instance().Get(devices[j]);
+        allocator_addresses[j][i] = tl_allocator_impl.get();
+        memory::Release(phi::GPUPlace(devices[j]));
+      }
+    });
+  }
+
+  {
+    std::lock_guard<std::mutex> lock(mutex);
+    flag = true;
+    cv.notify_all();
+  }
+
+  for (auto &th : threads) {
+    th.join();
+  }
+
+  for (auto &addresses : allocator_addresses) {
+    std::sort(addresses.begin(), addresses.end());
+    ASSERT_EQ(std::adjacent_find(
+                  addresses.begin(), addresses.end(), std::equal_to<>()),
+              addresses.end());
+  }
+
+  ::testing::FLAGS_gtest_death_test_style = "threadsafe";
+  ASSERT_EXIT(([&]() { thread_allocations.clear(); }(), exit(0)),
+              ::testing::ExitedWithCode(0),
+              ".*");
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/vmm_allocator_tail_test.cc

@@ -0,0 +1,55 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.h"
+#include "paddle/phi/core/memory/allocation/retry_allocator.h"
+#include "paddle/phi/core/memory/allocation/virtual_memory_auto_growth_best_fit_allocator.h"
+#include "paddle/phi/core/memory/memory.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+#include "glog/logging.h"
+#include "gtest/gtest.h"
+
+PD_DECLARE_bool(dump_vmm_allocation_info);
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(VirtualMemoryAutoGrowthBestFitAllocator, TestAllocatorVisitor) {
+  FLAGS_v = 4;
+  FLAGS_dump_vmm_allocation_info = true;
+  auto vmm_cuda_allocator =
+      std::make_shared<CUDAVirtualMemAllocator>(phi::GPUPlace());
+  auto vma_allocator =
+      std::make_shared<VirtualMemoryAutoGrowthBestFitAllocator>(
+          vmm_cuda_allocator, platform::GpuMinChunkSize(), phi::GPUPlace());
+  size_t mb = (1 << 20);
+  auto allocation1 = vma_allocator->Allocate(10 * mb);
+  auto allocation2 = vma_allocator->Allocate(20 * mb);
+  auto allocation_tiny = vma_allocator->Allocate(2 * mb - 1);
+  auto allocation3 = vma_allocator->Allocate(30 * mb);
+  auto allocation4 = vma_allocator->Allocate(40 * mb);
+  allocation2.reset();
+  allocation4.reset();
+  auto allocation5 = vma_allocator->Allocate(50 * mb);
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 114 * mb);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/vmm_auto_growth_best_fit_allocator_test.cu

@@ -0,0 +1,156 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <iostream>
+#include <memory>
+#include <string>
+
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.h"
+
+// Expose internals for white-box testing.
+#define private public
+#include "paddle/phi/core/memory/allocation/virtual_memory_auto_growth_best_fit_allocator.h"
+#undef private
+
+#include "gtest/gtest.h"
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/memory/memory.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class TestCUDAVirtualMemAllocator : public CUDAVirtualMemAllocator {
+ public:
+  using CUDAVirtualMemAllocator::CUDAVirtualMemAllocator;
+  using CUDAVirtualMemAllocator::FreeImpl;
+};
+
+TEST(test_vmm_allocator, test_mem_stats) {
+  size_t alignment = 256;
+  auto underlying_allocator =
+      std::make_shared<TestCUDAVirtualMemAllocator>(phi::GPUPlace());
+  auto allocation = underlying_allocator->Allocate(1024);
+  EXPECT_GT(DeviceMemoryStatCurrentValue("Reserved", 0), 1024);
+  allocation.reset();
+  EXPECT_EQ(DeviceMemoryStatCurrentValue("Reserved", 0), 0);
+}
+
+class DummyAllocator : public Allocator {
+ public:
+  bool IsAllocThreadSafe() const override { return true; }
+
+ protected:
+  phi::Allocation* AllocateImpl(size_t) override {
+    PADDLE_THROW(common::errors::Unavailable(
+        "DummyAllocator::AllocateImpl should not be called."));
+  }
+  void FreeImpl(phi::Allocation*) override {}
+};
+
+class AlwaysOOMAllocator : public Allocator {
+ public:
+  bool IsAllocThreadSafe() const override { return true; }
+
+ protected:
+  phi::Allocation* AllocateImpl(size_t size) override {
+    PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+        "AlwaysOOMAllocator failed to allocate %zu bytes.", size));
+  }
+  void FreeImpl(phi::Allocation*) override {}
+};
+
+// Expose FreeImpl for testing.
+class ExposedVmmAllocator : public VirtualMemoryAutoGrowthBestFitAllocator {
+ public:
+  using VirtualMemoryAutoGrowthBestFitAllocator::FreeImpl;
+  using VirtualMemoryAutoGrowthBestFitAllocator::
+      VirtualMemoryAutoGrowthBestFitAllocator;
+};
+
+TEST(test_vmm_allocator, free_impl_uses_allocation_block_iterator) {
+  auto underlying = std::make_shared<DummyAllocator>();
+  phi::GPUPlace place(0);
+  ExposedVmmAllocator allocator(underlying, 256, place);
+
+  // Manually construct blocks: [free-prev][used-target][free-next]
+  allocator.all_blocks_.clear();
+  auto prev = allocator.all_blocks_.emplace(
+      allocator.all_blocks_.end(), reinterpret_cast<void*>(0x1000), 1024, true);
+  auto target = allocator.all_blocks_.emplace(allocator.all_blocks_.end(),
+                                              reinterpret_cast<void*>(0x1400),
+                                              2048,
+                                              false);
+  auto next = allocator.all_blocks_.emplace(
+      allocator.all_blocks_.end(), reinterpret_cast<void*>(0x1C00), 4096, true);
+
+  allocator.free_blocks_.clear();
+  allocator.free_blocks_.emplace(std::make_pair(prev->size_, prev->ptr_), prev);
+  allocator.free_blocks_.emplace(std::make_pair(next->size_, next->ptr_), next);
+
+  auto allocation = std::make_unique<BlockAllocation>(target, place);
+
+  EXPECT_NO_THROW(allocator.FreeImpl(allocation.release()));
+  EXPECT_EQ(allocator.all_blocks_.size(), 1UL);
+  EXPECT_EQ(allocator.all_blocks_.front().ptr_,
+            reinterpret_cast<void*>(0x1000));
+  EXPECT_EQ(allocator.all_blocks_.front().size_, 7168UL);
+  EXPECT_TRUE(allocator.all_blocks_.front().is_free_);
+}
+
+TEST(test_vmm_allocator, oom_error_prints_pool_stats) {
+  auto underlying = std::make_shared<AlwaysOOMAllocator>();
+  phi::GPUPlace place(0);
+  ExposedVmmAllocator allocator(underlying, 256, place);
+
+  auto used1 = allocator.all_blocks_.emplace(allocator.all_blocks_.end(),
+                                             reinterpret_cast<void*>(0x1000),
+                                             1UL << 20,
+                                             false);
+  auto free1 = allocator.all_blocks_.emplace(allocator.all_blocks_.end(),
+                                             reinterpret_cast<void*>(0x101000),
+                                             2UL << 20,
+                                             true);
+  auto used2 = allocator.all_blocks_.emplace(allocator.all_blocks_.end(),
+                                             reinterpret_cast<void*>(0x301000),
+                                             1UL << 20,
+                                             false);
+  auto free2 = allocator.all_blocks_.emplace(allocator.all_blocks_.end(),
+                                             reinterpret_cast<void*>(0x401000),
+                                             4UL << 20,
+                                             true);
+
+  allocator.free_blocks_.emplace(std::make_pair(free1->size_, free1->ptr_),
+                                 free1);
+  allocator.free_blocks_.emplace(std::make_pair(free2->size_, free2->ptr_),
+                                 free2);
+
+  try {
+    allocator.Allocate(5UL << 20);
+    FAIL() << "Expected VMM allocator OOM.";
+  } catch (const BadAlloc& ex) {
+    std::string message = ex.what();
+    std::cout << "\n[VMM OOM MESSAGE]\n" << message << std::endl;
+    EXPECT_NE(message.find("VMM allocator stats (pool): "
+                           "total_free=6.000000MB, max_free=4.000000MB."),
+              std::string::npos);
+  }
+
+  static_cast<void>(used1);
+  static_cast<void>(used2);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

test/cpp/fluid/memory/vmm_auto_growth_best_fit_allocator_v2_test.cu

@@ -0,0 +1,644 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "gtest/gtest.h"
+
+#define private public
+#include "paddle/phi/core/memory/allocation/vmm_auto_growth_best_fit_allocator_v2.h"
+#undef private
+
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator_v2.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+namespace {
+
+std::shared_ptr<CUDAVirtualMemAllocatorV2> CreateUnderlyingAllocator() {
+  return std::make_shared<CUDAVirtualMemAllocatorV2>(
+      phi::GPUPlace(), 2UL << 20, PoolType::kSmall);
+}
+
+__global__ void DelayedStoreKernel(uint8_t* ptr, uint64_t cycles) {
+  uint64_t start = clock64();
+  while (clock64() - start < cycles) {
+  }
+  ptr[0] = 1;
+}
+
+void ExpectBlockView(const BlockV2& block) { EXPECT_GT(block.size_, 0UL); }
+
+}  // namespace
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, SplitFreeBlockOnReuse) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  auto large = allocator.Allocate(underlying->handle_size() * 2);
+  ASSERT_NE(large, nullptr);
+  large.reset();
+
+  auto small = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(small, nullptr);
+
+  ASSERT_EQ(allocator.all_blocks_.size(), 2UL);
+  size_t active_count = 0;
+  size_t free_count = 0;
+  size_t free_bytes = 0;
+  for (const auto& block : allocator.all_blocks_) {
+    if (block.type_ == BlockType::kActive) {
+      ++active_count;
+      EXPECT_EQ(block.size_, underlying->handle_size());
+    } else if (block.type_ == BlockType::kFree) {
+      ++free_count;
+      free_bytes += block.size_;
+      ExpectBlockView(block);
+    }
+  }
+  EXPECT_EQ(active_count, 1UL);
+  EXPECT_EQ(free_count, 1UL);
+  EXPECT_EQ(free_bytes, underlying->handle_size());
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, ReuseSmallestSufficientFreeBlock) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  // Layout after allocation:
+  //   [ACTIVE 4MB] [ACTIVE 2MB separator] [ACTIVE 2MB small]
+  // The separator prevents TryMerge from coalescing large and small on free.
+  auto large = allocator.Allocate(underlying->handle_size() * 2);
+  auto separator = allocator.Allocate(underlying->handle_size());
+  auto small = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(large, nullptr);
+  ASSERT_NE(separator, nullptr);
+  ASSERT_NE(small, nullptr);
+
+  auto* small_ptr = small->ptr();
+  large.reset();
+  small.reset();
+  // Layout: [FREE 4MB] [ACTIVE 2MB separator] [FREE 2MB]
+  // free_blocks_: {(2MB, ptr_small), (4MB, ptr_large)}
+
+  auto reused = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(reused, nullptr);
+
+  // lower_bound({2MB, nullptr}) picks the exact-fit 2MB free block over the
+  // larger 4MB one.
+  EXPECT_EQ(reused->ptr(), small_ptr);
+  // Layout: [FREE 4MB] [ACTIVE 2MB separator] [ACTIVE 2MB reused]
+  ASSERT_EQ(allocator.all_blocks_.size(), 3UL);
+  size_t free_block_count = 0;
+  for (const auto& block : allocator.all_blocks_) {
+    if (block.type_ == BlockType::kFree) {
+      ++free_block_count;
+      EXPECT_EQ(block.size_, underlying->handle_size() * 2);
+    }
+  }
+  EXPECT_EQ(free_block_count, 1UL);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, ReturnedAllocationSizeMatchesRequest) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  auto allocation = allocator.Allocate(256UL);
+  ASSERT_NE(allocation, nullptr);
+
+  EXPECT_EQ(allocation->size(), 256UL);
+  auto* alloc = static_cast<Allocation*>(allocation.get());
+  EXPECT_EQ(alloc->ptr(), alloc->base_ptr());
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, SplitGrowBlockAcrossTwoHandles) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  const size_t requested_size = underlying->handle_size() + 256UL;
+  auto allocation = allocator.Allocate(requested_size);
+  ASSERT_NE(allocation, nullptr);
+
+  ASSERT_EQ(allocator.all_blocks_.size(), 2UL);
+  auto it = allocator.all_blocks_.begin();
+  ASSERT_EQ(it->type_, BlockType::kActive);
+  EXPECT_EQ(it->size_, requested_size);
+  ExpectBlockView(*it);
+
+  ++it;
+  ASSERT_EQ(it, std::prev(allocator.all_blocks_.end()));
+  ASSERT_EQ(it->type_, BlockType::kFree);
+  EXPECT_EQ(it->size_, underlying->handle_size() - 256UL);
+  ExpectBlockView(*it);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, MergeSplitFreeSlicesAsBlockView) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  auto allocation = allocator.Allocate(256UL);
+  ASSERT_NE(allocation, nullptr);
+  allocation.reset();
+
+  ASSERT_EQ(allocator.all_blocks_.size(), 1UL);
+  const auto& merged = allocator.all_blocks_.front();
+  EXPECT_EQ(merged.type_, BlockType::kFree);
+  EXPECT_EQ(merged.size_, underlying->handle_size());
+  ExpectBlockView(merged);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, GrowExactHandleMultipleNoSplit) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  // Request exactly 1 handle_size — the bottom allocator returns the same
+  // amount, so grow-split should produce NO remaining FREE block.
+  auto allocation = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(allocation, nullptr);
+
+  EXPECT_EQ(allocator.all_blocks_.size(), 1UL);
+  EXPECT_EQ(allocator.all_blocks_.front().type_, BlockType::kActive);
+  EXPECT_EQ(allocator.all_blocks_.front().size_, underlying->handle_size());
+  EXPECT_EQ(allocator.free_blocks_.size(), 0UL);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, AlignmentRoundsUpRequestedSize) {
+  auto underlying = CreateUnderlyingAllocator();
+  const size_t alignment = 512;
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, alignment, phi::GPUPlace(), PoolType::kSmall);
+
+  // Request 100 bytes with alignment=512 → AlignedSize(100,512) = 512.
+  auto allocation = allocator.Allocate(100);
+  ASSERT_NE(allocation, nullptr);
+
+  // The returned allocation size must be the aligned 512, not 100.
+  EXPECT_EQ(allocation->size(), 512UL);
+
+  // The ACTIVE block in all_blocks_ should also be 512.
+  auto it = allocator.all_blocks_.begin();
+  ASSERT_EQ(it->type_, BlockType::kActive);
+  EXPECT_EQ(it->size_, 512UL);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, ExactFitReuseNoSplit) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  // Allocate and free one handle — creates one FREE block of handle_size.
+  auto allocation = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(allocation, nullptr);
+  auto* original_ptr = allocation->ptr();
+  allocation.reset();
+  ASSERT_EQ(allocator.free_blocks_.size(), 1UL);
+
+  // Re-allocate exactly the same size — exact fit, no split needed.
+  auto reused = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(reused, nullptr);
+  EXPECT_EQ(reused->ptr(), original_ptr);
+
+  // Only one block: ACTIVE, no FREE remainder.
+  EXPECT_EQ(allocator.all_blocks_.size(), 1UL);
+  EXPECT_EQ(allocator.all_blocks_.front().type_, BlockType::kActive);
+  EXPECT_EQ(allocator.free_blocks_.size(), 0UL);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, AllocFreeCycleConsistency) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  // Perform several alloc/free cycles and verify invariants after each.
+  for (int round = 0; round < 3; ++round) {
+    auto a1 = allocator.Allocate(underlying->handle_size());
+    auto a2 = allocator.Allocate(underlying->handle_size());
+    ASSERT_NE(a1, nullptr);
+    ASSERT_NE(a2, nullptr);
+    size_t active_before_free = 0;
+    for (const auto& block : allocator.all_blocks_) {
+      if (block.type_ == BlockType::kActive) {
+        ++active_before_free;
+      }
+    }
+    EXPECT_EQ(active_before_free, 2UL);
+
+    a1.reset();
+    a2.reset();
+    // After freeing all, adjacent blocks merge — should be exactly 1 FREE.
+    EXPECT_EQ(allocator.free_blocks_.size(), 1UL);
+
+    size_t total_free = 0;
+    for (const auto& block : allocator.all_blocks_) {
+      EXPECT_EQ(block.type_, BlockType::kFree);
+      total_free += block.size_;
+    }
+    EXPECT_EQ(total_free, underlying->handle_size() * 2);
+  }
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, FreeBlockTooSmallFallsBackToGrow) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  // Create a small free block (handle_size).
+  auto small = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(small, nullptr);
+  small.reset();
+  ASSERT_EQ(allocator.free_blocks_.size(), 1UL);
+
+  // Request 2*handle_size — free block is too small, must grow.
+  auto large = allocator.Allocate(underlying->handle_size() * 2);
+  ASSERT_NE(large, nullptr);
+
+  // The old full underlying allocation is idle and may be released before grow.
+  EXPECT_EQ(allocator.free_blocks_.size(), 0UL);
+
+  // Verify total layout: only the new ACTIVE block remains.
+  size_t active_count = 0;
+  size_t free_count = 0;
+  for (const auto& block : allocator.all_blocks_) {
+    if (block.type_ == BlockType::kActive) {
+      ++active_count;
+      EXPECT_EQ(block.size_, underlying->handle_size() * 2);
+    } else if (block.type_ == BlockType::kFree) {
+      ++free_count;
+    }
+  }
+  EXPECT_EQ(active_count, 1UL);
+  EXPECT_EQ(free_count, 0UL);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2,
+     ReleaseIdleMiddleChunkLeavesReusableUnmappedFreeBlock) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  auto first = allocator.Allocate(underlying->handle_size());
+  auto middle = allocator.Allocate(underlying->handle_size());
+  auto last = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(first, nullptr);
+  ASSERT_NE(middle, nullptr);
+  ASSERT_NE(last, nullptr);
+
+  auto* middle_ptr = middle->ptr();
+  const size_t tail_after_allocs = underlying->tail_offset();
+  middle.reset();
+
+  EXPECT_EQ(allocator.Release(phi::GPUPlace()), underlying->handle_size());
+  ASSERT_EQ(allocator.all_blocks_.size(), 3UL);
+  auto it = allocator.all_blocks_.begin();
+  EXPECT_TRUE(it->IsActive());
+  ++it;
+  ASSERT_TRUE(it->IsUnmappedFree());
+  EXPECT_EQ(it->ptr_, middle_ptr);
+  EXPECT_EQ(it->size_, underlying->handle_size());
+  ++it;
+  EXPECT_TRUE(it->IsActive());
+
+  auto reused = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(reused, nullptr);
+  EXPECT_EQ(reused->ptr(), middle_ptr);
+  EXPECT_EQ(underlying->tail_offset(), tail_after_allocs);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, ReleaseTailChunkRetreatsTailOffset) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+  const size_t handle_size = underlying->handle_size();
+
+  auto head = allocator.Allocate(handle_size);
+  auto tail = allocator.Allocate(handle_size);
+  ASSERT_NE(head, nullptr);
+  ASSERT_NE(tail, nullptr);
+  auto* expected_next_tail =
+      reinterpret_cast<uint8_t*>(head->ptr()) + handle_size;
+
+  tail.reset();
+  EXPECT_EQ(allocator.Release(phi::GPUPlace()), handle_size);
+  EXPECT_EQ(underlying->tail_offset(), handle_size);
+  EXPECT_EQ(allocator.all_blocks_.size(), 1UL);
+  EXPECT_EQ(allocator.unmapped_free_blocks_.size(), 0UL);
+
+  auto grow = allocator.Allocate(handle_size * 2);
+  ASSERT_NE(grow, nullptr);
+  EXPECT_EQ(grow->ptr(), expected_next_tail);
+  EXPECT_EQ(underlying->tail_offset(), handle_size * 3);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, ReleaseWithNoIdleChunkReturnsZero) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  auto active = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(active, nullptr);
+  const size_t tail_before_release = underlying->tail_offset();
+
+  EXPECT_EQ(allocator.Release(phi::GPUPlace()), 0UL);
+  EXPECT_EQ(underlying->tail_offset(), tail_before_release);
+  ASSERT_EQ(allocator.all_blocks_.size(), 1UL);
+  EXPECT_TRUE(allocator.all_blocks_.front().IsActive());
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2,
+     RangeOverlapsUnderlyingCoversRegistryCases) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+  const size_t handle_size = underlying->handle_size();
+
+  auto allocation = allocator.Allocate(handle_size);
+  ASSERT_NE(allocation, nullptr);
+  auto* ptr = reinterpret_cast<uint8_t*>(allocation->ptr());
+
+  EXPECT_TRUE(allocator.RangeOverlapsUnderlying(ptr, handle_size));
+  EXPECT_TRUE(allocator.RangeOverlapsUnderlying(ptr + handle_size / 2,
+                                                handle_size / 2));
+  EXPECT_FALSE(
+      allocator.RangeOverlapsUnderlying(ptr + handle_size, handle_size));
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, ReleasePredicatesRejectActiveAllocation) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  auto allocation = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(allocation, nullptr);
+  auto* ptr = reinterpret_cast<uint8_t*>(allocation->ptr());
+
+  EXPECT_FALSE(allocator.IsRangeEntirelyFree(ptr, underlying->handle_size()));
+  EXPECT_FALSE(
+      allocator.CanReleaseIdleUnderlying(ptr, underlying->handle_size()));
+  uint64_t released = 0;
+  auto it = allocator.underlying_allocations_.begin();
+  EXPECT_FALSE(allocator.TryReleaseIdleUnderlying(&it, &released));
+  EXPECT_EQ(released, 0UL);
+  EXPECT_EQ(allocator.FreeIdleChunks(), 0UL);
+
+  allocation.reset();
+  EXPECT_TRUE(allocator.IsRangeEntirelyFree(ptr, underlying->handle_size()));
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, FreeIndexHelpersIgnoreWrongBlockTypes) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  BlockV2 active = BlockV2::MakeMappedBlock(BlockType::kActive,
+                                            reinterpret_cast<void*>(0x30000000),
+                                            underlying->handle_size(),
+                                            PoolType::kSmall);
+  auto active_it = allocator.all_blocks_.insert(allocator.all_blocks_.end(),
+                                                std::move(active));
+  allocator.InsertFreeBlock(active_it);
+  EXPECT_TRUE(allocator.free_blocks_.empty());
+
+  BlockV2 mapped_free =
+      BlockV2::MakeMappedBlock(BlockType::kFree,
+                               reinterpret_cast<void*>(0x32000000),
+                               underlying->handle_size(),
+                               PoolType::kSmall);
+  auto mapped_free_it = allocator.all_blocks_.insert(
+      allocator.all_blocks_.end(), std::move(mapped_free));
+  allocator.InsertUnmappedFreeBlock(mapped_free_it);
+  EXPECT_TRUE(allocator.unmapped_free_blocks_.empty());
+
+  allocator.TryMergeUnmappedFree(allocator.all_blocks_.end());
+  EXPECT_TRUE(allocator.unmapped_free_blocks_.empty());
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2,
+     ReleaseAdjacentMiddleChunksMergeIntoSingleUnmappedFreeBlock) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+  const size_t handle_size = underlying->handle_size();
+
+  auto first = allocator.Allocate(handle_size);
+  auto middle_left = allocator.Allocate(handle_size);
+  auto middle_right = allocator.Allocate(handle_size);
+  auto last = allocator.Allocate(handle_size);
+  ASSERT_NE(first, nullptr);
+  ASSERT_NE(middle_left, nullptr);
+  ASSERT_NE(middle_right, nullptr);
+  ASSERT_NE(last, nullptr);
+
+  auto* middle_ptr = middle_left->ptr();
+  middle_left.reset();
+  middle_right.reset();
+
+  EXPECT_EQ(allocator.Release(phi::GPUPlace()), handle_size * 2);
+  ASSERT_EQ(allocator.all_blocks_.size(), 3UL);
+  auto it = allocator.all_blocks_.begin();
+  EXPECT_TRUE(it->IsActive());
+  ++it;
+  ASSERT_TRUE(it->IsUnmappedFree());
+  EXPECT_EQ(it->ptr_, middle_ptr);
+  EXPECT_EQ(it->size_, handle_size * 2);
+  ++it;
+  EXPECT_TRUE(it->IsActive());
+  EXPECT_EQ(allocator.unmapped_free_blocks_.size(), 1UL);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2,
+     ReuseUnmappedFreeBlockWithMappedAndUnmappedRemainders) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+  const size_t handle_size = underlying->handle_size();
+
+  auto first = allocator.Allocate(handle_size);
+  auto middle = allocator.Allocate(handle_size * 2);
+  auto last = allocator.Allocate(handle_size);
+  ASSERT_NE(first, nullptr);
+  ASSERT_NE(middle, nullptr);
+  ASSERT_NE(last, nullptr);
+
+  auto* middle_ptr = middle->ptr();
+  middle.reset();
+  EXPECT_EQ(allocator.Release(phi::GPUPlace()), handle_size * 2);
+  ASSERT_EQ(allocator.unmapped_free_blocks_.size(), 1UL);
+
+  auto reused = allocator.Allocate(256UL);
+  ASSERT_NE(reused, nullptr);
+  EXPECT_EQ(reused->ptr(), middle_ptr);
+
+  size_t active_count = 0;
+  size_t mapped_free_bytes = 0;
+  size_t unmapped_free_bytes = 0;
+  for (const auto& block : allocator.all_blocks_) {
+    if (block.IsActive()) {
+      ++active_count;
+    } else if (block.IsMappedFree()) {
+      mapped_free_bytes += block.size_;
+    } else if (block.IsUnmappedFree()) {
+      unmapped_free_bytes += block.size_;
+    }
+  }
+  EXPECT_EQ(active_count, 3UL);
+  EXPECT_EQ(mapped_free_bytes, handle_size - 256UL);
+  EXPECT_EQ(unmapped_free_bytes, handle_size);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2,
+     ReplaceRangeWithUnmappedFreeSplitsContainingFreeBlock) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+  const size_t handle_size = underlying->handle_size();
+
+  auto allocation = allocator.Allocate(handle_size * 3);
+  ASSERT_NE(allocation, nullptr);
+  auto* base = reinterpret_cast<uint8_t*>(allocation->ptr());
+  allocation.reset();
+  ASSERT_EQ(allocator.all_blocks_.size(), 1UL);
+
+  allocator.ReplaceRangeWithUnmappedFree(base + handle_size, handle_size);
+
+  ASSERT_EQ(allocator.all_blocks_.size(), 3UL);
+  auto it = allocator.all_blocks_.begin();
+  EXPECT_TRUE(it->IsMappedFree());
+  EXPECT_EQ(it->size_, handle_size);
+  ++it;
+  ASSERT_TRUE(it->IsUnmappedFree());
+  EXPECT_EQ(it->ptr_, base + handle_size);
+  EXPECT_EQ(it->size_, handle_size);
+  ++it;
+  EXPECT_TRUE(it->IsMappedFree());
+  EXPECT_EQ(it->size_, handle_size);
+  EXPECT_EQ(allocator.free_blocks_.size(), 2UL);
+  EXPECT_EQ(allocator.unmapped_free_blocks_.size(), 1UL);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2,
+     ReplaceRangeWithUnmappedFreeKeepsLeftRemainder) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+  const size_t handle_size = underlying->handle_size();
+
+  auto allocation = allocator.Allocate(handle_size * 3);
+  ASSERT_NE(allocation, nullptr);
+  auto* base = reinterpret_cast<uint8_t*>(allocation->ptr());
+  allocation.reset();
+
+  allocator.ReplaceRangeWithUnmappedFree(base + handle_size, handle_size * 2);
+
+  ASSERT_EQ(allocator.all_blocks_.size(), 2UL);
+  auto it = allocator.all_blocks_.begin();
+  EXPECT_TRUE(it->IsMappedFree());
+  EXPECT_EQ(it->ptr_, base);
+  EXPECT_EQ(it->size_, handle_size);
+  ++it;
+  EXPECT_TRUE(it->IsUnmappedFree());
+  EXPECT_EQ(it->ptr_, base + handle_size);
+  EXPECT_EQ(it->size_, handle_size * 2);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2,
+     ReplaceRangeWithUnmappedFreeKeepsRightRemainder) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+  const size_t handle_size = underlying->handle_size();
+
+  auto allocation = allocator.Allocate(handle_size * 3);
+  ASSERT_NE(allocation, nullptr);
+  auto* base = reinterpret_cast<uint8_t*>(allocation->ptr());
+  allocation.reset();
+
+  allocator.ReplaceRangeWithUnmappedFree(base, handle_size * 2);
+
+  ASSERT_EQ(allocator.all_blocks_.size(), 2UL);
+  auto it = allocator.all_blocks_.begin();
+  EXPECT_TRUE(it->IsUnmappedFree());
+  EXPECT_EQ(it->ptr_, base);
+  EXPECT_EQ(it->size_, handle_size * 2);
+  ++it;
+  EXPECT_TRUE(it->IsMappedFree());
+  EXPECT_EQ(it->ptr_, base + handle_size * 2);
+  EXPECT_EQ(it->size_, handle_size);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, AdoptBackingBlockRejectsNullInput) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  EXPECT_THROW(allocator.AdoptBackingBlock(nullptr),
+               common::enforce::EnforceNotMet);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, ReleaseWaitsBeforeUnmappingBacking) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  auto allocation = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(allocation, nullptr);
+  auto* ptr = reinterpret_cast<uint8_t*>(allocation->ptr());
+
+  DelayedStoreKernel<<<1, 1>>>(ptr, 20000000ULL);
+  ASSERT_EQ(cudaGetLastError(), cudaSuccess);
+
+  allocation.reset();
+  EXPECT_EQ(allocator.Release(phi::GPUPlace()), underlying->handle_size());
+  EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+}
+
+TEST(VMMAutoGrowthBestFitAllocatorV2, ThreeWayMerge) {
+  auto underlying = CreateUnderlyingAllocator();
+  VMMAutoGrowthBestFitAllocatorV2 allocator(
+      underlying, 256, phi::GPUPlace(), PoolType::kSmall);
+
+  // Allocate 3 consecutive handle-sized blocks.
+  auto a = allocator.Allocate(underlying->handle_size());
+  auto b = allocator.Allocate(underlying->handle_size());
+  auto c = allocator.Allocate(underlying->handle_size());
+  ASSERT_NE(a, nullptr);
+  ASSERT_NE(b, nullptr);
+  ASSERT_NE(c, nullptr);
+  ASSERT_EQ(allocator.all_blocks_.size(), 3UL);
+
+  // Free first and last — creates 2 non-adjacent FREE blocks.
+  a.reset();
+  c.reset();
+  EXPECT_EQ(allocator.free_blocks_.size(), 2UL);
+
+  // Free middle — TryMerge merges prev+it (left), then merged+next (right)
+  // into a single block spanning all 3 handles.
+  b.reset();
+  EXPECT_EQ(allocator.all_blocks_.size(), 1UL);
+  EXPECT_EQ(allocator.free_blocks_.size(), 1UL);
+
+  const auto& merged = allocator.all_blocks_.front();
+  EXPECT_EQ(merged.type_, BlockType::kFree);
+  EXPECT_EQ(merged.size_, underlying->handle_size() * 3);
+  ExpectBlockView(merged);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle