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 7, 11:30 UTC
[torch.compile] Improve cold and warm start compile tests (#35709)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Mar 2, 2026 at 14:27 UTC
d1a6e96d9e0b76cc2a0af33e014b4bd8b860f1e4
5 files changed
+82
-48
tests/compile/test_cold_start.py
deleted
-48
@@ -1,48 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from torch._dynamo.utils import counters
5
-
6
-from vllm import LLM
7
-from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
8
-
9
-
10
-def test_moe_compilation_cold_start(monkeypatch, use_fresh_inductor_cache):
11
- # Run in same process so we can access PyTorch's internal counters
12
- monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
13
-
14
- # I'm not sure if this is going to affect the numbers
15
- monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
16
-
17
- # Force cold compilation
18
- monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
19
-
20
- compilation_config = CompilationConfig(
21
- mode=CompilationMode.VLLM_COMPILE,
22
- cudagraph_mode=CUDAGraphMode.NONE, # make the model loading faster
23
- )
24
-
25
- counters.clear()
26
-
27
- _ = LLM(
28
- model="microsoft/Phi-tiny-MoE-instruct",
29
- max_model_len=256,
30
- load_format="dummy", # make the model loading faster
31
- compilation_config=compilation_config,
32
- num_gpu_blocks_override=8, # make the model loading faster
33
- )
34
-
35
- # vLLM-compile cold start is special. By default, we do
36
- # one full dynamo capture of the entire forward pass.
37
- # The forward pass consists of 32 transformer layers.
38
- # Then, we split on the attention operation. This results in
39
- # 33 subgraphs (not including the attention operation).
40
- # We then generate compiled artifacts for the unique subgraphs.
41
- #
42
- # There are actually only 3 unique subgraphs for this model
43
- # (all of its transformer layers are the same modulo weights);
44
- # this is true for most vLLM models.
45
- # So we test that during cold start, we are only compling
46
- # for 3 unique subgraphs.
47
- assert counters["aot_autograd"]["autograd_cache_miss"] == 3
48
- assert counters["aot_autograd"]["autograd_cache_hit"] == 0
tests/compile/test_startup.py
new
+71
@@ -0,0 +1,71 @@
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
+from torch._dynamo.utils import counters
13
+
14
+from vllm.compilation.counter import compilation_counter
15
+from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
16
+
17
+MODEL = "microsoft/Phi-tiny-MoE-instruct"
18
+
19
+
20
+def _run_vllm(vllm_runner):
21
+ with vllm_runner(
22
+ MODEL,
23
+ trust_remote_code=False,
24
+ max_model_len=256,
25
+ max_num_batched_tokens=1024,
26
+ load_format="dummy",
27
+ compilation_config=CompilationConfig(
28
+ mode=CompilationMode.VLLM_COMPILE,
29
+ cudagraph_mode=CUDAGraphMode.NONE,
30
+ ),
31
+ num_gpu_blocks_override=8,
32
+ ):
33
+ pass
34
+
35
+
36
+def _cold_start(vllm_runner):
37
+ counters.clear()
38
+ with compilation_counter.expect(
39
+ num_compiled_artifacts_saved=3,
40
+ num_compiled_artifacts_loaded=0,
41
+ ):
42
+ _run_vllm(vllm_runner)
43
+ assert counters["aot_autograd"]["total"] == 33
44
+ assert counters["aot_autograd"]["autograd_cache_miss"] == 3
45
+ assert counters["aot_autograd"]["autograd_cache_hit"] == 0
46
+
47
+
48
+def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache):
49
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
50
+
51
+ # Cold start in a forked child (must fork before CUDA init).
52
+ # This model has 32 identical transformer layers which produce
53
+ # 33 subgraphs after splitting on attention — only 3 are unique.
54
+ ctx = mp.get_context("fork")
55
+ p = ctx.Process(target=_cold_start, args=(vllm_runner,))
56
+ p.start()
57
+ p.join()
58
+ assert p.exitcode == 0, "Cold-start child failed"
59
+
60
+ # Warm start — compiled artifacts loaded from disk cache.
61
+ counters.clear()
62
+ with compilation_counter.expect(
63
+ num_compiled_artifacts_loaded=3,
64
+ # TODO: warm start should not save any artifacts
65
+ # https://github.com/vllm-project/vllm/issues/35708
66
+ num_compiled_artifacts_saved=1,
67
+ ):
68
+ _run_vllm(vllm_runner)
69
+ assert counters["aot_autograd"]["total"] == 30
70
+ assert counters["aot_autograd"]["autograd_cache_miss"] == 0
71
+ assert counters["aot_autograd"]["autograd_cache_hit"] == 1
tests/conftest.py
+8
@@ -1548,6 +1548,14 @@ def use_fresh_inductor_cache():
1548
yield
1549
1550
1551
+@pytest.fixture
1552
+def fresh_vllm_cache(monkeypatch, use_fresh_inductor_cache):
1553
+ """Temporary VLLM_CACHE_ROOT combined with a fresh inductor cache."""
1554
+ with tempfile.TemporaryDirectory() as tmp_dir:
1555
+ monkeypatch.setenv("VLLM_CACHE_ROOT", tmp_dir)
1556
+ yield tmp_dir
1557
+
1558
+
1559
@pytest.fixture(scope="function")
1560
def enable_pickle(monkeypatch):
1561
"""`LLM.apply_model` requires pickling a function."""
vllm/compilation/compiler_interface.py
+1
@@ -368,6 +368,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
368
inductor_compiled_graph = torch._inductor.CompiledArtifact.load(
369
path=path, format=self.save_format
370
)
371
+ compilation_counter.num_compiled_artifacts_loaded += 1
372
from torch._inductor.compile_fx import graph_returns_tuple
373
374
returns_tuple = graph_returns_tuple(graph)
vllm/compilation/counter.py
+2
@@ -29,6 +29,8 @@ class CompilationCounter:
29
num_cache_entries_updated: int = 0
30
# The number of standalone_compile compiled artifacts saved
31
num_compiled_artifacts_saved: int = 0
32
+ # The number of standalone_compile compiled artifacts loaded from cache
33
+ num_compiled_artifacts_loaded: int = 0
34
# Number of times a model was loaded with CompilationMode.STOCK_TORCH_COMPILE
35
stock_torch_compile_count: int = 0
36