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

Merge fd593c596de73fcb98ee8636150a385ea896b43d into 254f6b986720c92ddf97fbb1a6a6465da8e87e29

Browse Source
This commit is contained in:
Amir Samani 2025-12-25 00:06:49 +00:00 committed by GitHub
commit 3bb2e7950a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 16 additions and 8 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)

16
tests/v1/cudagraph/test_cudagraph_dispatch.py
View File

@ -189,13 +189,16 @@ class TestCUDAGraphWrapper:
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
) )
batch_descriptor = BatchDescriptor(num_tokens=10) batch_descriptor = BatchDescriptor(num_tokens=10)
stream = torch.cuda.Stream()
# 0. global warmup # 0. global warmup
with set_forward_context( with (
attn_metadata=None, set_forward_context(
vllm_config=self.vllm_config, attn_metadata=None,
cudagraph_runtime_mode=CUDAGraphMode.NONE, vllm_config=self.vllm_config,
batch_descriptor=None, cudagraph_runtime_mode=CUDAGraphMode.NONE,
batch_descriptor=None,
),
torch.cuda.stream(stream),
): ):
wrapper(self.input_tensor) wrapper(self.input_tensor)
@ -208,6 +211,7 @@ class TestCUDAGraphWrapper:
batch_descriptor=batch_descriptor, batch_descriptor=batch_descriptor,
), ),
patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_cuda_graph, patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_cuda_graph,
torch.cuda.stream(stream),
): ):
output1 = wrapper(self.input_tensor) output1 = wrapper(self.input_tensor)
# capturing phase should generate a zero output # capturing phase should generate a zero output

6
vllm/compilation/cuda_graph.py
View File

@ -263,7 +263,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: