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 6, 00:45 UTC
[compile] Add some more startup tests for top models (#38046)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Mar 25, 2026 at 12:02 UTC
6e37c46b35e2ee799fb280180f4d582219bea3f0
4 files changed
+260
-88
.buildkite/test_areas/pytorch.yaml
+11
-1
@@ -17,6 +17,16 @@ steps:
17
# (using -0 for proper path handling)
18
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
19
20
+- label: PyTorch Compilation Unit Tests (H100)
21
+ timeout_in_minutes: 30
22
+ device: h100
23
+ num_devices: 1
24
+ source_file_dependencies:
25
+ - vllm/
26
+ - tests/compile/h100/
27
+ commands:
28
+ - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
29
+
30
- label: PyTorch Compilation Passes Unit Tests
31
timeout_in_minutes: 20
32
source_file_dependencies:
@@ -54,4 +64,4 @@ steps:
64
source_file_dependencies:
65
- requirements/nightly_torch_test.txt
66
commands:
57
- - bash standalone_tests/pytorch_nightly_dependency.sh
\ No newline at end of file
67
+ - bash standalone_tests/pytorch_nightly_dependency.sh
tests/compile/h100/__init__.py
tests/compile/h100/test_startup.py
new
+249
@@ -0,0 +1,249 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Cold start and warm start tests for vLLM-compile.
4
+
5
+Cold start runs in a forked child (must fork before CUDA init) which
6
+populates on-disk caches and asserts cold-start counters. Warm start
7
+then runs in the parent with clean in-memory state but populated caches.
8
+"""
9
+
10
+import multiprocessing as mp
11
+from typing import NamedTuple
12
+
13
+import pytest
14
+from torch._dynamo.utils import counters
15
+
16
+import vllm.envs as envs
17
+from vllm.compilation.counter import compilation_counter
18
+from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
19
+from vllm.utils.torch_utils import is_torch_equal_or_newer
20
+
21
+from ...utils import fork_new_process_for_each_test
22
+
23
+MODEL = "microsoft/Phi-tiny-MoE-instruct"
24
+
25
+
26
+def _run_vllm(vllm_runner):
27
+ with vllm_runner(
28
+ MODEL,
29
+ trust_remote_code=False,
30
+ max_model_len=256,
31
+ max_num_batched_tokens=1024,
32
+ load_format="dummy",
33
+ compilation_config=CompilationConfig(
34
+ mode=CompilationMode.VLLM_COMPILE,
35
+ cudagraph_mode=CUDAGraphMode.NONE,
36
+ ),
37
+ num_gpu_blocks_override=8,
38
+ ):
39
+ pass
40
+
41
+
42
+def _cold_start(vllm_runner):
43
+ counters.clear()
44
+ with compilation_counter.expect(
45
+ num_compiled_artifacts_saved=3,
46
+ num_compiled_artifacts_loaded=0,
47
+ ):
48
+ _run_vllm(vllm_runner)
49
+ assert counters["aot_autograd"]["total"] == 33
50
+ assert counters["aot_autograd"]["autograd_cache_miss"] == 3
51
+ assert counters["aot_autograd"]["autograd_cache_hit"] == 0
52
+
53
+
54
+@fork_new_process_for_each_test
55
+@pytest.mark.parametrize("mega_aot_artifact", ["0", "1"])
56
+def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache, mega_aot_artifact):
57
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
58
+ monkeypatch.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", mega_aot_artifact)
59
+
60
+ # Cold start in a forked child (must fork before CUDA init).
61
+ # This model has 32 identical transformer layers which produce
62
+ # 33 subgraphs after splitting on attention — only 3 are unique.
63
+ ctx = mp.get_context("fork")
64
+ p = ctx.Process(target=_cold_start, args=(vllm_runner,))
65
+ p.start()
66
+ p.join()
67
+ assert p.exitcode == 0, "Cold-start child failed"
68
+
69
+ # Warm start — compiled artifacts loaded from disk cache.
70
+ counters.clear()
71
+ with compilation_counter.expect(
72
+ num_compiled_artifacts_loaded=3,
73
+ num_compiled_artifacts_saved=0,
74
+ ):
75
+ _run_vllm(vllm_runner)
76
+ mega_aot_active = envs.VLLM_USE_MEGA_AOT_ARTIFACT and is_torch_equal_or_newer(
77
+ "2.10.0"
78
+ )
79
+ if mega_aot_active:
80
+ # MEGA_AOT_ARTIFACT is enabled, so we expect no aot_autograd running on
81
+ # subgraphs.
82
+ assert counters["aot_autograd"]["total"] == 0
83
+ else:
84
+ assert counters["aot_autograd"]["total"] == 30
85
+ assert counters["aot_autograd"]["autograd_cache_miss"] == 0
86
+ assert (
87
+ counters["aot_autograd"]["autograd_cache_hit"] == 0
88
+ ) # No miss at aot_autograd level causing disk I/O.
89
+
90
+
91
+# ---------------------------------------------------------------------------
92
+# Parametrized model startup tests
93
+# ---------------------------------------------------------------------------
94
+
95
+
96
+class ModelStartupSpec(NamedTuple):
97
+ model: str
98
+ hf_overrides: dict
99
+ cold_artifacts_saved: int
100
+ warm_artifacts_saved: int
101
+ warm_artifacts_loaded: int
102
+
103
+
104
+_SMALL_MOE_OVERRIDES = {
105
+ "num_hidden_layers": 8,
106
+ "hidden_size": 256,
107
+ "intermediate_size": 512,
108
+ "num_attention_heads": 8,
109
+ "num_key_value_heads": 1,
110
+ "n_routed_experts": 8,
111
+}
112
+
113
+MODEL_SPECS = [
114
+ pytest.param(
115
+ ModelStartupSpec(
116
+ model="openai/gpt-oss-120b",
117
+ hf_overrides={
118
+ "num_hidden_layers": 8,
119
+ "hidden_size": 256,
120
+ "intermediate_size": 512,
121
+ "num_attention_heads": 8,
122
+ "num_key_value_heads": 1,
123
+ "num_local_experts": 8,
124
+ },
125
+ cold_artifacts_saved=3,
126
+ warm_artifacts_saved=0,
127
+ warm_artifacts_loaded=3,
128
+ ),
129
+ id="gpt_oss_120b",
130
+ ),
131
+ # NOTE: DeepSeek-V3.2 requires sparse MLA (index_topk) which needs
132
+ # Hopper+ GPUs. This test must run on H100 (see pytorch.yaml).
133
+ pytest.param(
134
+ ModelStartupSpec(
135
+ model="deepseek-ai/DeepSeek-V3.2",
136
+ hf_overrides=_SMALL_MOE_OVERRIDES,
137
+ cold_artifacts_saved=4,
138
+ # TODO: https://github.com/vllm-project/vllm/issues/38051
139
+ # We shouldn't be saving any artifacts on warm start.
140
+ warm_artifacts_saved=4,
141
+ warm_artifacts_loaded=0,
142
+ ),
143
+ id="deepseek_v3.2",
144
+ ),
145
+ pytest.param(
146
+ ModelStartupSpec(
147
+ model="moonshotai/Kimi-K2.5",
148
+ hf_overrides={"text_config": _SMALL_MOE_OVERRIDES},
149
+ cold_artifacts_saved=4,
150
+ # TODO: https://github.com/vllm-project/vllm/issues/38051
151
+ # We shouldn't be saving any artifacts on warm start.
152
+ warm_artifacts_saved=4,
153
+ warm_artifacts_loaded=0,
154
+ ),
155
+ id="kimi_k2.5",
156
+ ),
157
+ pytest.param(
158
+ ModelStartupSpec(
159
+ model="zai-org/GLM-4.5",
160
+ hf_overrides=_SMALL_MOE_OVERRIDES,
161
+ cold_artifacts_saved=4,
162
+ warm_artifacts_saved=0,
163
+ warm_artifacts_loaded=4,
164
+ ),
165
+ id="glm_4.5",
166
+ ),
167
+ pytest.param(
168
+ ModelStartupSpec(
169
+ model="MiniMaxAI/MiniMax-M2.5",
170
+ hf_overrides=_SMALL_MOE_OVERRIDES,
171
+ cold_artifacts_saved=3,
172
+ warm_artifacts_saved=0,
173
+ warm_artifacts_loaded=3,
174
+ ),
175
+ id="minimax_m2.5",
176
+ ),
177
+]
178
+
179
+
180
+def _run_model(vllm_runner, spec: ModelStartupSpec):
181
+ with vllm_runner(
182
+ spec.model,
183
+ trust_remote_code=True,
184
+ max_model_len=256,
185
+ max_num_batched_tokens=1024,
186
+ block_size=64,
187
+ load_format="dummy",
188
+ hf_overrides=spec.hf_overrides,
189
+ compilation_config=CompilationConfig(
190
+ mode=CompilationMode.VLLM_COMPILE,
191
+ cudagraph_mode=CUDAGraphMode.NONE,
192
+ pass_config=PassConfig(fuse_allreduce_rms=False),
193
+ ),
194
+ num_gpu_blocks_override=8,
195
+ ):
196
+ pass
197
+
198
+
199
+def _check_model_run(vllm_runner, spec: ModelStartupSpec, is_cold_start: bool):
200
+ """Runs a model and checks the number of compiled artifacts."""
201
+ old = compilation_counter.clone()
202
+ _run_model(vllm_runner, spec)
203
+ saved = (
204
+ compilation_counter.num_compiled_artifacts_saved
205
+ - old.num_compiled_artifacts_saved
206
+ )
207
+ loaded = (
208
+ compilation_counter.num_compiled_artifacts_loaded
209
+ - old.num_compiled_artifacts_loaded
210
+ )
211
+
212
+ start_type = "COLD" if is_cold_start else "WARM"
213
+ # Print actual values for debugging — intentional, helps diagnose
214
+ # failures and calibrate expected counts when adding new models.
215
+ print(f"\n=== {start_type} START for {spec.model} ===")
216
+ print(f" num_compiled_artifacts_saved={saved}")
217
+ print(f" num_compiled_artifacts_loaded={loaded}")
218
+
219
+ if is_cold_start:
220
+ expected_saved = spec.cold_artifacts_saved
221
+ expected_loaded = 0
222
+ else:
223
+ expected_saved = spec.warm_artifacts_saved
224
+ expected_loaded = spec.warm_artifacts_loaded
225
+
226
+ assert saved == expected_saved, f"{start_type.lower()}_artifacts_saved: got {saved}"
227
+ assert loaded == expected_loaded, (
228
+ f"{start_type.lower()}_artifacts_loaded: got {loaded}"
229
+ )
230
+
231
+
232
+def _cold_start_model(vllm_runner, spec: ModelStartupSpec):
233
+ _check_model_run(vllm_runner, spec, is_cold_start=True)
234
+
235
+
236
+@pytest.mark.parametrize("spec", MODEL_SPECS)
237
+@fork_new_process_for_each_test
238
+def test_model_startup(monkeypatch, vllm_runner, fresh_vllm_cache, spec):
239
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
240
+
241
+ # Cold start in a forked child (must fork before CUDA init).
242
+ ctx = mp.get_context("fork")
243
+ p = ctx.Process(target=_cold_start_model, args=(vllm_runner, spec))
244
+ p.start()
245
+ p.join()
246
+ assert p.exitcode == 0, "Cold-start child failed"
247
+
248
+ # Warm start — compiled artifacts loaded from disk cache.
249
+ _check_model_run(vllm_runner, spec, is_cold_start=False)
tests/compile/test_startup.py
deleted
-87
@@ -1,87 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Cold start and warm start tests for vLLM-compile.
4
-
5
-Cold start runs in a forked child (must fork before CUDA init) which
6
-populates on-disk caches and asserts cold-start counters. Warm start
7
-then runs in the parent with clean in-memory state but populated caches.
8
-"""
9
-
10
-import multiprocessing as mp
11
-
12
-import pytest
13
-from torch._dynamo.utils import counters
14
-
15
-import vllm.envs as envs
16
-from vllm.compilation.counter import compilation_counter
17
-from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
18
-from vllm.utils.torch_utils import is_torch_equal_or_newer
19
-
20
-from ..utils import fork_new_process_for_each_test
21
-
22
-MODEL = "microsoft/Phi-tiny-MoE-instruct"
23
-
24
-
25
-def _run_vllm(vllm_runner):
26
- with vllm_runner(
27
- MODEL,
28
- trust_remote_code=False,
29
- max_model_len=256,
30
- max_num_batched_tokens=1024,
31
- load_format="dummy",
32
- compilation_config=CompilationConfig(
33
- mode=CompilationMode.VLLM_COMPILE,
34
- cudagraph_mode=CUDAGraphMode.NONE,
35
- ),
36
- num_gpu_blocks_override=8,
37
- ):
38
- pass
39
-
40
-
41
-def _cold_start(vllm_runner):
42
- counters.clear()
43
- with compilation_counter.expect(
44
- num_compiled_artifacts_saved=3,
45
- num_compiled_artifacts_loaded=0,
46
- ):
47
- _run_vllm(vllm_runner)
48
- assert counters["aot_autograd"]["total"] == 33
49
- assert counters["aot_autograd"]["autograd_cache_miss"] == 3
50
- assert counters["aot_autograd"]["autograd_cache_hit"] == 0
51
-
52
-
53
-@fork_new_process_for_each_test
54
-@pytest.mark.parametrize("mega_aot_artifact", ["0", "1"])
55
-def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache, mega_aot_artifact):
56
- monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
57
- monkeypatch.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", mega_aot_artifact)
58
-
59
- # Cold start in a forked child (must fork before CUDA init).
60
- # This model has 32 identical transformer layers which produce
61
- # 33 subgraphs after splitting on attention — only 3 are unique.
62
- ctx = mp.get_context("fork")
63
- p = ctx.Process(target=_cold_start, args=(vllm_runner,))
64
- p.start()
65
- p.join()
66
- assert p.exitcode == 0, "Cold-start child failed"
67
-
68
- # Warm start — compiled artifacts loaded from disk cache.
69
- counters.clear()
70
- with compilation_counter.expect(
71
- num_compiled_artifacts_loaded=3,
72
- num_compiled_artifacts_saved=0,
73
- ):
74
- _run_vllm(vllm_runner)
75
- mega_aot_active = envs.VLLM_USE_MEGA_AOT_ARTIFACT and is_torch_equal_or_newer(
76
- "2.10.0"
77
- )
78
- if mega_aot_active:
79
- # MEGA_AOT_ARTIFACT is enabled, so we expect no aot_autograd running on
80
- # subgraphs.
81
- assert counters["aot_autograd"]["total"] == 0
82
- else:
83
- assert counters["aot_autograd"]["total"] == 30
84
- assert counters["aot_autograd"]["autograd_cache_miss"] == 0
85
- assert (
86
- counters["aot_autograd"]["autograd_cache_hit"] == 0
87
- ) # No miss at aot_autograd level causing disk I/O.