biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[compile] Apply stored functorch config while finalizing loaded artifacts. (#36582)

Browse Source 
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
This commit is contained in:
Zhengxu Chen
2026-03-10 12:34:35 -04:00
committed by GitHub
parent f088a831dd
commit bdd8981dab
2 changed files with 15 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/compilation/caching.py
View File

@@ -369,8 +369,14 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
from vllm.compilation.backends import VllmBackend
saved_aot_autograd_config = self.aot_autograd_config
if saved_aot_autograd_config is not None:
functorch_ctx = torch._functorch.config.patch(saved_aot_autograd_config)
else:
functorch_ctx = contextlib.nullcontext()
vllm_backend = VllmBackend(vllm_config, self.prefix, self.is_encoder)
with tracing(TracingContext(self._fake_mode)):
with tracing(TracingContext(self._fake_mode)), functorch_ctx:
result = vllm_backend(self.graph_module, list(self.example_inputs))
self.optimized_call = result.optimized_call
self.vllm_backend = vllm_backend

32
vllm/compilation/piecewise_backend.py
View File

@@ -258,31 +258,15 @@ class PiecewiseBackend:
else:
args_list = get_fake_args_from_graph(self.graph)
# TODO(https://github.com/vllm-project/vllm/issues/35766)
# Can we remove strict_autograd_cache and
# force_non_lazy_backward_lowering overrides?
# I added them explicitly because this is what they are
# set to before the refactor
# (https://github.com/vllm-project/vllm/pull/35472).
# They affect the aotautograd cache key computation
# but they shouldn't have any effect on the actual
# compilation.
config_patches = dict(
bundled_autograd_cache=True,
strict_autograd_cache=False,
range_entry.runnable = self.vllm_backend.compiler_manager.compile(
self.graph,
args_list,
self.vllm_backend.inductor_config,
self.compilation_config,
compile_range=range_entry.compile_range,
graph_index=self.piecewise_compile_index,
num_graphs=self.total_piecewise_compiles,
)
if hasattr(torch._functorch.config, "force_non_lazy_backward_lowering"):
config_patches["force_non_lazy_backward_lowering"] = False
with torch._functorch.config.patch(**config_patches):
range_entry.runnable = self.vllm_backend.compiler_manager.compile(
self.graph,
args_list,
self.vllm_backend.inductor_config,
self.compilation_config,
compile_range=range_entry.compile_range,
graph_index=self.piecewise_compile_index,
num_graphs=self.total_piecewise_compiles,
)
range_entry.compiled = True