frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-13 18:25:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update cudaGraphInstantiate to use node priority flag

Browse Source
This commit is contained in:
eqy
2025-11-07 10:54:11 -08:00
committed by GitHub
parent e401a56b96
commit e10ddca827
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
aten/src/ATen/cuda/CUDAGraph.cpp
View File

@ -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;
}