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 14, 05:30 UTC
AOT Compilation for torch.compile (Bundled) (#24274)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Oct 10, 2025 at 19:02 UTC
eef921f45e7d3efb2ed2ccab80ee20ee2e4ebe38
9 files changed
+484
-41
.buildkite/test-pipeline.yaml
+1
@@ -403,6 +403,7 @@ steps:
403
- pytest -v -s compile/test_fusion_all_reduce.py
404
- pytest -v -s compile/test_decorator.py
405
- pytest -v -s compile/test_noop_elimination.py
406
+ - pytest -v -s compile/test_aot_compile.py
407
408
- label: PyTorch Fullgraph Smoke Test # 15min
409
timeout_in_minutes: 30
tests/compile/test_aot_compile.py
new
+139
@@ -0,0 +1,139 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import tempfile
5
+from contextlib import contextmanager
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.compilation.decorators import support_torch_compile
11
+from vllm.config import (
12
+ CompilationConfig,
13
+ CompilationLevel,
14
+ VllmConfig,
15
+ set_current_vllm_config,
16
+)
17
+from vllm.forward_context import set_forward_context
18
+from vllm.utils import is_torch_equal_or_newer
19
+
20
+
21
+def reference_fn(x: torch.Tensor):
22
+ assert x.shape[0] <= 42
23
+ assert x.shape[0] % 2 == 0
24
+ for _ in range(3000):
25
+ x = x + x.shape[0]
26
+ return x
27
+
28
+
29
+@support_torch_compile
30
+class CompiledMod(torch.nn.Module):
31
+ def __init__(self, **kwargs):
32
+ super().__init__()
33
+
34
+ def forward(self, x: torch.Tensor):
35
+ return reference_fn(x)
36
+
37
+
38
+def make_vllm_config() -> VllmConfig:
39
+ return VllmConfig(
40
+ compilation_config=CompilationConfig(
41
+ level=CompilationLevel.PIECEWISE,
42
+ )
43
+ )
44
+
45
+
46
+@contextmanager
47
+def use_vllm_config(vllm_config: VllmConfig):
48
+ with set_forward_context({}, vllm_config), set_current_vllm_config(vllm_config):
49
+ yield
50
+
51
+
52
+@pytest.mark.skipif(
53
+ not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
54
+)
55
+def test_no_dynamo_cache_entry(monkeypatch: pytest.MonkeyPatch):
56
+ with monkeypatch.context() as m:
57
+ vllm_config = make_vllm_config()
58
+ args = (torch.randn(10, 10),)
59
+ expected = reference_fn(*args)
60
+ with use_vllm_config(vllm_config):
61
+ m.setenv("VLLM_USE_AOT_COMPILE", "0")
62
+ with (
63
+ pytest.raises(RuntimeError, match="Detected recompile"),
64
+ torch.compiler.set_stance("fail_on_recompile"),
65
+ ):
66
+ CompiledMod(vllm_config=vllm_config)(*args)
67
+
68
+ m.setenv("VLLM_USE_AOT_COMPILE", "1")
69
+ torch._dynamo.reset()
70
+ with torch.compiler.set_stance("fail_on_recompile"):
71
+ actual = CompiledMod(vllm_config=vllm_config)(*args)
72
+ assert torch.allclose(actual, expected)
73
+
74
+
75
+@pytest.mark.skipif(
76
+ not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
77
+)
78
+def test_force_aot_load(monkeypatch: pytest.MonkeyPatch):
79
+ with tempfile.TemporaryDirectory() as tmpdirname, monkeypatch.context() as m:
80
+ args = (torch.randn(10, 10),)
81
+ m.setenv("VLLM_USE_AOT_COMPILE", "1")
82
+ m.setenv("VLLM_FORCE_AOT_LOAD", "1")
83
+ m.setenv("VLLM_CACHE_ROOT", tmpdirname)
84
+ vllm_config = make_vllm_config()
85
+ with use_vllm_config(vllm_config), pytest.raises(FileNotFoundError):
86
+ CompiledMod(vllm_config=vllm_config)(*args)
87
+
88
+
89
+@pytest.mark.skipif(
90
+ not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
91
+)
92
+def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
93
+ with monkeypatch.context() as m:
94
+ args = (torch.randn(10, 10),)
95
+
96
+ with tempfile.TemporaryDirectory() as tmpdirname:
97
+ m.setenv("VLLM_CACHE_ROOT", tmpdirname)
98
+ m.setenv("VLLM_USE_AOT_COMPILE", "1")
99
+ vllm_config = make_vllm_config()
100
+ with use_vllm_config(vllm_config):
101
+ expected = CompiledMod(vllm_config=vllm_config)(*args)
102
+
103
+ m.setenv("VLLM_FORCE_AOT_LOAD", "1")
104
+ vllm_config = make_vllm_config()
105
+ with use_vllm_config(vllm_config):
106
+ ret = CompiledMod(vllm_config=vllm_config)(*args)
107
+ assert torch.allclose(ret, expected)
108
+
109
+
110
+@pytest.mark.skipif(
111
+ not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
112
+)
113
+def test_shape_env(monkeypatch: pytest.MonkeyPatch):
114
+ """
115
+ Test that the shape environment is correctly serialized and preserved
116
+ when loading from cache.
117
+ """
118
+ with monkeypatch.context() as m:
119
+ args = (torch.randn(10, 10),)
120
+
121
+ with tempfile.TemporaryDirectory() as tmpdirname:
122
+ m.setenv("VLLM_CACHE_ROOT", tmpdirname)
123
+ m.setenv("VLLM_USE_AOT_COMPILE", "1")
124
+ vllm_config = make_vllm_config()
125
+ with use_vllm_config(vllm_config):
126
+ compiled_mod = CompiledMod(vllm_config=vllm_config)
127
+ compiled_mod(*args)
128
+ artifacts = compiled_mod.aot_compiled_fn._artifacts
129
+ guards_string = artifacts.compiled_fn.shape_env.format_guards()
130
+ assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
131
+
132
+ m.setenv("VLLM_FORCE_AOT_LOAD", "1")
133
+ vllm_config = make_vllm_config()
134
+ with use_vllm_config(vllm_config):
135
+ compiled_mod = CompiledMod(vllm_config=vllm_config)
136
+ compiled_mod(*args)
137
+ artifacts = compiled_mod.aot_compiled_fn._artifacts
138
+ guards_string = artifacts.compiled_fn.shape_env.format_guards()
139
+ assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
tools/pre_commit/check_pickle_imports.py
+1
@@ -22,6 +22,7 @@ ALLOWED_FILES = {
22
"vllm/multimodal/hasher.py",
23
"vllm/transformers_utils/config.py",
24
"vllm/model_executor/models/registry.py",
25
+ "vllm/compilation/caching.py",
26
"tests/utils_/test_utils.py",
27
"tests/tokenization/test_cached_tokenizer.py",
28
"vllm/distributed/utils.py",
vllm/compilation/backends.py
+16
-33
@@ -3,6 +3,7 @@
3
4
import ast
5
import dataclasses
6
+import hashlib
7
import os
8
import pprint
9
import time
@@ -25,6 +26,7 @@ from vllm.logger import init_logger
26
from vllm.platforms import current_platform
27
from vllm.utils import is_torch_equal_or_newer, resolve_obj_by_qualname
28
29
+from .caching import VllmSerializableFunction
30
from .compiler_interface import (
31
CompilerInterface,
32
EagerAdaptor,
@@ -195,6 +197,7 @@ class CompilerManager:
197
# there can be multiple graphs due to piecewise compilation.
198
now = time.time()
199
elapsed = now - compilation_start_time
200
+ compilation_config.compilation_time += elapsed
201
if runtime_shape is None:
202
logger.info(
203
"Directly load the compiled graph(s) for dynamic shape "
@@ -549,7 +552,11 @@ class VllmBackend:
552
self.post_grad_pass_manager.add(inductor_config[PASS_KEY])
553
inductor_config[PASS_KEY] = self.post_grad_pass_manager
554
552
- def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:
555
+ def __call__(
556
+ self, graph: fx.GraphModule, example_inputs
557
+ ) -> VllmSerializableFunction:
558
+ from .caching import _compute_code_hash, compilation_config_hash_factors
559
+
560
vllm_config = self.vllm_config
561
if not self.compilation_config.cache_dir:
562
# no provided cache dir, generate one based on the known factors
@@ -557,39 +564,11 @@ class VllmBackend:
564
# the cache dir will be the same so that we can reuse the compiled
565
# graph.
566
560
- factors = []
561
- # 0. factors come from the env, for example, The values of
562
- # VLLM_PP_LAYER_PARTITION will affect the computation graph.
563
- env_hash = envs.compute_hash()
564
- factors.append(env_hash)
565
-
566
- # 1. factors come from the vllm_config (it mainly summarizes how the
567
- # model is created)
568
- config_hash = vllm_config.compute_hash()
569
- factors.append(config_hash)
570
-
567
+ factors = compilation_config_hash_factors(vllm_config)
568
# 2. factors come from the code files that are traced by Dynamo (
569
# it mainly summarizes how the model is used in forward pass)
573
- forward_code_files = list(sorted(self.compilation_config.traced_files))
570
+ code_hash = _compute_code_hash(self.compilation_config.traced_files)
571
self.compilation_config.traced_files.clear()
575
- logger.debug(
576
- "Traced files (to be considered for compilation cache):\n%s",
577
- "\n".join(forward_code_files),
578
- )
579
- hash_content = []
580
- for filepath in forward_code_files:
581
- hash_content.append(filepath)
582
- if filepath == "<string>":
583
- # This means the function was dynamically generated, with
584
- # e.g. exec(). We can't actually check these.
585
- continue
586
- with open(filepath) as f:
587
- hash_content.append(f.read())
588
- import hashlib
589
-
590
- code_hash = hashlib.md5(
591
- "\n".join(hash_content).encode(), usedforsecurity=False
592
- ).hexdigest()
572
factors.append(code_hash)
573
574
# 3. compiler hash
@@ -695,7 +674,9 @@ class VllmBackend:
674
self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
675
or not self.compilation_config.cudagraph_copy_inputs
676
):
698
- return self.split_gm
677
+ return VllmSerializableFunction(
678
+ graph, example_inputs, self.prefix, self.split_gm
679
+ )
680
681
# if we need to copy input buffers for cudagraph
682
from torch._guards import detect_fake_mode
@@ -740,4 +721,6 @@ class VllmBackend:
721
list_args[index] = static_tensor
722
return self.split_gm(*list_args)
723
743
- return copy_and_call
724
+ return VllmSerializableFunction(
725
+ graph, example_inputs, self.prefix, copy_and_call
726
+ )
vllm/compilation/caching.py
new
+176
@@ -0,0 +1,176 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import hashlib
5
+import inspect
6
+import pickle
7
+from unittest.mock import patch
8
+
9
+import torch
10
+from torch.utils import _pytree as pytree
11
+
12
+import vllm.envs as envs
13
+from vllm.config import VllmConfig, get_current_vllm_config
14
+from vllm.logger import init_logger
15
+
16
+try:
17
+ from torch._dynamo.aot_compile import SerializableCallable
18
+except ImportError:
19
+ SerializableCallable = object
20
+
21
+assert isinstance(SerializableCallable, type)
22
+
23
+logger = init_logger(__name__)
24
+
25
+
26
+class VllmSerializableFunction(SerializableCallable):
27
+ """
28
+ A wrapper around a compiled function by vllm. It will forward the tensor
29
+ inputs to the compiled function and return the result.
30
+ It also implements a serialization interface to support PyTorch's precompile
31
+ with custom backend, so that we can save and load the compiled function on
32
+ disk. There's no need to wrap around the compiled function if we don't want
33
+ to serialize them in particular cases.
34
+ Right now serialization for the custom backend is done via
35
+ serializing the Dynamo fx graph plus example inputs.
36
+ """
37
+
38
+ def __init__(self, graph_module, example_inputs, prefix, optimized_call):
39
+ assert isinstance(graph_module, torch.fx.GraphModule)
40
+ self.graph_module = graph_module
41
+ self.example_inputs = example_inputs
42
+ self.prefix = prefix
43
+ self.optimized_call = optimized_call
44
+ self.shape_env = None
45
+ sym_input = next(
46
+ (i for i in self.example_inputs if isinstance(i, torch.SymInt)), None
47
+ )
48
+ if sym_input is not None:
49
+ self.shape_env = sym_input.node.shape_env
50
+
51
+ def __call__(self, *args, **kwargs):
52
+ return self.optimized_call(*args, **kwargs)
53
+
54
+ @classmethod
55
+ def serialize_compile_artifacts(
56
+ cls, compiled_fn: "VllmSerializableFunction"
57
+ ) -> bytes:
58
+ import sympy
59
+ from torch._subclasses import FakeTensorMode
60
+ from torch.fx._graph_pickler import GraphPickler, Options
61
+
62
+ state = compiled_fn.__dict__.copy()
63
+ state.pop("optimized_call")
64
+ state.pop("shape_env")
65
+ for node in state["graph_module"].graph.nodes:
66
+ node.meta.pop("source_fn_stack", None)
67
+ node.meta.pop("nn_module_stack", None)
68
+
69
+ graph_reducer_override = GraphPickler.reducer_override
70
+
71
+ def _graph_reducer_override(self, obj):
72
+ if (
73
+ inspect.isclass(obj)
74
+ and issubclass(obj, sympy.Function)
75
+ and hasattr(obj, "_torch_unpickler")
76
+ ):
77
+ return obj._torch_unpickler, (obj._torch_handler_name,)
78
+ if isinstance(obj, FakeTensorMode):
79
+ return type(None), ()
80
+ return graph_reducer_override(self, obj)
81
+
82
+ # Mask off tensor inputs since they are large and not needed.
83
+ state["example_inputs"] = pytree.tree_map_only(
84
+ torch.Tensor, lambda _: None, state["example_inputs"]
85
+ )
86
+ with patch.object(GraphPickler, "reducer_override", _graph_reducer_override):
87
+ state["graph_module"] = GraphPickler.dumps(
88
+ state["graph_module"], Options(ops_filter=None)
89
+ )
90
+ state["example_inputs"] = GraphPickler.dumps(state["example_inputs"])
91
+ return pickle.dumps(state)
92
+
93
+ @classmethod
94
+ def deserialize_compile_artifacts(cls, data: bytes) -> "VllmSerializableFunction":
95
+ from torch._guards import TracingContext, tracing
96
+ from torch._subclasses import FakeTensorMode
97
+ from torch.fx._graph_pickler import GraphPickler
98
+ from torch.fx.experimental.symbolic_shapes import ShapeEnv
99
+
100
+ from vllm.compilation.backends import VllmBackend
101
+
102
+ state = pickle.loads(data)
103
+ fake_mode = FakeTensorMode(shape_env=ShapeEnv())
104
+ state["graph_module"] = GraphPickler.loads(state["graph_module"], fake_mode)
105
+ state["example_inputs"] = GraphPickler.loads(state["example_inputs"], fake_mode)
106
+ vllm_backend = VllmBackend(get_current_vllm_config(), state["prefix"])
107
+
108
+ def optimized_call(*example_inputs):
109
+ """
110
+ On the first run of the optimized call, we rerun the compiler
111
+ backend which should result in a cache hit. After the backend
112
+ call returns, we just do a one-time replacement of the optimized
113
+ call with the compiled function, so that subsequent calls are on
114
+ the AOT compiled path.
115
+ """
116
+ compile_inputs = [
117
+ inp or example_inputs[i] for i, inp in enumerate(fn.example_inputs)
118
+ ]
119
+ with tracing(TracingContext(fake_mode)):
120
+ fn.optimized_call = vllm_backend(
121
+ state["graph_module"], compile_inputs
122
+ ).optimized_call
123
+ return fn.optimized_call(*example_inputs)
124
+
125
+ fn = cls(**state, optimized_call=optimized_call)
126
+ return fn
127
+
128
+ @property
129
+ def co_name(self):
130
+ """
131
+ Used for depyf debugging.
132
+ """
133
+ return "VllmSerializableFunction"
134
+
135
+
136
+def compilation_config_hash_factors(vllm_config: VllmConfig) -> list[str]:
137
+ factors = []
138
+ # 0. factors come from the env, for example, The values of
139
+ # VLLM_PP_LAYER_PARTITION will affect the computation graph.
140
+ env_hash = envs.compute_hash()
141
+ factors.append(env_hash)
142
+
143
+ # 1. factors come from the vllm_config (it mainly summarizes how the
144
+ # model is created)
145
+ config_hash = vllm_config.compute_hash()
146
+ factors.append(config_hash)
147
+ return factors
148
+
149
+
150
+def _compute_code_hash_with_content(file_contents: dict[str, str]) -> str:
151
+ items = list(sorted(file_contents.items(), key=lambda x: x[0]))
152
+ hash_content = []
153
+ for filepath, content in items:
154
+ hash_content.append(filepath)
155
+ if filepath == "<string>":
156
+ # This means the function was dynamically generated, with
157
+ # e.g. exec(). We can't actually check these.
158
+ continue
159
+ hash_content.append(content)
160
+ return hashlib.md5(
161
+ "\n".join(hash_content).encode(), usedforsecurity=False
162
+ ).hexdigest()
163
+
164
+
165
+def _compute_code_hash(files: set[str]) -> str:
166
+ logger.debug(
167
+ "Traced files (to be considered for compilation cache):\n%s", "\n".join(files)
168
+ )
169
+ file_contents = {}
170
+ for filepath in files:
171
+ if filepath == "<string>":
172
+ file_contents[filepath] = ""
173
+ else:
174
+ with open(filepath) as f:
175
+ file_contents[filepath] = f.read()
176
+ return _compute_code_hash_with_content(file_contents)
vllm/compilation/compiler_interface.py
+6
@@ -199,6 +199,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
199
if compiler_config is not None:
200
current_config.update(compiler_config)
201
set_inductor_config(current_config, runtime_shape)
202
+ set_functorch_config()
203
204
if isinstance(runtime_shape, int):
205
dynamic_shapes = "from_example_inputs"
@@ -307,6 +308,7 @@ class InductorAdaptor(CompilerInterface):
308
current_config["fx_graph_remote_cache"] = False
309
310
set_inductor_config(current_config, runtime_shape)
311
+ set_functorch_config()
312
313
# inductor can inplace modify the graph, so we need to copy it
314
# see https://github.com/pytorch/pytorch/issues/138980
@@ -596,6 +598,10 @@ def set_inductor_config(config, runtime_shape):
598
)
599
600
601
+def set_functorch_config():
602
+ torch._functorch.config.bundled_autograd_cache = False
603
+
604
+
605
class EagerAdaptor(CompilerInterface):
606
name = "eager"
607
vllm/compilation/decorators.py
+105
-8
@@ -2,7 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import contextlib
5
+import hashlib
6
import inspect
7
+import os
8
+import sys
9
from typing import Callable, Optional, TypeVar, Union, overload
10
from unittest.mock import patch
11
@@ -11,9 +14,10 @@ import torch.nn as nn
14
from packaging import version
15
from torch._dynamo.symbolic_convert import InliningInstructionTranslator
16
17
+import vllm.envs as envs
18
from vllm.compilation.counter import compilation_counter
19
from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
16
-from vllm.config import CompilationLevel, VllmConfig
20
+from vllm.config import CompilationLevel, VllmConfig, set_current_vllm_config
21
from vllm.logger import init_logger
22
from vllm.sequence import IntermediateTensors
23
from vllm.utils import resolve_obj_by_qualname, supports_dynamo
@@ -176,6 +180,33 @@ def support_torch_compile(
180
return cls_decorator_helper
181
182
183
+def _model_hash_key(fn) -> str:
184
+ import vllm
185
+
186
+ sha256_hash = hashlib.sha256()
187
+ sha256_hash.update(vllm.__version__.encode())
188
+ sha256_hash.update(fn.__qualname__.encode())
189
+ sha256_hash.update(str(fn.__code__.co_firstlineno).encode())
190
+ return sha256_hash.hexdigest()
191
+
192
+
193
+def _verify_source_unchanged(source_info, vllm_config) -> None:
194
+ from .caching import _compute_code_hash, _compute_code_hash_with_content
195
+
196
+ file_contents = {}
197
+ for source in source_info.inlined_sources:
198
+ module = sys.modules[source.module]
199
+ file = inspect.getfile(module)
200
+ vllm_config.compilation_config.traced_files.add(file)
201
+ file_contents[file] = source.content
202
+ expected_checksum = _compute_code_hash_with_content(file_contents)
203
+ actual_checksum = _compute_code_hash(set(file_contents.keys()))
204
+ if expected_checksum != actual_checksum:
205
+ raise RuntimeError(
206
+ "Source code has changed since the last compilation. Recompiling the model."
207
+ )
208
+
209
+
210
def _support_torch_compile(
211
cls: _T,
212
dynamic_arg_dims: dict[str, Union[int, list[int]]],
@@ -227,6 +258,64 @@ def _support_torch_compile(
258
if self.do_not_compile or torch.compiler.is_compiling():
259
return self.forward(*args, **kwargs)
260
261
+ if getattr(self, "aot_compiled_fn", None) is not None:
262
+ return self.aot_compiled_fn(self, *args, **kwargs)
263
+
264
+ cache_dir = None
265
+ aot_compilation_path = None
266
+ if envs.VLLM_USE_AOT_COMPILE:
267
+ """
268
+ When using torch.compile in AOT mode, we store the cache artifacts
269
+ under VLLM_CACHE_ROOT/torch_aot_compile/{hash}/rank_i_j. The {hash}
270
+ contains all of the factors except for the source files being
271
+ traced through, because we don't actually know which source files
272
+ to check at this point (before dynamo runs).
273
+ On loading we will actually look at the source files being traced
274
+ through. If any source file have changed (compared with the
275
+ serialized backend artifacts), then we need to generate a new AOT
276
+ compile artifact from scratch.
277
+ """
278
+ from .caching import compilation_config_hash_factors
279
+
280
+ factors: list[str] = compilation_config_hash_factors(self.vllm_config)
281
+
282
+ factors.append(_model_hash_key(self.forward))
283
+ hash_key = hashlib.sha256(str(factors).encode()).hexdigest()
284
+
285
+ cache_dir = os.path.join(
286
+ envs.VLLM_CACHE_ROOT,
287
+ "torch_aot_compile",
288
+ hash_key,
289
+ )
290
+
291
+ rank = self.vllm_config.parallel_config.rank
292
+ dp_rank = self.vllm_config.parallel_config.data_parallel_rank
293
+ cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
294
+ aot_compilation_path = os.path.join(cache_dir, "model")
295
+ try:
296
+ with (
297
+ set_current_vllm_config(self.vllm_config),
298
+ open(aot_compilation_path, "rb") as f,
299
+ ):
300
+ start_monitoring_torch_compile(self.vllm_config)
301
+ loaded_fn = torch.compiler.load_compiled_function(f)
302
+ _verify_source_unchanged(loaded_fn.source_info(), self.vllm_config)
303
+ self.aot_compiled_fn = loaded_fn
304
+ except Exception as e:
305
+ if os.path.exists(aot_compilation_path):
306
+ logger.warning(
307
+ "Cannot load aot compilation from path %s, error: %s",
308
+ aot_compilation_path,
309
+ str(e),
310
+ )
311
+ if envs.VLLM_FORCE_AOT_LOAD:
312
+ raise e
313
+ if getattr(self, "aot_compiled_fn", None) is not None:
314
+ logger.info(
315
+ "Directly load AOT compilation from path %s", aot_compilation_path
316
+ )
317
+ return self.aot_compiled_fn(self, *args, **kwargs)
318
+
319
# the first compilation needs to have dynamic shapes marked
320
if len(self.compiled_codes) < 1:
321
sig = inspect.signature(self.__class__.forward)
@@ -275,15 +364,15 @@ def _support_torch_compile(
364
)
365
366
# 2. every time Dynamo sees a function call, it will inline
278
- # the function by calling InliningInstructionTranslator.inline_call
367
+ # the function by calling InliningInstructionTranslator.inline_call_
368
# we hijack this function to know all the functions called
369
# during Dynamo tracing, and their corresponding files
281
- inline_call = InliningInstructionTranslator.inline_call
370
+ inline_call = InliningInstructionTranslator.inline_call_
371
283
- def patched_inline_call(parent, func, args, kwargs):
284
- code = func.get_code()
372
+ def patched_inline_call(self_):
373
+ code = self_.f_code
374
self.vllm_config.compilation_config.traced_files.add(code.co_filename)
286
- return inline_call(parent, func, args, kwargs)
375
+ return inline_call(self_)
376
377
# Disable the C++ compilation of symbolic shape guards. C++-fication
378
# of symbolic shape guards can improve guard overhead. But, since
@@ -300,13 +389,21 @@ def _support_torch_compile(
389
390
with (
391
patch.object(
303
- InliningInstructionTranslator, "inline_call", patched_inline_call
392
+ InliningInstructionTranslator, "inline_call_", patched_inline_call
393
),
394
torch._dynamo.config.patch(**dynamo_config_patches),
395
maybe_use_cudagraph_partition_wrapper(self.vllm_config),
396
_torch27_patch_tensor_subclasses(),
397
):
309
- output = self.compiled_callable(*args, **kwargs)
398
+ if envs.VLLM_USE_AOT_COMPILE:
399
+ self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
400
+ output = self.aot_compiled_fn(self, *args, **kwargs)
401
+ assert aot_compilation_path is not None
402
+ assert cache_dir is not None
403
+ os.makedirs(cache_dir, exist_ok=True)
404
+ self.aot_compiled_fn.save_compiled_function(aot_compilation_path)
405
+ else:
406
+ output = self.compiled_callable(*args, **kwargs)
407
return output
408
409
# usually, capturing the model once is enough, and then we can
vllm/compilation/wrapper.py
+23
@@ -10,6 +10,7 @@ from typing import Callable, Optional
10
11
import torch
12
13
+import vllm.envs as envs
14
from vllm.config import CompilationLevel, CUDAGraphMode, get_current_vllm_config
15
from vllm.logger import init_logger
16
@@ -44,6 +45,19 @@ class TorchCompileWrapperWithCustomDispatcher:
45
options = (
46
get_current_vllm_config().compilation_config.inductor_compile_config
47
)
48
+ if envs.VLLM_USE_AOT_COMPILE:
49
+ options = options or {}
50
+ # This effectively drop all the guards.
51
+ # We need this because bytecode hook is not used any more to
52
+ # drop guards in the AOT compile mode.
53
+ options["guard_filter_fn"] = lambda guards: [False for _ in guards]
54
+ if hasattr(torch._dynamo.config, "enable_aot_compile"):
55
+ torch._dynamo.config.enable_aot_compile = True
56
+ else:
57
+ msg = "torch._dynamo.config.enable_aot_compile is not "
58
+ msg += "available. AOT compile is disabled and please "
59
+ msg += "upgrade PyTorch version to use AOT compile."
60
+ logger.warning(msg)
61
62
compiled_callable = torch.compile(
63
self.forward, fullgraph=True, backend=backend, options=options
@@ -61,6 +75,15 @@ class TorchCompileWrapperWithCustomDispatcher:
75
compilation_level >= CompilationLevel.DYNAMO_ONCE
76
)
77
78
+ def aot_compile(self, *args, **kwargs):
79
+ if not hasattr(self.compiled_callable, "aot_compile"):
80
+ raise RuntimeError(
81
+ "aot_compile is not supported by the current configuration. "
82
+ + "Please make sure torch.compile is enabled with the latest "
83
+ + f"version of PyTorch (current using torch: {torch.__version__})"
84
+ )
85
+ return self.compiled_callable.aot_compile((args, kwargs))
86
+
87
def __call__(self, *args, **kwargs):
88
"""Implement the dispatch logic here, beyond the torch.compile level.
89
NOTE: this function can have additional arguments beyond the forward
vllm/envs.py
+17
@@ -89,6 +89,8 @@ if TYPE_CHECKING:
89
VLLM_TORCH_PROFILER_DIR: Optional[str] = None
90
VLLM_TORCH_PROFILER_RECORD_SHAPES: bool = False
91
VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
92
+ VLLM_USE_AOT_COMPILE: bool = False
93
+ VLLM_FORCE_AOT_LOAD: bool = False
94
VLLM_TORCH_PROFILER_WITH_STACK: bool = True
95
VLLM_TORCH_PROFILER_WITH_FLOPS: bool = False
96
VLLM_USE_TRITON_AWQ: bool = False
@@ -235,6 +237,13 @@ def maybe_convert_bool(value: Optional[str]) -> Optional[bool]:
237
return bool(int(value))
238
239
240
+def use_aot_compile() -> bool:
241
+ from vllm.utils import is_torch_equal_or_newer
242
+
243
+ default_value = "1" if is_torch_equal_or_newer("2.10.0.dev") else "0"
244
+ return os.environ.get("VLLM_USE_AOT_COMPILE", default_value) == "1"
245
+
246
+
247
def env_with_choices(
248
env_name: str,
249
default: Optional[str],
@@ -494,6 +503,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
503
# Dump fx graphs to the given directory.
504
# It will override CompilationConfig.debug_dump_path if set.
505
"VLLM_DEBUG_DUMP_PATH": lambda: os.environ.get("VLLM_DEBUG_DUMP_PATH", None),
506
+ # Feature flag to enable/disable AOT compilation. This will ensure
507
+ # compilation is done in warmup phase and the compilation will be
508
+ # reused in subsequent calls.
509
+ "VLLM_USE_AOT_COMPILE": use_aot_compile,
510
+ # Force vllm to always load AOT compiled models from disk. Failure
511
+ # to load will result in a hard error when this is enabled.
512
+ # Will be ignored when VLLM_USE_AOT_COMPILE is disabled.
513
+ "VLLM_FORCE_AOT_LOAD": lambda: os.environ.get("VLLM_FORCE_AOT_LOAD", "0") == "1",
514
# local rank of the process in the distributed setting, used to determine
515
# the GPU device id
516
"LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),