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 7, 04:00 UTC
[compile] Consistent compiler config for saved/loaded vllm backends. (#35810)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Mar 5, 2026 at 15:08 UTC
a97954b6a8fa41a162ebf58f80a1460a98e0baf0
2 files changed
+65
-8
tests/compile/test_aot_compile.py
+42
@@ -14,6 +14,7 @@ import pytest
14
import torch
15
16
import vllm.model_executor.layers.activation
17
+from vllm.compilation.backends import VllmBackend
18
from vllm.compilation.caching import (
19
StandaloneCompiledArtifacts,
20
VllmSerializableFunction,
@@ -721,3 +722,44 @@ class TestStandaloneCompiledArtifactsIntegration:
722
("mod3", "shape3"),
723
]:
724
assert cache.get(submod, shape) == shared_data
725
+
726
+ def test_functorch_config(self):
727
+ vllm_config = make_vllm_config()
728
+ example_inputs = (torch.randn(10, 10),)
729
+
730
+ def add_1(x: torch.Tensor):
731
+ return x + 1
732
+
733
+ gm = torch._dynamo.functional_export.dynamo_graph_capture_for_export(add_1)(
734
+ *example_inputs
735
+ )
736
+
737
+ gm.graph._codegen = torch.fx.graph.CodeGen()
738
+ gm._dynamo_bytecode_flatten = None
739
+ gm._dynamo_bytecode_unflatten = None
740
+
741
+ with (
742
+ torch._functorch.config.patch(bundled_autograd_cache=False),
743
+ set_current_vllm_config(vllm_config),
744
+ ):
745
+ with torch._functorch.config.patch(bundled_autograd_cache=True):
746
+ fn = VllmSerializableFunction(gm, example_inputs, "", add_1)
747
+
748
+ payload = VllmSerializableFunction.serialize_compile_artifacts(fn)
749
+
750
+ config = None
751
+
752
+ def backend(*args, **kwargs) -> VllmSerializableFunction:
753
+ nonlocal config
754
+ # bundled_autograd_cache should be True even compiler backend
755
+ # runs with bundled_autograd_cache=False in ambient context.
756
+ config = torch._functorch.config.save_config_portable()
757
+ return fn
758
+
759
+ loaded_fn = VllmSerializableFunction.deserialize_compile_artifacts(payload)
760
+ with patch.object(VllmBackend, "__call__", backend):
761
+ loaded_fn(*example_inputs)
762
+
763
+ assert isinstance(config, dict)
764
+ assert "bundled_autograd_cache" in config
765
+ assert config["bundled_autograd_cache"] is True
vllm/compilation/caching.py
+23
-8
@@ -178,6 +178,7 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
178
is_encoder: bool = False,
179
vllm_backend: Any | None = None,
180
sym_tensor_indices: list[int] | None = None,
181
+ aot_autograd_config: dict[str, Any] | None = None,
182
) -> None:
183
assert isinstance(graph_module, torch.fx.GraphModule)
184
self.graph_module = graph_module
@@ -188,6 +189,13 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
189
self.shape_env = None
190
self.vllm_backend = vllm_backend
191
self.sym_tensor_indices = sym_tensor_indices
192
+
193
+ import torch._functorch.config as functorch_config
194
+
195
+ self.aot_autograd_config = (
196
+ aot_autograd_config or functorch_config.save_config_portable()
197
+ )
198
+
199
sym_input = next(
200
(i for i in self.example_inputs if isinstance(i, torch.SymInt)), None
201
)
@@ -286,6 +294,12 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
294
sym_shape_indices_map = state.pop("sym_shape_indices_map", {})
295
returns_tuple_map = state.pop("returns_tuple_map", {})
296
297
+ saved_aot_autograd_config = state["aot_autograd_config"]
298
+ if saved_aot_autograd_config is not None:
299
+ functorch_ctx = torch._functorch.config.patch(saved_aot_autograd_config)
300
+ else:
301
+ functorch_ctx = contextlib.nullcontext()
302
+
303
if envs.VLLM_USE_MEGA_AOT_ARTIFACT:
304
assert standalone_compile_artifacts is not None
305
submod_names = standalone_compile_artifacts.submodule_names()
@@ -299,13 +313,14 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
313
num_submods,
314
)
315
302
- fn = reconstruct_serializable_fn_from_mega_artifact(
303
- state=state,
304
- standalone_compile_artifacts=standalone_compile_artifacts,
305
- vllm_config=get_current_vllm_config(),
306
- sym_shape_indices_map=sym_shape_indices_map,
307
- returns_tuple_map=returns_tuple_map,
308
- )
316
+ with functorch_ctx:
317
+ fn = reconstruct_serializable_fn_from_mega_artifact(
318
+ state=state,
319
+ standalone_compile_artifacts=standalone_compile_artifacts,
320
+ vllm_config=get_current_vllm_config(),
321
+ sym_shape_indices_map=sym_shape_indices_map,
322
+ returns_tuple_map=returns_tuple_map,
323
+ )
324
325
logger.info(
326
"reconstructed serializable fn from standalone compile artifacts"
@@ -328,7 +343,7 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
343
vllm_backend: VllmBackend = VllmBackend(
344
vllm_config, state["prefix"], is_encoder
345
)
331
- with tracing(TracingContext(fake_mode)):
346
+ with tracing(TracingContext(fake_mode)), functorch_ctx:
347
fn.optimized_call = vllm_backend(
348
state["graph_module"], compile_inputs
349
).optimized_call