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, 21:15 UTC
Add evaluate_guards option to DynamicShapesConfig (#27432)
Signed-off-by: Laith Sakka <lsakka@meta.com>
Laith Sakka committed
Dec 8, 2025 at 07:46 UTC
87aee9ed2b65b9ef7cdb83cc69dc1fc15f39ff36
6 files changed
+219
-28
tests/compile/test_dynamic_shapes_compilation.py
+132
-7
@@ -2,12 +2,21 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import gc
5
+import tempfile
6
+from contextlib import contextmanager
7
8
import pytest
9
import torch
10
11
from vllm import LLM, SamplingParams
10
-from vllm.config.compilation import CompilationMode, DynamicShapesType
12
+from vllm.compilation.decorators import support_torch_compile
13
+from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
14
+from vllm.config.compilation import (
15
+ CompilationMode,
16
+ DynamicShapesConfig,
17
+ DynamicShapesType,
18
+)
19
+from vllm.forward_context import set_forward_context
20
from vllm.tokenizers import get_tokenizer
21
from vllm.utils.torch_utils import is_torch_equal_or_newer
22
@@ -29,18 +38,19 @@ def get_test_models():
38
)
39
@pytest.mark.parametrize("use_aot_compile", ["0"])
40
@pytest.mark.parametrize("use_bytecode_hook", [True, False])
41
+@pytest.mark.parametrize("evaluate_guards", [False, True])
42
@pytest.mark.skipif(
43
not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
44
)
45
def test_dynamic_shapes_compilation(
36
- monkeypatch, model_name, shapes_type, use_aot_compile, use_bytecode_hook
46
+ monkeypatch,
47
+ model_name,
48
+ shapes_type,
49
+ use_aot_compile,
50
+ use_bytecode_hook,
51
+ evaluate_guards,
52
):
53
"""Test that all dynamic shapes types compile successfully"""
39
- print(
40
- f"\nTesting model: {model_name} with {shapes_type.name}, "
41
- f"AOT compile: {use_aot_compile}, "
42
- f"Bytecode hook: {use_bytecode_hook}"
43
- )
54
if use_bytecode_hook and shapes_type == DynamicShapesType.UNBACKED:
55
pytest.skip("UNBACKED dynamic shapes require VLLM_USE_BYTECODE_HOOK=0")
56
@@ -58,6 +68,7 @@ def test_dynamic_shapes_compilation(
68
"mode": CompilationMode.VLLM_COMPILE,
69
"dynamic_shapes_config": {
70
"type": shapes_type.value,
71
+ "evaluate_guards": evaluate_guards,
72
},
73
},
74
)
@@ -86,3 +97,117 @@ def test_dynamic_shapes_compilation(
97
torch.cuda.empty_cache()
98
torch.cuda.synchronize()
99
print("GPU memory cleared")
100
+
101
+
102
+@pytest.mark.parametrize("use_aot_compile", ["0", "1"])
103
+@pytest.mark.parametrize(
104
+ "dynamic_shapes_type",
105
+ [
106
+ DynamicShapesType.BACKED,
107
+ DynamicShapesType.BACKED_SIZE_OBLIVIOUS,
108
+ ],
109
+)
110
+@pytest.mark.parametrize("evaluate_guards", [False, True])
111
+def test_model_specialization_with_evaluate_guards(
112
+ monkeypatch, use_aot_compile, dynamic_shapes_type, evaluate_guards
113
+):
114
+ """Test that evaluate_guards correctly detects shape specialization
115
+ violations.
116
+ """
117
+
118
+ if (
119
+ use_aot_compile == "1"
120
+ and dynamic_shapes_type == DynamicShapesType.BACKED
121
+ and evaluate_guards
122
+ ):
123
+ pytest.skip("evaluate_guards for backed does not work with aot_compile =1")
124
+
125
+ @support_torch_compile
126
+ class ModelWithSizeCheck(torch.nn.Module):
127
+ def __init__(self, **kwargs):
128
+ super().__init__()
129
+
130
+ def forward(self, x: torch.Tensor):
131
+ # This will cause specialization - torch.compile will guard on
132
+ # sx.shape[0]
133
+ if x.shape[0] >= 10:
134
+ return x * 10
135
+ else:
136
+ return x * 10
137
+
138
+ @support_torch_compile
139
+ class ModelWithOneSizeCheck(torch.nn.Module):
140
+ def __init__(self, **kwargs):
141
+ super().__init__()
142
+
143
+ def forward(self, x: torch.Tensor):
144
+ # This will cause 0/1 specializations.
145
+ if x.shape[0] == 0:
146
+ return x * 10
147
+ if x.shape[0] == 1:
148
+ return x * 10
149
+ else:
150
+ return x * 10
151
+
152
+ @contextmanager
153
+ def use_vllm_config(vllm_config: VllmConfig):
154
+ with set_forward_context({}, vllm_config), set_current_vllm_config(vllm_config):
155
+ yield
156
+
157
+ monkeypatch.setenv("TOKENIZERS_PARALLELISM", "true")
158
+ monkeypatch.setenv("VLLM_USE_AOT_COMPILE", use_aot_compile)
159
+ monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "0")
160
+
161
+ # Create vllm config with the desired settings
162
+ from vllm.config import CompilationMode
163
+
164
+ vllm_config = VllmConfig(
165
+ compilation_config=CompilationConfig(
166
+ mode=CompilationMode.VLLM_COMPILE,
167
+ dynamic_shapes_config=DynamicShapesConfig(
168
+ type=dynamic_shapes_type,
169
+ evaluate_guards=evaluate_guards,
170
+ ),
171
+ )
172
+ )
173
+
174
+ def test(model_class, input1, input2, is_01_specialization=False):
175
+ with (
176
+ torch.no_grad(),
177
+ use_vllm_config(vllm_config),
178
+ tempfile.TemporaryDirectory() as tmpdirname,
179
+ ):
180
+ monkeypatch.setenv("VLLM_CACHE_ROOT", tmpdirname)
181
+
182
+ model = model_class(vllm_config=vllm_config).cuda()
183
+
184
+ model(input1)
185
+
186
+ if evaluate_guards and (
187
+ not (
188
+ is_01_specialization
189
+ and dynamic_shapes_type == DynamicShapesType.BACKED
190
+ )
191
+ ):
192
+ # This should fail because guards were added.
193
+ with pytest.raises(RuntimeError) as excinfo:
194
+ model(input2)
195
+
196
+ # Expected failure - guard was violated
197
+ error_msg = str(excinfo.value)
198
+ assert (
199
+ "GuardManager check failed" in error_msg
200
+ or "Detected recompile when torch.compile stance" in error_msg
201
+ ), error_msg
202
+
203
+ else:
204
+ model(input2)
205
+
206
+ test(ModelWithSizeCheck, torch.randn(20, 10).cuda(), torch.randn(5, 10).cuda())
207
+ test(ModelWithSizeCheck, torch.randn(5, 10).cuda(), torch.randn(20, 10).cuda())
208
+ test(
209
+ ModelWithOneSizeCheck,
210
+ torch.randn(20, 10).cuda(),
211
+ torch.randn(1, 10).cuda(),
212
+ is_01_specialization=True,
213
+ )
vllm/compilation/backends.py
+24
-2
@@ -26,6 +26,7 @@ from vllm.compilation.partition_rules import (
26
should_split,
27
)
28
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
29
+from vllm.config.compilation import DynamicShapesType
30
from vllm.config.utils import Range, hash_factors
31
from vllm.logger import init_logger
32
from vllm.logging_utils import lazy
@@ -722,6 +723,29 @@ class VllmBackend:
723
self.split_gm, submod_names_to_compile, self.vllm_config, self
724
).run(*fake_args)
725
726
+ from torch._guards import detect_fake_mode
727
+
728
+ fake_mode = detect_fake_mode()
729
+
730
+ if (
731
+ self.compilation_config.dynamic_shapes_config.evaluate_guards
732
+ and self.compilation_config.dynamic_shapes_config.type
733
+ == DynamicShapesType.BACKED
734
+ ):
735
+ from torch.utils._sympy.value_ranges import ValueRanges
736
+
737
+ # Drop counter-0/1 specializations guards; for backed dynamic shapes,
738
+ # torch.compile will specialize for 0/1 inputs or otherwise guards that
739
+ # shape is >= 2. This is because it's really hard not to hit a check
740
+ # against 0/1. When we evaluate shape guards, we exclude checking those
741
+ # guards (We would fail always otherwise).
742
+
743
+ # We avoid that by updating the ranges of backed sizes when the min is
744
+ # 2 for any, we assume it's 0.
745
+ for s, r in fake_mode.shape_env.var_to_range.items():
746
+ if r.lower == 2:
747
+ fake_mode.shape_env.var_to_range[s] = ValueRanges(0, r.upper)
748
+
749
graph_path = os.path.join(local_cache_dir, "computation_graph.py")
750
if not os.path.exists(graph_path):
751
# code adapted from
@@ -749,8 +773,6 @@ class VllmBackend:
773
graph, example_inputs, self.prefix, self.split_gm
774
)
775
752
- # if we need to copy input buffers for cudagraph
753
- #
776
# index of tensors that have symbolic shapes (batch size)
777
# for weights and static buffers, they will have concrete shapes.
778
# symbolic shape only happens for input tensors.
vllm/compilation/decorators.py
+2
-2
@@ -392,7 +392,6 @@ def _support_torch_compile(
392
393
factors.append(_model_hash_key(self.forward))
394
hash_key = hashlib.sha256(str(factors).encode()).hexdigest()
395
-
395
cache_dir = os.path.join(
396
envs.VLLM_CACHE_ROOT,
397
"torch_aot_compile",
@@ -413,7 +412,8 @@ def _support_torch_compile(
412
f, f_globals=self.forward.__globals__
413
)
414
_verify_source_unchanged(loaded_fn.source_info(), self.vllm_config)
416
- loaded_fn.disable_guard_check()
415
+ if not self.compilation_config.dynamic_shapes_config.evaluate_guards:
416
+ loaded_fn.disable_guard_check()
417
self.aot_compiled_fn = loaded_fn
418
except Exception as e:
419
if os.path.exists(aot_compilation_path):
vllm/compilation/wrapper.py
+46
-13
@@ -4,7 +4,7 @@
4
import os
5
import sys
6
from abc import abstractmethod
7
-from contextlib import contextmanager
7
+from contextlib import contextmanager, nullcontext
8
from types import CodeType
9
from typing import Any
10
@@ -13,6 +13,7 @@ import torch._C._dynamo.guards
13
14
import vllm.envs as envs
15
from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config
16
+from vllm.config.compilation import DynamicShapesType
17
from vllm.logger import init_logger
18
from vllm.utils.nvtx_pytorch_hooks import layerwise_nvtx_marker_context
19
@@ -125,23 +126,49 @@ class TorchCompileWithNoGuardsWrapper:
126
if isinstance(backend, str) and backend == "inductor":
127
options = vllm_config.compilation_config.inductor_compile_config
128
129
+ self.first_compile = True
130
+ self.evaluate_guards = (
131
+ vllm_config.compilation_config.dynamic_shapes_config.evaluate_guards
132
+ )
133
+
134
+ ds_type = vllm_config.compilation_config.dynamic_shapes_config.type
135
+
136
if mode != CompilationMode.STOCK_TORCH_COMPILE:
137
# Drop all the guards.
130
- options["guard_filter_fn"] = lambda x: [False for _ in x]
138
+ if self.evaluate_guards:
139
+ assert not envs.VLLM_USE_BYTECODE_HOOK, (
140
+ "compilation_config.dynamic_shapes_config.evaluate_guards "
141
+ "requires VLLM_USE_BYTECODE_HOOK=0. "
142
+ )
143
132
- # Validate that unbacked dynamic shapes require VLLM_USE_BYTECODE_HOOK=False
133
- from vllm.compilation.decorators import DynamicShapesType
144
+ if envs.VLLM_USE_AOT_COMPILE:
145
+ # disabled until https://github.com/pytorch/pytorch/pull/169239
146
+ # is picked up.
147
+ assert ds_type != DynamicShapesType.BACKED, (
148
+ "evaluate_guards for backed shapes requires "
149
+ "VLLM_USE_AOT_COMPILE=False. "
150
+ )
151
+
152
+ options["guard_filter_fn"] = lambda x: [
153
+ entry.guard_type == "SHAPE_ENV" for entry in x
154
+ ]
155
+ else:
156
+ options["guard_filter_fn"] = lambda x: [False for _ in x]
157
135
- ds_type = vllm_config.compilation_config.dynamic_shapes_config.type
158
compiled_ptr: Any = self.forward
159
+ # Validate that unbacked dynamic shapes require VLLM_USE_BYTECODE_HOOK=False
160
+
161
if ds_type == DynamicShapesType.UNBACKED:
138
- if envs.VLLM_USE_BYTECODE_HOOK:
139
- # reason is that bytecode does this hack torch._dynamo.eval_frame.
140
- # remove_from_cache(self.original_code_object()) to force a new
141
- # re-compilation.
142
- raise ValueError(
143
- "UNBACKED dynamic shapes require VLLM_USE_BYTECODE_HOOK=0. "
144
- )
162
+ # reason is that bytecode does torch._dynamo.eval_frame.
163
+ # remove_from_cache(self.original_code_object()) to force a new
164
+ # re-compilation. And if we use
165
+ # compiled_ptr = self.check_invariants_and_forward
166
+ # it will reset all entries.
167
+ assert not envs.VLLM_USE_BYTECODE_HOOK, (
168
+ "UNBACKED dynamic shapes requires VLLM_USE_BYTECODE_HOOK=0. "
169
+ )
170
+ assert not self.evaluate_guards, "UNBACKED dynamic shapes do not add guards"
171
+
172
compiled_ptr = self.check_invariants_and_forward
173
174
if envs.VLLM_USE_AOT_COMPILE:
@@ -195,7 +222,13 @@ class TorchCompileWithNoGuardsWrapper:
222
self.forward, *args, **kwargs
223
)
224
else:
198
- with _compilation_context():
225
+ ctx = (
226
+ nullcontext()
227
+ if self.first_compile or not self.evaluate_guards
228
+ else torch.compiler.set_stance("fail_on_recompile")
229
+ )
230
+ self.first_compile = False
231
+ with _compilation_context(), ctx:
232
return self._call_with_optional_nvtx_range(
233
self._compiled_callable, *args, **kwargs
234
)
vllm/config/compilation.py
+14
-3
@@ -344,7 +344,18 @@ class DynamicShapesConfig:
344
backed/unbacked.
345
"""
346
347
- # TODO add a debug mode to fail
347
+ evaluate_guards: bool = False
348
+ """
349
+ A debug mode to detect and fail if Dynamo ever specializes a dynamic shape by
350
+ guarding on it. When True, dynamic shape guards are not dropped from dynamo.
351
+ And a failure will be triggered if a recompilation ever happens due to that.
352
+ This mode requires VLLM_USE_BYTECODE_HOOK to be 0.
353
+ Enabling this allow observing the dynamic shapes guards in the tlparse
354
+ artifacts also.
355
+ When type is backed, aot_compile must be disabled for this mode to work.
356
+ until this change picked up https://github.com/pytorch/pytorch/pull/169239.
357
+
358
+ """
359
360
def compute_hash(self) -> str:
361
"""
@@ -455,8 +466,8 @@ class CompilationConfig:
466
We use string to avoid serialization issues when using compilation in a
467
distributed setting. When the compilation mode is 1 or 2, the backend is
468
used for the compilation directly (it sees the whole graph). When the
458
- compilation mode is 3, the backend supports both whole graph and piecewise
459
- compilation, available backends include eager, inductor, and custom backends,
469
+ compilation mode is 3, the backend supports both whole graph and piecewise
470
+ compilation, available backends include eager, inductor, and custom backends,
471
the latter of which can be defined via `get_compile_backend`. Furthermore,
472
compilation is only piecewise if splitting ops is set accordingly and
473
use_inductor_graph_partition is off. Note that the default options for
vllm/config/vllm.py
+1
-1
@@ -66,7 +66,7 @@ class OptimizationLevel(IntEnum):
66
"""O0 : No optimization. no compilation, no cudagraphs, no other
67
optimization, just starting up immediately"""
68
O1 = 1
69
- """O1: Quick optimizations. Dynamo+Inductor compilation and Piecewise
69
+ """O1: Quick optimizations. Dynamo+Inductor compilation and Piecewise
70
cudagraphs"""
71
O2 = 2
72
"""O2: Full optimizations. -O1 as well as Full and Piecewise cudagraphs."""