[4/N][torch.compile] clean up set_torch_compile_backend (#10401)

Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao
2024-11-17 23:57:20 -08:00
committed by GitHub
parent 47826cacf0
commit 51bb12d17b
7 changed files with 49 additions and 42 deletions

View File

@@ -3,8 +3,6 @@ from typing import TYPE_CHECKING
import torch
from vllm.plugins import set_torch_compile_backend
from .interface import Platform, PlatformEnum
if TYPE_CHECKING:
@@ -12,8 +10,6 @@ if TYPE_CHECKING:
else:
VllmConfig = None
set_torch_compile_backend("openxla")
class TpuPlatform(Platform):
_enum = PlatformEnum.TPU
@@ -38,3 +34,6 @@ class TpuPlatform(Platform):
compilation_config.level = CompilationLevel.DYNAMO_ONCE
assert compilation_config.level < CompilationLevel.PIECEWISE,\
"TPU does not support Inductor."
if compilation_config.backend == "":
compilation_config.backend = "openxla"