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

[TPU] add tpu_inference (#27277)

Browse Source 
Signed-off-by: Johnny Yang <johnnyyang@google.com>
This commit is contained in:
Johnny Yang
2025-11-26 14:46:36 -08:00
committed by GitHub
parent 56539cddac
commit ba1fcd84a7
4 changed files with 5 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
requirements/tpu.txt
View File

@@ -12,6 +12,4 @@ ray[data]
setuptools==78.1.0 setuptools==78.1.0
nixl==0.3.0 nixl==0.3.0
tpu_info==0.4.0 tpu_info==0.4.0
tpu-inference==0.11.1
# Install torch_xla
torch_xla[tpu, pallas]==2.8.0

8
vllm/distributed/device_communicators/tpu_communicator.py
View File

@@ -97,11 +97,3 @@ class TpuCommunicator(DeviceCommunicatorBase):
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
assert dim == -1, "TPUs only support dim=-1 for all-gather." assert dim == -1, "TPUs only support dim=-1 for all-gather."
return xm.all_gather(input_, dim=dim) return xm.all_gather(input_, dim=dim)
if USE_TPU_INFERENCE:
from tpu_inference.distributed.device_communicators import (
TpuCommunicator as TpuInferenceCommunicator,
)
TpuCommunicator = TpuInferenceCommunicator # type: ignore

4
vllm/platforms/tpu.py
View File

@@ -267,7 +267,9 @@ class TpuPlatform(Platform):
try: try:
from tpu_inference.platforms import TpuPlatform as TpuInferencePlatform from tpu_inference.platforms.tpu_platforms import (
TpuPlatform as TpuInferencePlatform,
)
TpuPlatform = TpuInferencePlatform # type: ignore TpuPlatform = TpuInferencePlatform # type: ignore
USE_TPU_INFERENCE = True USE_TPU_INFERENCE = True

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

@@ -346,6 +346,6 @@ class TPUWorker:
if USE_TPU_INFERENCE: if USE_TPU_INFERENCE:
from tpu_inference.worker import TPUWorker as TpuInferenceWorker from tpu_inference.worker.tpu_worker import TPUWorker as TpuInferenceWorker
TPUWorker = TpuInferenceWorker # type: ignore TPUWorker = TpuInferenceWorker # type: ignore