Skip to content

GitLab

Unverified Commit 92fea56f authored by Zhengxu Chen's avatar Zhengxu Chen Committed by GitHub
Browse files

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


Signed-off-by: default avatarzhxchen17 <zhxchen17@fb.com>
parent e458270a
Show whitespace changes
Inline Side-by-side
Showing with 10 additions and 8 deletions
vllm/compilation/wrapper.py
View file @ 92fea56f
...@@ -171,15 +171,17 @@ class TorchCompileWithNoGuardsWrapper: ...@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment