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 11, 15:00 UTC
[compile] Stop one-off setting enable_aot_compile and use context manager instead. (#30503)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Dec 11, 2025 at 15:28 UTC
92fea56fd1e148a5650160427d6b5c733ff211b8
1 file changed
+10
-8
vllm/compilation/wrapper.py
+10
-8
@@ -171,22 +171,24 @@ class TorchCompileWithNoGuardsWrapper:
171
172
compiled_ptr = self.check_invariants_and_forward
173
174
+ aot_context = nullcontext()
175
if envs.VLLM_USE_AOT_COMPILE:
176
if hasattr(torch._dynamo.config, "enable_aot_compile"):
176
- torch._dynamo.config.enable_aot_compile = True
177
+ aot_context = torch._dynamo.config.patch(enable_aot_compile=True)
178
else:
179
msg = "torch._dynamo.config.enable_aot_compile is not "
180
msg += "available. AOT compile is disabled and please "
181
msg += "upgrade PyTorch version to use AOT compile."
182
logger.warning(msg)
183
183
- self._compiled_callable = torch.compile(
184
- compiled_ptr,
185
- fullgraph=True,
186
- dynamic=False,
187
- backend=backend,
188
- options=options,
189
- )
184
+ with aot_context:
185
+ self._compiled_callable = torch.compile(
186
+ compiled_ptr,
187
+ fullgraph=True,
188
+ dynamic=False,
189
+ backend=backend,
190
+ options=options,
191
+ )
192
193
if envs.VLLM_USE_BYTECODE_HOOK and mode != CompilationMode.STOCK_TORCH_COMPILE:
194
torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook)