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

18
vllm/compilation/wrapper.py
View File

@@ -171,22 +171,24 @@ 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)
self._compiled_callable = torch.compile( with aot_context:
compiled_ptr, self._compiled_callable = torch.compile(
fullgraph=True, compiled_ptr,
dynamic=False, fullgraph=True,
backend=backend, dynamic=False,
options=options, backend=backend,
) options=options,
)
if envs.VLLM_USE_BYTECODE_HOOK and mode != CompilationMode.STOCK_TORCH_COMPILE: if envs.VLLM_USE_BYTECODE_HOOK and mode != CompilationMode.STOCK_TORCH_COMPILE:
torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook) torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook)