apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2025-12-24 13:28:13 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

[CUDAGraph]Add support for custom all-reduce operators under SOT mode (#4386)

Browse Source
This commit is contained in:
Ryan
2025-10-16 19:31:19 +08:00
committed by GitHub
parent 26ff2f8683
commit b87e2c6184
3 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
fastdeploy/model_executor/graph_optimization/cudagraph_piecewise_backend.py
View File

@@ -133,8 +133,9 @@ class CudaGraphPiecewiseBackend:
self.cuda_graph_manager.state = jit_utils.CUDAGraphState.CAPTURE
self.cuda_graph_manager.batch_size = entry.real_shape
entry.captured = True
with self.cuda_graph_manager.run_impl_guard():
entry.runnable(**kwargs)
with capture_custom_allreduce():
with self.cuda_graph_manager.run_impl_guard():
entry.runnable(**kwargs)
# Replay
self.cuda_graph_manager.state = jit_utils.CUDAGraphState.REPLAY