xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-03-23 01:05:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

add stream to cuda graph catpure

Browse Source 
Signed-off-by: Amir Samani <asamani@nvidia.com>
This commit is contained in:
Amir Samani 2025-11-21 14:20:08 -08:00
parent 57430fc95c
commit 28579b55fa
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
benchmarks/kernels/benchmark_device_communicators.py
View File

@ -293,7 +293,7 @@ class CommunicatorBenchmark:
graph = torch.cuda.CUDAGraph()
graph_pool = torch.cuda.graph_pool_handle()
set_graph_pool_id(graph_pool)
with torch.cuda.graph(graph, pool=graph_pool):
with torch.cuda.graph(graph, pool=graph_pool, stream=stream):
for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):
allreduce_fn(graph_input)

6
vllm/compilation/cuda_graph.py
View File

@ -169,7 +169,11 @@ class CUDAGraphWrapper:
else:
set_graph_pool_id(current_platform.graph_pool_handle())
# mind-exploding: carefully manage the reference and memory.
with torch.cuda.graph(cudagraph, pool=self.graph_pool):
with torch.cuda.graph(
cudagraph,
pool=self.graph_pool,
stream=torch.cuda.current_stream(),
):
# `output` is managed by pytorch's cudagraph pool
output = self.runnable(*args, **kwargs)
if self.cudagraph_options.weak_ref_output: