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 11, 06:15 UTC
[compile] Fix CI for test_gpt2_cache_hit (#30902)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Dec 17, 2025 at 23:22 UTC
5f2f3fba1d9ed0aa433171b86c415a5f02055035
2 files changed
+15
-6
tests/compile/test_aot_compile.py
+10
-1
@@ -9,6 +9,7 @@ from contextlib import contextmanager
9
import pytest
10
import torch
11
12
+import vllm.model_executor.layers.activation
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import (
15
CompilationConfig,
@@ -16,9 +17,12 @@ from vllm.config import (
17
VllmConfig,
18
set_current_vllm_config,
19
)
20
+from vllm.envs import disable_envs_cache
21
from vllm.forward_context import set_forward_context
22
from vllm.utils.torch_utils import is_torch_equal_or_newer
23
24
+from ..utils import create_new_process_for_each_test
25
+
26
27
def reference_fn(x: torch.Tensor):
28
assert x.shape[0] <= 42
@@ -66,6 +70,7 @@ def test_no_dynamo_cache_entry(monkeypatch: pytest.MonkeyPatch):
70
torch.compiler.set_stance("fail_on_recompile"),
71
):
72
CompiledMod(vllm_config=vllm_config)(*args)
73
+ disable_envs_cache()
74
75
m.setenv("VLLM_USE_AOT_COMPILE", "1")
76
torch._dynamo.reset()
@@ -101,6 +106,7 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
106
vllm_config = make_vllm_config()
107
with use_vllm_config(vllm_config):
108
expected = CompiledMod(vllm_config=vllm_config)(*args)
109
+ disable_envs_cache()
110
111
m.setenv("VLLM_FORCE_AOT_LOAD", "1")
112
vllm_config = make_vllm_config()
@@ -130,6 +136,7 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
136
artifacts = compiled_mod.aot_compiled_fn._artifacts
137
guards_string = artifacts.compiled_fn.shape_env.format_guards()
138
assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
139
+ disable_envs_cache()
140
141
m.setenv("VLLM_FORCE_AOT_LOAD", "1")
142
vllm_config = make_vllm_config()
@@ -144,7 +151,7 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
151
@pytest.mark.skipif(
152
not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
153
)
147
-@use_vllm_config(make_vllm_config())
154
+@create_new_process_for_each_test("spawn")
155
def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
156
"""
157
Test that compiling gpt2 twice results in a cache hit and
@@ -186,6 +193,8 @@ def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
193
194
# Clean up first model
195
del llm_model
196
+ disable_envs_cache()
197
+ vllm.model_executor.layers.activation._ACTIVATION_REGISTRY._dict.clear()
198
199
# Second compilation - should hit cache
200
m.setenv("VLLM_FORCE_AOT_LOAD", "1")
vllm/config/compilation.py
+5
-5
@@ -437,14 +437,14 @@ class CompilationConfig:
437
438
compile_ranges_split_points: list[int] | None = None
439
"""Split points that represent compile ranges for inductor.
440
- The compile ranges are
441
- [1, split_points[0]],
442
- [split_points[0] + 1, split_points[1]], ...,
440
+ The compile ranges are
441
+ [1, split_points[0]],
442
+ [split_points[0] + 1, split_points[1]], ...,
443
[split_points[-1] + 1, max_num_batched_tokens].
444
Compile sizes are also used single element ranges,
445
the range is represented as [compile_sizes[i], compile_sizes[i]].
446
-
447
- If a range overlaps with the compile size, graph for compile size
446
+
447
+ If a range overlaps with the compile size, graph for compile size
448
will be prioritized, i.e. if we have a range [1, 8] and a compile size 4,
449
graph for compile size 4 will be compiled and used instead of the graph
450
for range [1, 8].