A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[CI] Add tests for cudagraph (#27391)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Oct 25, 2025 at 10:37 UTC
29c9cb8007bfc865ec9f8d02a5630eab3d7b6577
4 files changed
+54
-18
.buildkite/test-pipeline.yaml
+12
@@ -435,6 +435,18 @@ steps:
435
- pytest -v -s compile/test_full_graph.py
436
- pytest -v -s compile/test_fusions_e2e.py
437
438
+- label: Cudagraph test
439
+ timeout_in_minutes: 20
440
+ mirror_hardwares: [amdexperimental]
441
+ source_file_dependencies:
442
+ - tests/v1/cudagraph
443
+ - vllm/v1/cudagraph_dispatcher.py
444
+ - vllm/config/compilation.py
445
+ - vllm/compilation
446
+ commands:
447
+ - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
448
+ - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
449
+
450
- label: Kernels Core Operation Test # 48min
451
timeout_in_minutes: 75
452
mirror_hardwares: [amdexperimental]
tests/utils.py
+5
@@ -1111,6 +1111,11 @@ def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]
1111
# `cloudpickle` allows pickling complex functions directly
1112
input_bytes = cloudpickle.dumps((f, output_filepath))
1113
1114
+ repo_root = str(VLLM_PATH.resolve())
1115
+
1116
+ env = dict(env or os.environ)
1117
+ env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")
1118
+
1119
cmd = [sys.executable, "-m", f"{module_name}"]
1120
1121
returned = subprocess.run(
tests/v1/cudagraph/test_cudagraph_dispatch.py
+34
-15
@@ -34,13 +34,16 @@ class SimpleMLP(nn.Module):
34
35
36
def _create_vllm_config(
37
- compilation_config: CompilationConfig, max_num_seqs: int = 8
37
+ compilation_config: CompilationConfig,
38
+ max_num_seqs: int = 8,
39
+ lora_config: bool = False,
40
) -> MagicMock:
41
mock_config = MagicMock(spec=VllmConfig)
42
mock_config.compilation_config = compilation_config
43
mock_config.scheduler_config = SchedulerConfig(max_num_seqs=max_num_seqs)
44
mock_config.parallel_config = ParallelConfig()
43
-
45
+ if not lora_config:
46
+ mock_config.lora_config = None
47
# Mimic the behavior of VllmConfig.__post_init__()
48
if compilation_config.mode == CompilationMode.VLLM_COMPILE:
49
compilation_config.set_splitting_ops_for_v1()
@@ -50,19 +53,21 @@ def _create_vllm_config(
53
54
class TestCudagraphDispatcher:
55
@pytest.mark.parametrize(
53
- "case_id,cudagraph_mode_str,compilation_mode",
56
+ "cudagraph_mode_str,compilation_mode,lora_config",
57
[
58
# Test case 0: Full CG for mixed batches, no separate routine
56
- (0, "FULL", CompilationMode.NONE),
59
+ ("FULL", CompilationMode.NONE, False),
60
# Test case 1: Full CG for uniform batches, piecewise for mixed
58
- (1, "FULL_AND_PIECEWISE", CompilationMode.NONE),
61
+ ("FULL_AND_PIECEWISE", CompilationMode.NONE, False),
62
# Test case 2: Full CG for uniform batches, no CG for mixed
60
- (2, "FULL_DECODE_ONLY", CompilationMode.NONE),
63
+ ("FULL_DECODE_ONLY", CompilationMode.NONE, False),
64
# Test case 3: PIECEWISE for all
62
- (3, "PIECEWISE", CompilationMode.VLLM_COMPILE),
65
+ ("PIECEWISE", CompilationMode.VLLM_COMPILE, False),
66
+ # Test case 4: PIECEWISE for all, specialize LoRA cases
67
+ ("PIECEWISE", CompilationMode.VLLM_COMPILE, True),
68
],
69
)
65
- def test_dispatcher(self, cudagraph_mode_str, compilation_mode):
70
+ def test_dispatcher(self, cudagraph_mode_str, compilation_mode, lora_config):
71
# Setup dispatcher
72
comp_config = CompilationConfig(
73
cudagraph_mode=cudagraph_mode_str,
@@ -70,7 +75,17 @@ class TestCudagraphDispatcher:
75
cudagraph_capture_sizes=[1, 8],
76
)
77
73
- config = _create_vllm_config(comp_config, max_num_seqs=8)
78
+ config = _create_vllm_config(
79
+ comp_config, max_num_seqs=8, lora_config=lora_config
80
+ )
81
+ if (
82
+ cudagraph_mode_str == "FULL_AND_PIECEWISE"
83
+ and compilation_mode == CompilationMode.NONE
84
+ ):
85
+ with pytest.raises(AssertionError):
86
+ dispatcher = CudagraphDispatcher(config)
87
+ return
88
+
89
dispatcher = CudagraphDispatcher(config)
90
dispatcher.initialize_cudagraph_keys(
91
cudagraph_mode=comp_config.cudagraph_mode, uniform_decode_query_len=1
@@ -78,17 +93,24 @@ class TestCudagraphDispatcher:
93
94
# Verify the key is initialized correctly
95
if cudagraph_mode_str in ["FULL_AND_PIECEWISE", "PIECEWISE"]:
81
- assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == 2
96
+ assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == (
97
+ 4 if lora_config else 2
98
+ )
99
else:
100
assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == 0
101
if cudagraph_mode_str not in ["NONE", "PIECEWISE"]:
85
- assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == 2
102
+ assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == (
103
+ 4 if lora_config else 2
104
+ )
105
else:
106
assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == 0
107
108
# Test dispatch logic
109
# 1. non-uniform batch, size in cudagraph size list
91
- desc_full_exact = BatchDescriptor(num_tokens=8, uniform_decode=False)
110
+ desc_full_exact = BatchDescriptor(
111
+ num_tokens=8,
112
+ uniform_decode=False,
113
+ )
114
rt_mode, key = dispatcher.dispatch(desc_full_exact)
115
if cudagraph_mode_str == "FULL":
116
assert rt_mode == CUDAGraphMode.FULL
@@ -138,7 +160,6 @@ class TestCUDAGraphWrapper:
160
self.persistent_input_buffer = torch.zeros(1, 10, device="cuda")
161
self.input_tensor = torch.randn(1, 10, device="cuda")
162
141
- @create_new_process_for_each_test("spawn")
163
def test_capture_and_replay(self):
164
wrapper = CUDAGraphWrapper(
165
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
@@ -192,7 +213,6 @@ class TestCUDAGraphWrapper:
213
eager_output = self.model(self.input_tensor)
214
torch.testing.assert_close(eager_output, output2)
215
195
- @create_new_process_for_each_test("spawn")
216
def test_bypass_on_mode_mismatch(self):
217
wrapper = CUDAGraphWrapper(
218
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
@@ -216,7 +236,6 @@ class TestCUDAGraphWrapper:
236
mock_forward.assert_called_once()
237
assert not wrapper.concrete_cudagraph_entries
238
219
- @create_new_process_for_each_test("spawn")
239
def test_bypass_on_mode_none(self):
240
wrapper = CUDAGraphWrapper(
241
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
tests/v1/cudagraph/test_cudagraph_mode.py
+3
-3
@@ -109,9 +109,9 @@ combo_cases_2 = [
109
@pytest.mark.parametrize(
110
"backend_name,cudagraph_mode,compilation_mode,supported", combo_cases_2
111
)
112
-def test_cudagraph_compilation_combo(combo_case):
113
- backend_name, cudagraph_mode, compilation_mode, supported = combo_case
114
-
112
+def test_cudagraph_compilation_combo(
113
+ backend_name, cudagraph_mode, compilation_mode, supported
114
+):
115
env_vars = backend_configs[backend_name].env_vars
116
117
with temporary_environ(env_vars), ExitStack() as stack: