[torch.compile] integration with compilation control (#9058)

2024-10-10 12:39:36 -07:00
parent 78c0b4166c
commit e4d652ea3e
22 changed files with 404 additions and 98 deletions
--- a/tests/tpu/test_compilation.py
+++ b/tests/tpu/test_compilation.py
@@ -5,9 +5,11 @@ import tempfile

 import depyf

+from vllm.compilation.levels import CompilationLevel
+
 # disable custom dispatcher, let Dynamo takes over
 # all the control
-os.environ['VLLM_DYNAMO_USE_CUSTOM_DISPATCHER'] = "0"
+os.environ['VLLM_TORCH_COMPILE_LEVEL'] = str(CompilationLevel.DYNAMO_AS_IS)

 temp_dir = tempfile.mkdtemp()
 with depyf.prepare_debug(temp_dir):