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, 23:45 UTC
[aot_compile]change VLLM backend to read fake args from example_value (#29104)
Signed-off-by: Laith Sakka <lsakka@meta.com>
Laith Sakka committed
Dec 4, 2025 at 14:33 UTC
1f0d1845909877b3492d725cb63b386d80c39b47
3 files changed
+81
-10
tests/compile/test_aot_compile.py
+66
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import functools
5
+import multiprocessing
6
import tempfile
7
from contextlib import contextmanager
8
@@ -137,3 +139,67 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
139
artifacts = compiled_mod.aot_compiled_fn._artifacts
140
guards_string = artifacts.compiled_fn.shape_env.format_guards()
141
assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
142
+
143
+
144
+@pytest.mark.skipif(
145
+ not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
146
+)
147
+@use_vllm_config(make_vllm_config())
148
+def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
149
+ """
150
+ Test that compiling gpt2 twice results in a cache hit and
151
+ capture torch dynamic symbol creations to ensure make_symbol
152
+ not called on cache hit.
153
+ """
154
+
155
+ import torch.fx.experimental.symbolic_shapes as symbolic_shapes_module
156
+ from torch.utils._sympy.symbol import make_symbol
157
+
158
+ from vllm import LLM
159
+
160
+ create_symbol_counter = multiprocessing.Value("i", 0)
161
+ original_make_symbol = make_symbol
162
+
163
+ @functools.wraps(original_make_symbol)
164
+ def counting_make_symbol(prefix, idx, **kwargs):
165
+ with create_symbol_counter.get_lock():
166
+ create_symbol_counter.value += 1
167
+ return original_make_symbol(prefix, idx, **kwargs)
168
+
169
+ symbolic_shapes_module.make_symbol = counting_make_symbol
170
+ try:
171
+ with monkeypatch.context() as m, tempfile.TemporaryDirectory() as tmpdirname:
172
+ m.setenv("VLLM_CACHE_ROOT", tmpdirname)
173
+ m.setenv("VLLM_USE_AOT_COMPILE", "1")
174
+ # First compilation - initialize model and generate
175
+ llm_model = LLM(
176
+ model="gpt2",
177
+ compilation_config=CompilationConfig(
178
+ mode=CompilationMode.VLLM_COMPILE,
179
+ ),
180
+ max_model_len=256,
181
+ )
182
+
183
+ llm_model.generate("Hello, my name is")
184
+ assert create_symbol_counter.value == 2
185
+ create_symbol_counter.value = 0
186
+
187
+ # Clean up first model
188
+ del llm_model
189
+
190
+ # Second compilation - should hit cache
191
+ m.setenv("VLLM_FORCE_AOT_LOAD", "1")
192
+ llm_model = LLM(
193
+ model="gpt2",
194
+ compilation_config=CompilationConfig(
195
+ mode=CompilationMode.VLLM_COMPILE,
196
+ ),
197
+ max_model_len=256,
198
+ )
199
+ llm_model.generate("Hello, my name is")
200
+
201
+ assert create_symbol_counter.value == 0
202
+
203
+ finally:
204
+ # Restore original method
205
+ symbolic_shapes_module.make_symbol = original_make_symbol
vllm/compilation/backends.py
+15
-9
@@ -402,6 +402,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
402
self.extra_traceback = False
403
404
def run(self, *args):
405
+ # maybe instead just assert inputs are fake?
406
fake_args = [
407
self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t
408
for t in args
@@ -416,11 +417,13 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
417
kwargs: dict[str, Any],
418
) -> Any:
419
assert isinstance(target, str)
420
+
421
output = super().call_module(target, args, kwargs)
422
423
if target in self.compile_submod_names:
424
index = self.compile_submod_names.index(target)
425
submod = self.fetch_attr(target)
426
+
427
sym_shape_indices = [
428
i for i, x in enumerate(args) if isinstance(x, torch.SymInt)
429
]
@@ -746,11 +749,21 @@ class VllmBackend:
749
if not item.is_splitting_graph
750
]
751
752
+ # Extract fake values from the graph to use them when needed.
753
+ all_fake_values = []
754
+ for i in graph.graph.find_nodes(op="placeholder"):
755
+ all_fake_values.append(i.meta["example_value"])
756
+
757
+ fake_args = [
758
+ all_fake_values[i] if isinstance(t, torch.Tensor) else t
759
+ for i, t in enumerate(example_inputs)
760
+ ]
761
+
762
# propagate the split graph to the piecewise backend,
763
# compile submodules with symbolic shapes
764
PiecewiseCompileInterpreter(
765
self.split_gm, submod_names_to_compile, self.vllm_config, self
753
- ).run(*example_inputs)
766
+ ).run(*fake_args)
767
768
graph_path = os.path.join(local_cache_dir, "computation_graph.py")
769
if not os.path.exists(graph_path):
@@ -780,14 +793,7 @@ class VllmBackend:
793
)
794
795
# if we need to copy input buffers for cudagraph
783
- from torch._guards import detect_fake_mode
784
-
785
- fake_mode = detect_fake_mode()
786
- fake_args = [
787
- fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t
788
- for t in example_inputs
789
- ]
790
-
796
+ #
797
# index of tensors that have symbolic shapes (batch size)
798
# for weights and static buffers, they will have concrete shapes.
799
# symbolic shape only happens for input tensors.
vllm/compilation/decorators.py
-1
@@ -433,7 +433,6 @@ def _support_torch_compile(
433
return TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs)
434
435
# This is the path for the first compilation.
436
-
436
# the first compilation needs to have dynamic shapes marked
437
_mark_dynamic_inputs(
438
self,