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

[TPU] Rename tpu_commons to tpu_inference (#26279)

Browse Source 
Signed-off-by: Utkarsh Sharma <utksharma@google.com>
Co-authored-by: Utkarsh Sharma <utksharma@google.com>
Co-authored-by: Chengji Yao <chengjiyao@google.com>
This commit is contained in:
Utkarsh Sharma
2025-10-08 12:00:52 +05:30
committed by GitHub
parent 5e65d6b2ad
commit 335b28f7d1
6 changed files with 22 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -10,7 +10,7 @@ from torch.distributed import ProcessGroup
from vllm.config import get_current_vllm_config
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.platforms.tpu import USE_TPU_COMMONS
from vllm.platforms.tpu import USE_TPU_INFERENCE
from .base_device_communicator import DeviceCommunicatorBase
@@ -20,8 +20,8 @@ USE_RAY = parallel_config = (
logger = init_logger(__name__)
if not USE_TPU_COMMONS:
logger.info("tpu_commons not found, using vLLM's TpuCommunicator")
if not USE_TPU_INFERENCE:
logger.info("tpu_inference not found, using vLLM's TpuCommunicator")
if current_platform.is_tpu():
import torch_xla
import torch_xla.core.xla_model as xm
@@ -100,9 +100,9 @@ class TpuCommunicator(DeviceCommunicatorBase):
return xm.all_gather(input_, dim=dim)
if USE_TPU_COMMONS:
from tpu_commons.distributed.device_communicators import (
TpuCommunicator as TpuCommonsCommunicator,
if USE_TPU_INFERENCE:
from tpu_inference.distributed.device_communicators import (
TpuCommunicator as TpuInferenceCommunicator,
)
TpuCommunicator = TpuCommonsCommunicator # type: ignore
TpuCommunicator = TpuInferenceCommunicator # type: ignore