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

[Hardware][TPU] Optionally import for TPU backend (#18269)

Browse Source 
Signed-off-by: Siyuan Liu <lsiyuan@google.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Co-authored-by: Carol Zheng <cazheng@google.com>
Co-authored-by: Jade Zheng <zheng.shoujian@outlook.com>
Co-authored-by: Hongmin Fan <fanhongmin@google.com>
This commit is contained in:
Siyuan Liu
2025-05-17 00:23:12 -07:00
committed by GitHub
parent 3e0d435027
commit 48ac2bed5b
3 changed files with 25 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/v1/worker/tpu_worker.py
View File

@@ -267,3 +267,11 @@ def init_tpu_worker_distributed_environment(
ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
parallel_config.pipeline_parallel_size,
parallel_config.enable_expert_parallel)
try:
from tpu_commons.worker import TPUWorker as TPUCommonsWorker
TPUWorker = TPUCommonsWorker # type: ignore
except ImportError:
logger.info("tpu_commons not found, using vLLM's TPUWorker.")
pass