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, 23:00 UTC
[compile] Apply stored functorch config while finalizing loaded artifacts. (#36582)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Mar 10, 2026 at 12:34 UTC
bdd8981dab8d8c6ae88a3f605d04ec5243088e5a
2 files changed
+15
-25
vllm/compilation/caching.py
+7
-1
@@ -369,8 +369,14 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
369
370
from vllm.compilation.backends import VllmBackend
371
372
+ saved_aot_autograd_config = self.aot_autograd_config
373
+ if saved_aot_autograd_config is not None:
374
+ functorch_ctx = torch._functorch.config.patch(saved_aot_autograd_config)
375
+ else:
376
+ functorch_ctx = contextlib.nullcontext()
377
+
378
vllm_backend = VllmBackend(vllm_config, self.prefix, self.is_encoder)
373
- with tracing(TracingContext(self._fake_mode)):
379
+ with tracing(TracingContext(self._fake_mode)), functorch_ctx:
380
result = vllm_backend(self.graph_module, list(self.example_inputs))
381
self.optimized_call = result.optimized_call
382
self.vllm_backend = vllm_backend
vllm/compilation/piecewise_backend.py
+8
-24
@@ -258,31 +258,15 @@ class PiecewiseBackend:
258
else:
259
args_list = get_fake_args_from_graph(self.graph)
260
261
- # TODO(https://github.com/vllm-project/vllm/issues/35766)
262
- # Can we remove strict_autograd_cache and
263
- # force_non_lazy_backward_lowering overrides?
264
- # I added them explicitly because this is what they are
265
- # set to before the refactor
266
- # (https://github.com/vllm-project/vllm/pull/35472).
267
- # They affect the aotautograd cache key computation
268
- # but they shouldn't have any effect on the actual
269
- # compilation.
270
- config_patches = dict(
271
- bundled_autograd_cache=True,
272
- strict_autograd_cache=False,
261
+ range_entry.runnable = self.vllm_backend.compiler_manager.compile(
262
+ self.graph,
263
+ args_list,
264
+ self.vllm_backend.inductor_config,
265
+ self.compilation_config,
266
+ compile_range=range_entry.compile_range,
267
+ graph_index=self.piecewise_compile_index,
268
+ num_graphs=self.total_piecewise_compiles,
269
)
274
- if hasattr(torch._functorch.config, "force_non_lazy_backward_lowering"):
275
- config_patches["force_non_lazy_backward_lowering"] = False
276
- with torch._functorch.config.patch(**config_patches):
277
- range_entry.runnable = self.vllm_backend.compiler_manager.compile(
278
- self.graph,
279
- args_list,
280
- self.vllm_backend.inductor_config,
281
- self.compilation_config,
282
- compile_range=range_entry.compile_range,
283
- graph_index=self.piecewise_compile_index,
284
- num_graphs=self.total_piecewise_compiles,
285
- )
270
271
range_entry.compiled = True
272