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 12, 12:15 UTC
[TPU] add tpu_inference (#27277)
Signed-off-by: Johnny Yang <johnnyyang@google.com>
Johnny Yang committed
Nov 26, 2025 at 14:46 UTC
ba1fcd84a7f1dc907c17bf4ba4fab6762a9f33a1
4 files changed
+5
-13
requirements/tpu.txt
+1
-3
@@ -12,6 +12,4 @@ ray[data]
12
setuptools==78.1.0
13
nixl==0.3.0
14
tpu_info==0.4.0
15
-
16
-# Install torch_xla
17
-torch_xla[tpu, pallas]==2.8.0
\ No newline at end of file
15
+tpu-inference==0.11.1
vllm/distributed/device_communicators/tpu_communicator.py
-8
@@ -97,11 +97,3 @@ class TpuCommunicator(DeviceCommunicatorBase):
97
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
98
assert dim == -1, "TPUs only support dim=-1 for all-gather."
99
return xm.all_gather(input_, dim=dim)
100
-
101
-
102
-if USE_TPU_INFERENCE:
103
- from tpu_inference.distributed.device_communicators import (
104
- TpuCommunicator as TpuInferenceCommunicator,
105
- )
106
-
107
- TpuCommunicator = TpuInferenceCommunicator # type: ignore
vllm/platforms/tpu.py
+3
-1
@@ -267,7 +267,9 @@ class TpuPlatform(Platform):
267
268
269
try:
270
- from tpu_inference.platforms import TpuPlatform as TpuInferencePlatform
270
+ from tpu_inference.platforms.tpu_platforms import (
271
+ TpuPlatform as TpuInferencePlatform,
272
+ )
273
274
TpuPlatform = TpuInferencePlatform # type: ignore
275
USE_TPU_INFERENCE = True
vllm/v1/worker/tpu_worker.py
+1
-1
@@ -346,6 +346,6 @@ class TPUWorker:
346
347
348
if USE_TPU_INFERENCE:
349
- from tpu_inference.worker import TPUWorker as TpuInferenceWorker
349
+ from tpu_inference.worker.tpu_worker import TPUWorker as TpuInferenceWorker
350
351
TPUWorker = TpuInferenceWorker # type: ignore