Update cudaGraphInstantiate to use node priority flag

2025-11-12 06:44:55 +08:00 · 2025-11-07 10:54:11 -08:00
1 changed files with 2 additions and 2 deletions
--- a/aten/src/ATen/cuda/CUDAGraph.cpp
+++ b/aten/src/ATen/cuda/CUDAGraph.cpp
@ -175,7 +175,7 @@ void CUDAGraph::instantiate() {
    // who prefer not to report error message through these arguments moving forward
    // (they prefer return value, or errors on api calls internal to the capture)
 #if (defined(CUDA_VERSION) && CUDA_VERSION >= 12000)
-    AT_CUDA_CHECK(cudaGraphInstantiate(&graph_exec_, graph_, 0));
+    AT_CUDA_CHECK(cudaGraphInstantiate(&graph_exec_, graph_, cudaGraphInstantiateFlagUseNodePriority));
 #else
    AT_CUDA_CHECK(cudaGraphInstantiate(&graph_exec_, graph_, NULL, NULL, 0));
 #endif
@ -184,7 +184,7 @@ void CUDAGraph::instantiate() {
  } else {
    AT_CUDA_CHECK(cudaGraphInstantiateWithFlags(&graph_exec_,
                                                graph_,
-                                                cudaGraphInstantiateFlagAutoFreeOnLaunch));
+                                                cudaGraphInstantiateFlagAutoFreeOnLaunch | cudaGraphInstantiateFlagUseNodePriority));
  }
  has_graph_exec_ = true;
 }