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

[compile] Stop one-off setting enable_aot_compile and use context manager instead. (#30503)

Browse Source 
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
This commit is contained in:
Zhengxu Chen
2025-12-11 15:28:03 -05:00
committed by GitHub
parent e458270a95
commit 92fea56fd1
1 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/compilation/wrapper.py
View File

@@ -171,15 +171,17 @@ class TorchCompileWithNoGuardsWrapper:
compiled_ptr = self.check_invariants_and_forward compiled_ptr = self.check_invariants_and_forward
aot_context = nullcontext()
if envs.VLLM_USE_AOT_COMPILE: if envs.VLLM_USE_AOT_COMPILE:
if hasattr(torch._dynamo.config, "enable_aot_compile"): if hasattr(torch._dynamo.config, "enable_aot_compile"):
torch._dynamo.config.enable_aot_compile = True aot_context = torch._dynamo.config.patch(enable_aot_compile=True)
else: else:
msg = "torch._dynamo.config.enable_aot_compile is not " msg = "torch._dynamo.config.enable_aot_compile is not "
msg += "available. AOT compile is disabled and please " msg += "available. AOT compile is disabled and please "
msg += "upgrade PyTorch version to use AOT compile." msg += "upgrade PyTorch version to use AOT compile."
logger.warning(msg) logger.warning(msg)
with aot_context:
self._compiled_callable = torch.compile( self._compiled_callable = torch.compile(
compiled_ptr, compiled_ptr,
fullgraph=True, fullgraph=True,