PaddlePaddle
diff --git a/‎paddle/common/flags.cc
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions b/‎paddle/common/flags.cc
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions
diff --git a/‎paddle/fluid/distributed/collective/process_group_nccl.cc
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions b/‎paddle/fluid/distributed/collective/process_group_nccl.cc
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions
diff --git a/‎paddle/fluid/distributed/collective/reducer.cc
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/distributed/collective/reducer.cc
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/fluid/memory/allocation/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/memory/allocation/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
@@ -1090,6 +1090,34 @@ PHI_DEFINE_EXPORTED_bool(new_executor_use_cuda_graph,
                          false,
                          "Use CUDA Graph in new executor");
 
+/*
+ * CUDA Graph / Allocator related FLAG
+ * Name: FLAGS_use_cuda_malloc_async_allocator
+ * Since Version: 2.7
+ * Value Range: bool, default=false
+ * Example: FLAGS_use_cuda_malloc_async_allocator=true would allow
+ * CUDAMallocAsyncAllocator replace StreamSafeCUDAAllocator.
+ */
+PHI_DEFINE_EXPORTED_bool(use_cuda_malloc_async_allocator,
+                         false,
+                         "Enable CUDAMallocAsyncAllocator");
+
+/*
+ * CUDA Graph / Allocator related FLAG
+ * Name: FLAGS_auto_free_cudagraph_allocations_on_launch
+ * Since Version: 2.7
+ * Value Range: bool, default=true
+ * Example: When enabling CUDA Graph with CUDAMallocAsyncAllocator, we add
+ * cudaGraphInstantiateFlagAutoFreeOnLaunch so it would automatically
+ * release graph-owned blocks that have not freed before relaunching.
+ */
+PHI_DEFINE_EXPORTED_bool(
+    auto_free_cudagraph_allocations_on_launch,
+    true,
+    "When enabling CUDA Graph with CUDAMallocAsyncAllocator, we add "
+    "cudaGraphInstantiateFlagAutoFreeOnLaunch so it would automatically "
+    "release graph-owned blocks that have not freed before relaunching.");
+
 /*
  * Executor related FLAG
  * Name: FLAGS_executor_log_deps_every_microseconds
 
@@ -35,6 +35,7 @@ COMMON_DECLARE_bool(benchmark);
 COMMON_DECLARE_bool(benchmark_nccl);
 COMMON_DECLARE_bool(nccl_blocking_wait);
 COMMON_DECLARE_bool(use_stream_safe_cuda_allocator);
+COMMON_DECLARE_bool(use_cuda_malloc_async_allocator);
 COMMON_DECLARE_bool(enable_async_trace);
 
 // set this flag to `true` and recompile to enable dynamic checks
@@ -859,7 +860,8 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupNCCL::Collective(
   }
 
   if (!use_calc_stream) {
-    if (FLAGS_use_stream_safe_cuda_allocator) {
+    if (FLAGS_use_stream_safe_cuda_allocator ||
+        FLAGS_use_cuda_malloc_async_allocator) {
       memory::RecordStream(tensor_tmp.Holder(), nccl_stream);
     }
     task->UpdateWaitChain(*comm_ctx);
@@ -974,7 +976,8 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupNCCL::Point2Point(
   }
 
   if (!use_calc_stream) {
-    if (FLAGS_use_stream_safe_cuda_allocator) {
+    if (FLAGS_use_stream_safe_cuda_allocator ||
+        FLAGS_use_cuda_malloc_async_allocator) {
       memory::RecordStream(tensor_tmp.Holder(), nccl_stream);
     }
     task->UpdateWaitChain(*comm_ctx);
 
@@ -25,8 +25,8 @@ namespace paddle {
 namespace distributed {
 
 static bool IsStreamSafeAllocator() {
-  return FLAGS_allocator_strategy == "auto_growth" &&
-         FLAGS_use_stream_safe_cuda_allocator;
+  return (FLAGS_allocator_strategy == "auto_growth" &&
+          FLAGS_use_stream_safe_cuda_allocator);
 }
 
 static Backend TransToBackend(platform::Place place) {
 
@@ -25,6 +25,7 @@ if(WITH_GPU OR WITH_ROCM)
     ALLOCATOR_SRCS
     cuda_allocator.cc
     cuda_managed_allocator.cc
+    cuda_malloc_async_allocator.cc
     pinned_allocator.cc
     stream_safe_cuda_allocator.cc
     thread_local_allocator.cc)
Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,7 @@ COMMON_DECLARE_bool(benchmark);`
`35`	`35`	`COMMON_DECLARE_bool(benchmark_nccl);`
`36`	`36`	`COMMON_DECLARE_bool(nccl_blocking_wait);`
`37`	`37`	`COMMON_DECLARE_bool(use_stream_safe_cuda_allocator);`
	`38`	`+COMMON_DECLARE_bool(use_cuda_malloc_async_allocator);`
`38`	`39`	`COMMON_DECLARE_bool(enable_async_trace);`
`39`	`40`
`40`	`41`	// set this flag to `true` and recompile to enable dynamic checks
`@@ -859,7 +860,8 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupNCCL::Collective(`
`859`	`860`	`}`
`860`	`861`
`861`	`862`	`if (!use_calc_stream) {`
`862`		`- if (FLAGS_use_stream_safe_cuda_allocator) {`
	`863`	`+ if (FLAGS_use_stream_safe_cuda_allocator \|\|`
	`864`	`+ FLAGS_use_cuda_malloc_async_allocator) {`
`863`	`865`	`memory::RecordStream(tensor_tmp.Holder(), nccl_stream);`
`864`	`866`	`}`
`865`	`867`	`task->UpdateWaitChain(*comm_ctx);`
`@@ -974,7 +976,8 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupNCCL::Point2Point(`
`974`	`976`	`}`
`975`	`977`
`976`	`978`	`if (!use_calc_stream) {`
`977`		`- if (FLAGS_use_stream_safe_cuda_allocator) {`
	`979`	`+ if (FLAGS_use_stream_safe_cuda_allocator \|\|`
	`980`	`+ FLAGS_use_cuda_malloc_async_allocator) {`
`978`	`981`	`memory::RecordStream(tensor_tmp.Holder(), nccl_stream);`
`979`	`982`	`}`
`980`	`983`	`task->UpdateWaitChain(*comm_ctx);`
Original file line number	Diff line number	Diff line change
`@@ -25,8 +25,8 @@ namespace paddle {`
`25`	`25`	`namespace distributed {`
`26`	`26`
`27`	`27`	`static bool IsStreamSafeAllocator() {`
`28`		`- return FLAGS_allocator_strategy == "auto_growth" &&`
`29`		`- FLAGS_use_stream_safe_cuda_allocator;`
	`28`	`+ return (FLAGS_allocator_strategy == "auto_growth" &&`
	`29`	`+ FLAGS_use_stream_safe_cuda_allocator);`
`30`	`30`	`}`
`31`	`31`
`32`	`32`	`static Backend TransToBackend(platform::Place place) {`