[CI] Prune `PyTorch Compilation Unit Tests` (#51069)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 4, 2026 at 21:53 UTC
d16e7f9da4c12c9487d4ccd6c97fb30d8bf74376
3 files changed
+124
-172
.buildkite/test_areas/pytorch.yaml
+2
-3
@@ -5,7 +5,7 @@ steps:
5
- label: PyTorch Compilation Unit Tests
6
device: h200_35gb
7
key: pytorch-compilation-unit-tests
8
- timeout_in_minutes: 150
8
+ timeout_in_minutes: 60
9
source_file_dependencies:
10
- vllm/__init__.py
11
- vllm/_aiter_ops.py
@@ -45,8 +45,7 @@ steps:
45
- label: PyTorch Compilation Unit Tests (H100)
46
key: pytorch-compilation-unit-tests-h100
47
timeout_in_minutes: 30
48
- device: h100
49
- num_devices: 1
48
+ device: h200_18gb
49
source_file_dependencies:
50
- vllm/__init__.py
51
- vllm/_aiter_ops.py
tests/compile/test_aot_compile.py
+21
-101
@@ -43,7 +43,7 @@ def vllm_tmp_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
43
def reference_fn(x: torch.Tensor):
44
assert x.shape[0] <= 42
45
assert x.shape[0] % 2 == 0
46
- for _ in range(3000):
46
+ for _ in range(30):
47
x = x + x.shape[0]
48
return x
49
@@ -52,7 +52,7 @@ def reference_fn_tuple(x: torch.Tensor):
52
"""Reference function that returns a tuple of tensors."""
53
assert x.shape[0] <= 42
54
assert x.shape[0] % 2 == 0
55
- for _ in range(3000):
55
+ for _ in range(30):
56
x = x + x.shape[0]
57
return x, x * 2
58
@@ -140,18 +140,35 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
140
m.setenv("VLLM_USE_AOT_COMPILE", "1")
141
m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
142
m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
143
+ disable_envs_cache()
144
vllm_config = make_vllm_config()
144
- with use_vllm_config(vllm_config):
145
+ with (
146
+ use_vllm_config(vllm_config),
147
+ compilation_counter.expect(
148
+ num_aot_compiles=1,
149
+ num_aot_artifacts_saved=1,
150
+ num_aot_artifacts_loaded=0,
151
+ ),
152
+ ):
153
compiled_mod = CompiledMod(vllm_config=vllm_config)
154
expected = compiled_mod(*args)
155
+ assert isinstance(expected, torch.Tensor)
156
157
disable_envs_cache()
158
159
m.setenv("VLLM_FORCE_AOT_LOAD", "1")
160
vllm_config = make_vllm_config()
152
- with use_vllm_config(vllm_config):
161
+ with (
162
+ use_vllm_config(vllm_config),
163
+ compilation_counter.expect(
164
+ num_aot_compiles=0,
165
+ num_aot_artifacts_saved=0,
166
+ num_aot_artifacts_loaded=1,
167
+ ),
168
+ ):
169
cached_mod = CompiledMod(vllm_config=vllm_config)
170
ret = cached_mod(*args)
171
+ assert isinstance(ret, torch.Tensor)
172
assert cached_mod.was_aot_compile_fn_loaded_from_disk, (
173
"Expected was_aot_compile_fn_loaded_from_disk to be True"
174
)
@@ -182,65 +199,6 @@ def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch):
199
assert gm.code == loaded_gm.code
200
201
185
-@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
186
-def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch):
187
- """
188
- Test that cache loading correctly handles the returns_tuple logic.
189
-
190
- This verifies that when a model returns a single tensor (not a tuple),
191
- the output type is consistent between fresh compilation and cache load.
192
- Without the fix, cached artifacts would return [tensor] instead of tensor.
193
- """
194
- with monkeypatch.context() as m:
195
- args = (torch.randn(10, 10),)
196
-
197
- with tempfile.TemporaryDirectory() as tmpdirname:
198
- m.setenv("VLLM_CACHE_ROOT", tmpdirname)
199
- m.setenv("VLLM_USE_AOT_COMPILE", "1")
200
- m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
201
- m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
202
- vllm_config = make_vllm_config()
203
-
204
- # Fresh compilation
205
- with use_vllm_config(vllm_config):
206
- compiled_mod = CompiledMod(vllm_config=vllm_config)
207
- fresh_result = compiled_mod(*args)
208
- fresh_result_type = type(fresh_result)
209
-
210
- # Verify fresh result is a tensor, not a tuple/list
211
- assert isinstance(fresh_result, torch.Tensor), (
212
- f"Fresh compile should return tensor, got {fresh_result_type}"
213
- )
214
-
215
- disable_envs_cache()
216
-
217
- # Load from cache
218
- m.setenv("VLLM_FORCE_AOT_LOAD", "1")
219
- vllm_config = make_vllm_config()
220
- with use_vllm_config(vllm_config):
221
- cached_mod = CompiledMod(vllm_config=vllm_config)
222
- cached_result = cached_mod(*args)
223
- cached_result_type = type(cached_result)
224
-
225
- # Verify cache was actually loaded
226
- assert cached_mod.was_aot_compile_fn_loaded_from_disk, (
227
- "Expected was_aot_compile_fn_loaded_from_disk to be True after "
228
- "loading from cache"
229
- )
230
-
231
- # Verify cached result has same type as fresh result
232
- assert isinstance(cached_result, torch.Tensor), (
233
- f"Cache load should return tensor, got {cached_result_type}. "
234
- "This indicates the returns_tuple logic is not being applied "
235
- "correctly when loading from cache."
236
- )
237
-
238
- # Verify values match
239
- assert torch.allclose(cached_result, fresh_result), (
240
- "Cached result values should match fresh compilation"
241
- )
242
-
243
-
202
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
203
def test_cache_load_returns_tuple_consistency_tuple_output(
204
monkeypatch: pytest.MonkeyPatch,
@@ -872,41 +830,3 @@ def test_disable_compile_cache_skips_aot_load(
830
mod(*args)
831
832
assert not mod.was_aot_compile_fn_loaded_from_disk
875
-
876
-
877
-@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
878
-def test_aot_counters_on_save_and_load(
879
- monkeypatch: pytest.MonkeyPatch, fresh_vllm_cache: str
880
-):
881
- """Verify AOT counters are incremented correctly on save and load."""
882
- monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
883
- disable_envs_cache()
884
-
885
- args = (torch.randn(10, 10),)
886
-
887
- # Phase 1: fresh compile + save
888
- vllm_config = make_vllm_config()
889
- with (
890
- use_vllm_config(vllm_config),
891
- compilation_counter.expect(
892
- num_aot_compiles=1,
893
- num_aot_artifacts_saved=1,
894
- num_aot_artifacts_loaded=0,
895
- ),
896
- ):
897
- CompiledMod(vllm_config=vllm_config)(*args)
898
-
899
- # Phase 2: load from cache
900
- monkeypatch.setenv("VLLM_FORCE_AOT_LOAD", "1")
901
- disable_envs_cache()
902
-
903
- vllm_config = make_vllm_config()
904
- with (
905
- use_vllm_config(vllm_config),
906
- compilation_counter.expect(
907
- num_aot_compiles=0,
908
- num_aot_artifacts_saved=0,
909
- num_aot_artifacts_loaded=1,
910
- ),
911
- ):
912
- CompiledMod(vllm_config=vllm_config)(*args)
tests/compile/test_dynamic_shapes_compilation.py
+101
-68
@@ -3,6 +3,7 @@
3
4
import tempfile
5
from contextlib import contextmanager
6
+from dataclasses import dataclass
7
8
import pytest
9
import torch
@@ -23,40 +24,101 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer
24
def get_test_models():
25
"""Get list of models to test based on PyTorch version"""
26
models = [
27
+ "Qwen/Qwen3-0.6B",
28
"openai-community/gpt2",
27
- "Qwen/Qwen2-7B-Instruct",
28
- "meta-llama/Llama-3.1-8B",
29
]
30
- if is_torch_equal_or_newer("2.12.0.dev"):
31
- models.append("Qwen/Qwen3-4B-Instruct-2507")
30
return models
31
32
35
-@pytest.mark.parametrize("model_name", get_test_models())
36
-@pytest.mark.parametrize(
37
- "shapes_type",
38
- [
39
- DynamicShapesType.BACKED,
40
- DynamicShapesType.UNBACKED,
41
- DynamicShapesType.BACKED_SIZE_OBLIVIOUS,
42
- ],
43
-)
44
-@pytest.mark.parametrize("use_aot_compile", ["0", "1"])
45
-@pytest.mark.parametrize("use_bytecode_hook", [True, False])
46
-@pytest.mark.parametrize("evaluate_guards", [False, True])
33
+@dataclass(frozen=True)
34
+class DynamicShapesTestCase:
35
+ model_name: str
36
+ shapes_type: DynamicShapesType
37
+ use_aot_compile: bool
38
+ use_bytecode_hook: bool
39
+
40
+ def __str__(self) -> str:
41
+ model_id = self.model_name.rsplit("/", 1)[-1]
42
+ return (
43
+ f"{model_id}-{self.shapes_type.value}-"
44
+ f"aot{int(self.use_aot_compile)}-hook{int(self.use_bytecode_hook)}"
45
+ )
46
+
47
+
48
+def get_dynamic_shapes_test_cases():
49
+ """Pairwise-cover compilation options and smoke-test model classes."""
50
+ reference_model = "Qwen/Qwen3-0.6B"
51
+ cases = [
52
+ DynamicShapesTestCase(reference_model, DynamicShapesType.BACKED, False, True),
53
+ DynamicShapesTestCase(reference_model, DynamicShapesType.BACKED, True, False),
54
+ DynamicShapesTestCase(
55
+ reference_model, DynamicShapesType.UNBACKED, False, False
56
+ ),
57
+ DynamicShapesTestCase(reference_model, DynamicShapesType.UNBACKED, True, True),
58
+ DynamicShapesTestCase(
59
+ reference_model,
60
+ DynamicShapesType.BACKED_SIZE_OBLIVIOUS,
61
+ False,
62
+ True,
63
+ ),
64
+ DynamicShapesTestCase(
65
+ reference_model,
66
+ DynamicShapesType.BACKED_SIZE_OBLIVIOUS,
67
+ True,
68
+ False,
69
+ ),
70
+ ]
71
+ cases.extend(
72
+ DynamicShapesTestCase(model_name, DynamicShapesType.BACKED, False, True)
73
+ for model_name in get_test_models()
74
+ if model_name != reference_model
75
+ )
76
+ return cases
77
+
78
+
79
+TEST_PROMPTS = ["Hello, my name is", "The capital of France is"]
80
+
81
+
82
+def generate_outputs(vllm_model):
83
+ sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10)
84
+ outputs = []
85
+ for prompt in TEST_PROMPTS:
86
+ output = vllm_model.llm.generate(prompt, sampling_params)[0].outputs[0]
87
+ assert len(output.text.strip()) > 0, "Model produced empty output"
88
+ outputs.append((output.token_ids, output.text, output.logprobs))
89
+ return outputs
90
+
91
+
92
+@pytest.fixture(scope="module")
93
+def get_eager_outputs(vllm_runner):
94
+ cache = {}
95
+
96
+ def get(model_name):
97
+ if model_name not in cache:
98
+ with vllm_runner(
99
+ model_name,
100
+ enforce_eager=True,
101
+ max_model_len=1024,
102
+ enable_chunked_prefill=None,
103
+ ) as vllm_model:
104
+ cache[model_name] = generate_outputs(vllm_model)
105
+ return cache[model_name]
106
+
107
+ return get
108
+
109
+
110
+@pytest.mark.parametrize("test_case", get_dynamic_shapes_test_cases(), ids=str)
111
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
112
def test_dynamic_shapes_compilation(
113
monkeypatch,
114
vllm_runner,
51
- model_name,
52
- shapes_type,
53
- use_aot_compile,
54
- use_bytecode_hook,
55
- evaluate_guards,
115
+ get_eager_outputs,
116
+ test_case,
117
):
57
- """Test that all dynamic shapes types compile successfully"""
58
- if shapes_type == DynamicShapesType.UNBACKED and not is_torch_equal_or_newer(
59
- "2.11.0"
118
+ """Test representative dynamic-shapes configurations end to end."""
119
+ if (
120
+ test_case.shapes_type == DynamicShapesType.UNBACKED
121
+ and not is_torch_equal_or_newer("2.11.0")
122
):
123
# NOTE[ROCm]: shape_id (used by Qwen2/Llama to relate input dims) only
124
# landed in torch 2.11, but the ROCm CI still runs torch 2.10.x. On
@@ -64,64 +126,38 @@ def test_dynamic_shapes_compilation(
126
# data-dependent and compilation blows up -- nothing to test.
127
pytest.skip("unbacked dynamic shapes with shape_id require torch>=2.11")
128
67
- if evaluate_guards and shapes_type == DynamicShapesType.UNBACKED:
68
- pytest.skip("unbacked dynamic shapes do not add guards")
69
-
70
- # TODO is this still a requirement?
71
- if evaluate_guards and use_aot_compile:
72
- pytest.skip("evaluate_guards requires use_aot_compile=0")
73
-
74
- monkeypatch.setenv("VLLM_USE_AOT_COMPILE", use_aot_compile)
75
- monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
76
-
77
- prompt = "Hello, my name is"
78
-
79
- print(f"Testing {shapes_type.name} dynamic shapes...")
129
+ monkeypatch.setenv(
130
+ "VLLM_USE_AOT_COMPILE", "1" if test_case.use_aot_compile else "0"
131
+ )
132
+ monkeypatch.setenv(
133
+ "VLLM_USE_BYTECODE_HOOK", "1" if test_case.use_bytecode_hook else "0"
134
+ )
135
81
- sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10)
82
- test_prompts = [prompt, "The capital of France is"]
136
+ print(f"Testing {test_case.shapes_type.name} dynamic shapes...")
137
84
- # VllmRunner shuts down the engine core on exit, so the eager model
85
- # below never races a lingering compiled engine for GPU memory.
138
+ # The compiled engine shuts down before an uncached eager baseline starts.
139
with vllm_runner(
87
- model_name,
140
+ test_case.model_name,
141
compilation_config={
142
"mode": CompilationMode.VLLM_COMPILE,
143
"dynamic_shapes_config": {
91
- "type": shapes_type.value,
92
- "evaluate_guards": evaluate_guards,
144
+ "type": test_case.shapes_type.value,
145
},
146
},
147
max_model_len=1024,
148
enable_chunked_prefill=None,
149
) as vllm_model:
98
- compiled_outputs = []
99
- for p in test_prompts:
100
- output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0]
101
- assert len(output.text.strip()) > 0, "Compiled model produced empty output"
102
- compiled_outputs.append((output.token_ids, output.text, output.logprobs))
103
-
104
- with vllm_runner(
105
- model_name,
106
- enforce_eager=True,
107
- max_model_len=1024,
108
- enable_chunked_prefill=None,
109
- ) as vllm_model:
110
- eager_outputs = []
111
- for p in test_prompts:
112
- output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0]
113
- assert len(output.text.strip()) > 0, "Eager model produced empty output"
114
- eager_outputs.append((output.token_ids, output.text, output.logprobs))
150
+ compiled_outputs = generate_outputs(vllm_model)
151
152
check_logprobs_close(
117
- outputs_0_lst=eager_outputs,
153
+ outputs_0_lst=get_eager_outputs(test_case.model_name),
154
outputs_1_lst=compiled_outputs,
155
name_0="eager",
156
name_1="compiled",
157
)
158
159
124
-@pytest.mark.parametrize("use_aot_compile", ["0", "1"])
160
+@pytest.mark.parametrize("use_aot_compile", [False, True])
161
@pytest.mark.parametrize(
162
"dynamic_shapes_type",
163
[
@@ -138,7 +174,7 @@ def test_model_specialization_with_evaluate_guards(
174
"""
175
176
if (
141
- use_aot_compile == "1"
177
+ use_aot_compile
178
and dynamic_shapes_type == DynamicShapesType.BACKED
179
and evaluate_guards
180
):
@@ -177,12 +213,9 @@ def test_model_specialization_with_evaluate_guards(
213
yield
214
215
monkeypatch.setenv("TOKENIZERS_PARALLELISM", "true")
180
- monkeypatch.setenv("VLLM_USE_AOT_COMPILE", use_aot_compile)
216
+ monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1" if use_aot_compile else "0")
217
monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "0")
218
183
- # Create vllm config with the desired settings
184
- from vllm.config import CompilationMode
185
-
219
vllm_config = VllmConfig(
220
compilation_config=CompilationConfig(
221
mode=CompilationMode.VLLM_COMPILE,