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 6, 08:15 UTC
[compile] Fix aot test failures with torch 2.12. (#37604)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Mar 20, 2026 at 12:06 UTC
c0f5fae601cf2649dec3cb06ad80008ced7a46ea
2 files changed
+54
-35
tests/compile/test_aot_compile.py
+13
-4
@@ -14,6 +14,7 @@ from unittest.mock import Mock, patch
14
import pytest
15
import torch
16
17
+import vllm.envs as envs
18
import vllm.model_executor.layers.activation
19
from vllm.compilation.backends import VllmBackend
20
from vllm.compilation.caching import (
@@ -162,6 +163,9 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
163
164
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
165
def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch):
166
+ from torch._subclasses import FakeTensorMode
167
+ from torch.fx.experimental.symbolic_shapes import ShapeEnv
168
+
169
def foo(x: torch.Tensor):
170
return x[slice(0, x.shape[0])]
171
@@ -172,12 +176,13 @@ def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch):
176
gm = torch.fx.symbolic_trace(foo)
177
assert "getitem_1 = x[slice(0, getitem, None)]" in gm.code
178
with use_vllm_config(vllm_config):
175
- payload = VllmSerializableFunction.serialize_compile_artifacts(
176
- VllmSerializableFunction(gm, (example_input,), "", foo)
179
+ payload = VllmSerializableFunction.serialize_graph_module(gm)
180
+ fake_mode = FakeTensorMode(shape_env=ShapeEnv())
181
+ loaded_gm = VllmSerializableFunction.deserialize_graph_module(
182
+ payload, fake_mode
183
)
178
- fn = VllmSerializableFunction.deserialize_compile_artifacts(payload)
184
180
- assert gm.code == fn.graph_module.code
185
+ assert gm.code == loaded_gm.code
186
187
188
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
@@ -725,6 +730,10 @@ class TestStandaloneCompiledArtifactsIntegration:
730
]:
731
assert cache.get(submod, shape) == shared_data
732
733
+ @pytest.mark.skipif(
734
+ envs.VLLM_USE_MEGA_AOT_ARTIFACT,
735
+ reason="There's no AOT Autograd run with mega artifact",
736
+ )
737
def test_functorch_config(self):
738
vllm_config = make_vllm_config()
739
example_inputs = (torch.randn(10, 10),)
vllm/compilation/caching.py
+41
-31
@@ -11,6 +11,8 @@ from typing import Any, Literal
11
from unittest.mock import patch
12
13
import torch
14
+from torch._subclasses import FakeTensorMode
15
+from torch.fx._graph_pickler import GraphPickler, Options
16
from torch.utils import _pytree as pytree
17
18
import vllm.envs as envs
@@ -206,26 +208,8 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
208
return self.optimized_call(*args, **kwargs)
209
210
@classmethod
209
- def serialize_compile_artifacts(
210
- cls, compiled_fn: "VllmSerializableFunction"
211
- ) -> bytes:
211
+ def serialize_graph_module(cls, graph_module: torch.fx.GraphModule) -> bytes:
212
import sympy
213
- from torch._subclasses import FakeTensorMode
214
- from torch.fx._graph_pickler import GraphPickler, Options
215
-
216
- state = compiled_fn.__dict__.copy()
217
- state.pop("optimized_call")
218
- state.pop("shape_env")
219
- state.pop("vllm_backend", None)
220
- state.pop("_fake_mode", None)
221
- for node in state["graph_module"].graph.nodes:
222
- node.meta.pop("source_fn_stack", None)
223
- node.meta.pop("nn_module_stack", None)
224
- for name, submod in state["graph_module"].named_children():
225
- if hasattr(submod, "graph"):
226
- for node in submod.graph.nodes:
227
- node.meta.pop("source_fn_stack", None)
228
- node.meta.pop("nn_module_stack", None)
213
214
graph_reducer_override = GraphPickler.reducer_override
215
@@ -242,6 +226,37 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
226
return type(None), ()
227
return graph_reducer_override(self, obj)
228
229
+ with (
230
+ patch.object(GraphPickler, "reducer_override", _graph_reducer_override),
231
+ patch_pytree_map_over_slice(),
232
+ ):
233
+ return GraphPickler.dumps(graph_module, Options(ops_filter=None))
234
+
235
+ @classmethod
236
+ def deserialize_graph_module(
237
+ cls, data: bytes, fake_mode: FakeTensorMode
238
+ ) -> torch.fx.GraphModule:
239
+ with patch_pytree_map_over_slice():
240
+ return GraphPickler.loads(data, fake_mode)
241
+
242
+ @classmethod
243
+ def serialize_compile_artifacts(
244
+ cls, compiled_fn: "VllmSerializableFunction"
245
+ ) -> bytes:
246
+ state = compiled_fn.__dict__.copy()
247
+ state.pop("optimized_call")
248
+ state.pop("shape_env")
249
+ state.pop("vllm_backend", None)
250
+ state.pop("_fake_mode", None)
251
+ for node in state["graph_module"].graph.nodes:
252
+ node.meta.pop("source_fn_stack", None)
253
+ node.meta.pop("nn_module_stack", None)
254
+ for name, submod in state["graph_module"].named_children():
255
+ if hasattr(submod, "graph"):
256
+ for node in submod.graph.nodes:
257
+ node.meta.pop("source_fn_stack", None)
258
+ node.meta.pop("nn_module_stack", None)
259
+
260
if state.get("sym_tensor_indices"):
261
# put tensor inputs on meta device since their data
262
# isn't needed, yet we need the meta for make_copy_and_call
@@ -257,14 +272,9 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
272
lambda inp: torch.empty_like(inp, device="meta"),
273
state["example_inputs"],
274
)
260
- with (
261
- patch.object(GraphPickler, "reducer_override", _graph_reducer_override),
262
- patch_pytree_map_over_slice(),
263
- ):
264
- state["graph_module"] = GraphPickler.dumps(
265
- state["graph_module"], Options(ops_filter=None)
266
- )
267
- state["example_inputs"] = GraphPickler.dumps(state["example_inputs"])
275
+
276
+ state["graph_module"] = cls.serialize_graph_module(state["graph_module"])
277
+ state["example_inputs"] = GraphPickler.dumps(state["example_inputs"])
278
279
if compiled_fn.vllm_backend:
280
(
@@ -280,14 +290,14 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
290
@classmethod
291
def deserialize_compile_artifacts(cls, data: bytes) -> "VllmSerializableFunction":
292
from torch._guards import TracingContext, tracing
283
- from torch._subclasses import FakeTensorMode
284
- from torch.fx._graph_pickler import GraphPickler
293
from torch.fx.experimental.symbolic_shapes import ShapeEnv
294
295
state = pickle.loads(data)
296
fake_mode = FakeTensorMode(shape_env=ShapeEnv())
289
- with patch_pytree_map_over_slice():
290
- state["graph_module"] = GraphPickler.loads(state["graph_module"], fake_mode)
297
+
298
+ state["graph_module"] = cls.deserialize_graph_module(
299
+ state["graph_module"], fake_mode
300
+ )
301
state["graph_module"].recompile()
302
state["example_inputs"] = GraphPickler.loads(state["example_inputs"], fake_mode)
303