A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
[Hardware][TPU] Optionally import for TPU backend (#18269)
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>
Siyuan Liu committed
May 17, 2025 at 00:23 UTC
48ac2bed5b6271f82187de61245a85d987197c6f
3 files changed
+25
vllm/distributed/device_communicators/tpu_communicator.py
+9
@@ -91,3 +91,12 @@ class TpuCommunicator(DeviceCommunicatorBase):
91
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
92
assert dim == -1, "TPUs only support dim=-1 for all-gather."
93
return xm.all_gather(input_, dim=dim)
94
+
95
+
96
+try:
97
+ from tpu_commons.distributed.device_communicators import (
98
+ TpuCommunicator as TpuCommonsCommunicator)
99
+ TpuCommunicator = TpuCommonsCommunicator # type: ignore
100
+except ImportError:
101
+ logger.info("tpu_commons not found, using vLLM's TpuCommunicator")
102
+ pass
vllm/platforms/tpu.py
+8
@@ -194,3 +194,11 @@ class TpuPlatform(Platform):
194
if params.sampling_type == SamplingType.RANDOM_SEED:
195
raise ValueError(
196
"Torch XLA does not support per-request seed.")
197
+
198
+
199
+try:
200
+ from tpu_commons.platforms import TpuPlatform as TpuCommonsPlatform
201
+ TpuPlatform = TpuCommonsPlatform # type: ignore
202
+except ImportError:
203
+ logger.info("tpu_commons not found, using vLLM's TpuPlatform")
204
+ pass
vllm/v1/worker/tpu_worker.py
+8
@@ -267,3 +267,11 @@ def init_tpu_worker_distributed_environment(
267
ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
268
parallel_config.pipeline_parallel_size,
269
parallel_config.enable_expert_parallel)
270
+
271
+
272
+try:
273
+ from tpu_commons.worker import TPUWorker as TPUCommonsWorker
274
+ TPUWorker = TPUCommonsWorker # type: ignore
275
+except ImportError:
276
+ logger.info("tpu_commons not found, using vLLM's TPUWorker.")
277
+ pass