xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-31 05:57:04 +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 11:23:23 -08:00
parent 28579b55fa
commit 695d78c471
1 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -17,6 +17,7 @@ from vllm.config import (
SchedulerConfig, SchedulerConfig,
VllmConfig, VllmConfig,
) )
from vllm.distributed.parallel_state import graph_capture
from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.forward_context import BatchDescriptor, set_forward_context
from vllm.platforms import current_platform from vllm.platforms import current_platform
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
@ -167,11 +168,14 @@ class TestCUDAGraphWrapper:
batch_descriptor = BatchDescriptor(num_tokens=10) batch_descriptor = BatchDescriptor(num_tokens=10)
# 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,
),
graph_capture(device=torch.device("cuda")) as graph_ctx,
): ):
wrapper(self.input_tensor) wrapper(self.input_tensor)
@ -184,6 +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,
): ):
output1 = wrapper(self.input_tensor) output1 = wrapper(self.input_tensor)
# capturing phase should generate a zero output # capturing phase should generate a zero output