@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 247 lines 8.03 KB
Raw
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 separate child process 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 create_new_process_for_each_test, requires_spawn_multiprocessing
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 # Phi-tiny-MoE uses SWA, whose admission cap is `cdiv(L, block_size) + 1`
38 # at default block_size=16 — i.e. 17 blocks for max_model_len=256. Use
39 # 32 for headroom.
40 num_gpu_blocks_override=32,
41 ):
42 pass
43
44
45 def _cold_start(vllm_runner):
46 counters.clear()
47 with compilation_counter.expect(
48 num_compiled_artifacts_saved=3,
49 num_compiled_artifacts_loaded=0,
50 ):
51 _run_vllm(vllm_runner)
52 assert counters["aot_autograd"]["total"] == 33
53 assert counters["aot_autograd"]["autograd_cache_miss"] == 3
54 assert counters["aot_autograd"]["autograd_cache_hit"] == 0
55
56
57 @create_new_process_for_each_test()
58 @pytest.mark.parametrize("mega_aot_artifact", ["0", "1"])
59 def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache, mega_aot_artifact):
60 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
61 monkeypatch.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", mega_aot_artifact)
62 monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
63
64 # ROCm and XPU cannot fork safely after device discovery.
65 # This model has 32 identical transformer layers which produce
66 # 33 subgraphs after splitting on attention — only 3 are unique.
67 ctx = (
68 mp.get_context("spawn")
69 if requires_spawn_multiprocessing()
70 else mp.get_context("fork")
71 )
72 p = ctx.Process(target=_cold_start, args=(vllm_runner,))
73 p.start()
74 p.join()
75 assert p.exitcode == 0, "Cold-start child failed"
76
77 # Warm start — compiled artifacts loaded from disk cache.
78 counters.clear()
79 with compilation_counter.expect(
80 num_compiled_artifacts_loaded=3,
81 num_compiled_artifacts_saved=0,
82 ):
83 _run_vllm(vllm_runner)
84 mega_aot_active = envs.VLLM_USE_MEGA_AOT_ARTIFACT and is_torch_equal_or_newer(
85 "2.10.0"
86 )
87 if mega_aot_active:
88 # MEGA_AOT_ARTIFACT is enabled, so we expect no aot_autograd running on
89 # subgraphs.
90 assert counters["aot_autograd"]["total"] == 0
91 else:
92 assert counters["aot_autograd"]["total"] == 30
93 assert counters["aot_autograd"]["autograd_cache_miss"] == 0
94 assert (
95 counters["aot_autograd"]["autograd_cache_hit"] == 0
96 ) # No miss at aot_autograd level causing disk I/O.
97
98
99 # ---------------------------------------------------------------------------
100 # Parametrized model startup tests
101 # ---------------------------------------------------------------------------
102
103
104 class ModelStartupSpec(NamedTuple):
105 model: str
106 hf_overrides: dict
107 cold_artifacts_saved: int
108 warm_artifacts_saved: int
109 warm_artifacts_loaded: int
110
111
112 _SMALL_MOE_OVERRIDES = {
113 "num_hidden_layers": 8,
114 "hidden_size": 256,
115 "intermediate_size": 512,
116 "num_attention_heads": 8,
117 "num_key_value_heads": 1,
118 "n_routed_experts": 8,
119 }
120
121 MODEL_SPECS = [
122 pytest.param(
123 ModelStartupSpec(
124 model="openai/gpt-oss-120b",
125 hf_overrides={
126 "num_hidden_layers": 8,
127 "hidden_size": 256,
128 "intermediate_size": 512,
129 "num_attention_heads": 8,
130 "num_key_value_heads": 1,
131 "num_local_experts": 8,
132 },
133 cold_artifacts_saved=3,
134 warm_artifacts_saved=0,
135 warm_artifacts_loaded=3,
136 ),
137 id="gpt_oss_120b",
138 ),
139 pytest.param(
140 ModelStartupSpec(
141 model="moonshotai/Kimi-K2.5",
142 hf_overrides={"text_config": _SMALL_MOE_OVERRIDES},
143 cold_artifacts_saved=4,
144 # https://github.com/vllm-project/vllm/issues/38051
145 warm_artifacts_saved=0 if is_torch_equal_or_newer("2.12.0") else 4,
146 warm_artifacts_loaded=4 if is_torch_equal_or_newer("2.12.0") else 0,
147 ),
148 id="kimi_k2.5",
149 ),
150 pytest.param(
151 ModelStartupSpec(
152 model="zai-org/GLM-4.5",
153 hf_overrides=_SMALL_MOE_OVERRIDES,
154 cold_artifacts_saved=4,
155 warm_artifacts_saved=0,
156 warm_artifacts_loaded=4,
157 ),
158 id="glm_4.5",
159 ),
160 pytest.param(
161 ModelStartupSpec(
162 model="MiniMaxAI/MiniMax-M2.5",
163 hf_overrides=_SMALL_MOE_OVERRIDES,
164 cold_artifacts_saved=3,
165 warm_artifacts_saved=0,
166 warm_artifacts_loaded=3,
167 ),
168 id="minimax_m2.5",
169 ),
170 ]
171
172
173 def _run_model(vllm_runner, spec: ModelStartupSpec):
174 with vllm_runner(
175 spec.model,
176 trust_remote_code=True,
177 max_model_len=256,
178 max_num_batched_tokens=1024,
179 block_size=64,
180 load_format="dummy",
181 hf_overrides=spec.hf_overrides,
182 compilation_config=CompilationConfig(
183 mode=CompilationMode.VLLM_COMPILE,
184 cudagraph_mode=CUDAGraphMode.NONE,
185 pass_config=PassConfig(fuse_allreduce_rms=False),
186 ),
187 num_gpu_blocks_override=16,
188 ):
189 pass
190
191
192 def _check_model_run(vllm_runner, spec: ModelStartupSpec, is_cold_start: bool):
193 """Runs a model and checks the number of compiled artifacts."""
194 old = compilation_counter.clone()
195 _run_model(vllm_runner, spec)
196 saved = (
197 compilation_counter.num_compiled_artifacts_saved
198 - old.num_compiled_artifacts_saved
199 )
200 loaded = (
201 compilation_counter.num_compiled_artifacts_loaded
202 - old.num_compiled_artifacts_loaded
203 )
204
205 start_type = "COLD" if is_cold_start else "WARM"
206 # Print actual values for debugging — intentional, helps diagnose
207 # failures and calibrate expected counts when adding new models.
208 print(f"\n=== {start_type} START for {spec.model} ===")
209 print(f" num_compiled_artifacts_saved={saved}")
210 print(f" num_compiled_artifacts_loaded={loaded}")
211
212 if is_cold_start:
213 expected_saved = spec.cold_artifacts_saved
214 expected_loaded = 0
215 else:
216 expected_saved = spec.warm_artifacts_saved
217 expected_loaded = spec.warm_artifacts_loaded
218
219 assert saved == expected_saved, f"{start_type.lower()}_artifacts_saved: got {saved}"
220 assert loaded == expected_loaded, (
221 f"{start_type.lower()}_artifacts_loaded: got {loaded}"
222 )
223
224
225 def _cold_start_model(vllm_runner, spec: ModelStartupSpec):
226 _check_model_run(vllm_runner, spec, is_cold_start=True)
227
228
229 @pytest.mark.parametrize("spec", MODEL_SPECS)
230 @create_new_process_for_each_test()
231 def test_model_startup(monkeypatch, vllm_runner, fresh_vllm_cache, spec):
232 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
233 monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
234
235 # ROCm and XPU cannot fork safely after device discovery.
236 ctx = (
237 mp.get_context("spawn")
238 if requires_spawn_multiprocessing()
239 else mp.get_context("fork")
240 )
241 p = ctx.Process(target=_cold_start_model, args=(vllm_runner, spec))
242 p.start()
243 p.join()
244 assert p.exitcode == 0, "Cold-start child failed"
245
246 # Warm start — compiled artifacts loaded from disk cache.
247 _check_model_run(vllm_runner, spec, is_cold_start=False)