xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-28 12:07:04 +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 = torch.cuda.CUDAGraph()
graph_pool = torch.cuda.graph_pool_handle() graph_pool = torch.cuda.graph_pool_handle()
set_graph_pool_id(graph_pool) 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): for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):
allreduce_fn(graph_input) allreduce_fn(graph_input)

6
vllm/compilation/cuda_graph.py
View File

@ -169,7 +169,11 @@ class CUDAGraphWrapper:
else: else:
set_graph_pool_id(current_platform.graph_pool_handle()) set_graph_pool_id(current_platform.graph_pool_handle())
# mind-exploding: carefully manage the reference and memory. # 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` is managed by pytorch's cudagraph pool
output = self.runnable(*args, **kwargs) output = self.runnable(*args, **kwargs)
if self.cudagraph_options.weak_ref_output: if self.cudagraph_options.weak_ref_output: