xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-03-25 21:25:46 +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
)
batch_descriptor = BatchDescriptor(num_tokens=10)
stream = torch.cuda.Stream()
# 0. global warmup
with (
set_forward_context(
@ -175,7 +175,7 @@ class TestCUDAGraphWrapper:
cudagraph_runtime_mode=CUDAGraphMode.NONE,
batch_descriptor=None,
),
graph_capture(device=torch.device("cuda")) as graph_ctx,
torch.cuda.stream(stream),
):
wrapper(self.input_tensor)
@ -188,7 +188,7 @@ class TestCUDAGraphWrapper:
batch_descriptor=batch_descriptor,
),
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)
# capturing phase should generate a zero output