xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-31 22:47:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

wip

Browse Source 
Signed-off-by: Amir Samani <asamani@nvidia.com>
This commit is contained in:
Amir Samani 2025-11-24 13:48:31 -08:00
parent 695d78c471
commit dfb5ce44dc
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -166,7 +166,7 @@ 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 ( with (
set_forward_context( set_forward_context(
@ -175,7 +175,7 @@ class TestCUDAGraphWrapper:
cudagraph_runtime_mode=CUDAGraphMode.NONE, cudagraph_runtime_mode=CUDAGraphMode.NONE,
batch_descriptor=None, batch_descriptor=None,
), ),
graph_capture(device=torch.device("cuda")) as graph_ctx, torch.cuda.stream(stream),
): ):
wrapper(self.input_tensor) wrapper(self.input_tensor)
@ -188,7 +188,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,
graph_capture(device=torch.device("cuda")) as graph_ctx, 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