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, 08:15 UTC
[compile] Add compiled artifact counter for VLLM_USE_MEGA_AOT_ARTIFACT=1. (#37589)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Mar 20, 2026 at 12:22 UTC
2e089b96a8a69d921b3d3a127c0c9f84caca6f5e
2 files changed
+17
-2
tests/compile/test_startup.py
+14
-2
@@ -9,11 +9,15 @@ then runs in the parent with clean in-memory state but populated caches.
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
19
+from ..utils import fork_new_process_for_each_test
20
+
21
MODEL = "microsoft/Phi-tiny-MoE-instruct"
22
23
@@ -45,8 +49,11 @@ def _cold_start(vllm_runner):
49
assert counters["aot_autograd"]["autograd_cache_hit"] == 0
50
51
48
-def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache):
52
+@fork_new_process_for_each_test
53
+@pytest.mark.parametrize("mega_aot_artifact", ["0", "1"])
54
+def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache, mega_aot_artifact):
55
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
56
+ monkeypatch.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", mega_aot_artifact)
57
58
# Cold start in a forked child (must fork before CUDA init).
59
# This model has 32 identical transformer layers which produce
@@ -64,7 +71,12 @@ def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache):
71
num_compiled_artifacts_saved=0,
72
):
73
_run_vllm(vllm_runner)
67
- assert counters["aot_autograd"]["total"] == 30
74
+ if envs.VLLM_USE_MEGA_AOT_ARTIFACT:
75
+ # MEGA_AOT_ARTIFACT is enabled, so we expect no aot_autograd running on
76
+ # subgraphs.
77
+ assert counters["aot_autograd"]["total"] == 0
78
+ else:
79
+ assert counters["aot_autograd"]["total"] == 30
80
assert counters["aot_autograd"]["autograd_cache_miss"] == 0
81
assert (
82
counters["aot_autograd"]["autograd_cache_hit"] == 0
vllm/compilation/caching.py
+3
@@ -17,6 +17,7 @@ from torch.utils import _pytree as pytree
17
18
import vllm.envs as envs
19
from vllm.compilation.compiler_interface import get_inductor_factors
20
+from vllm.compilation.counter import compilation_counter
21
from vllm.config import VllmConfig, get_current_vllm_config
22
from vllm.config.utils import hash_factors
23
from vllm.logger import init_logger
@@ -61,6 +62,7 @@ class StandaloneCompiledArtifacts:
62
self.submodule_bytes[f"{submod_name}_{shape}"] = hex_digest
63
if hex_digest not in self.submodule_bytes_store:
64
self.submodule_bytes_store[hex_digest] = entry
65
+ compilation_counter.num_compiled_artifacts_saved += 1
66
logger.debug(
67
"inserting new artifact for submod %s with shape %s "
68
"(%s bytes) at hash %s",
@@ -124,6 +126,7 @@ class StandaloneCompiledArtifacts:
126
127
def _load_entry(entry_bytes: bytes) -> AOTCompiledArtifact:
128
entry = pickle.loads(entry_bytes)
129
+ compilation_counter.num_compiled_artifacts_loaded += 1
130
return AOTCompiledArtifact.deserialize(entry)
131
132
with concurrent.futures.ThreadPoolExecutor() as executor: