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 14, 05:30 UTC
[TPU] Rename tpu_commons to tpu_inference (#26279)
Signed-off-by: Utkarsh Sharma <utksharma@google.com> Co-authored-by: Utkarsh Sharma <utksharma@google.com> Co-authored-by: Chengji Yao <chengjiyao@google.com>
Utkarsh Sharma committed
Oct 8, 2025 at 12:00 UTC
335b28f7d1020911af7974f50e27b708113c5e3f
6 files changed
+22
-22
vllm/distributed/device_communicators/tpu_communicator.py
+7
-7
@@ -10,7 +10,7 @@ from torch.distributed import ProcessGroup
10
from vllm.config import get_current_vllm_config
11
from vllm.logger import init_logger
12
from vllm.platforms import current_platform
13
-from vllm.platforms.tpu import USE_TPU_COMMONS
13
+from vllm.platforms.tpu import USE_TPU_INFERENCE
14
15
from .base_device_communicator import DeviceCommunicatorBase
16
@@ -20,8 +20,8 @@ USE_RAY = parallel_config = (
20
21
logger = init_logger(__name__)
22
23
-if not USE_TPU_COMMONS:
24
- logger.info("tpu_commons not found, using vLLM's TpuCommunicator")
23
+if not USE_TPU_INFERENCE:
24
+ logger.info("tpu_inference not found, using vLLM's TpuCommunicator")
25
if current_platform.is_tpu():
26
import torch_xla
27
import torch_xla.core.xla_model as xm
@@ -100,9 +100,9 @@ class TpuCommunicator(DeviceCommunicatorBase):
100
return xm.all_gather(input_, dim=dim)
101
102
103
-if USE_TPU_COMMONS:
104
- from tpu_commons.distributed.device_communicators import (
105
- TpuCommunicator as TpuCommonsCommunicator,
103
+if USE_TPU_INFERENCE:
104
+ from tpu_inference.distributed.device_communicators import (
105
+ TpuCommunicator as TpuInferenceCommunicator,
106
)
107
108
- TpuCommunicator = TpuCommonsCommunicator # type: ignore
108
+ TpuCommunicator = TpuInferenceCommunicator # type: ignore
vllm/model_executor/model_loader/default_loader.py
+2
-2
@@ -223,9 +223,9 @@ class DefaultModelLoader(BaseModelLoader):
223
)
224
225
if current_platform.is_tpu():
226
- from vllm.platforms.tpu import USE_TPU_COMMONS
226
+ from vllm.platforms.tpu import USE_TPU_INFERENCE
227
228
- if not USE_TPU_COMMONS:
228
+ if not USE_TPU_INFERENCE:
229
# In PyTorch XLA, we should call `torch_xla.sync`
230
# frequently so that not too many ops are accumulated
231
# in the XLA program.
vllm/platforms/__init__.py
+1
-1
@@ -37,7 +37,7 @@ def tpu_platform_plugin() -> Optional[str]:
37
# Check for Pathways TPU proxy
38
if envs.VLLM_TPU_USING_PATHWAYS:
39
logger.debug("Confirmed TPU platform is available via Pathways proxy.")
40
- return "tpu_commons.platforms.tpu_jax.TpuPlatform"
40
+ return "tpu_inference.platforms.tpu_jax.TpuPlatform"
41
42
# Check for libtpu installation
43
try:
vllm/platforms/tpu.py
+5
-5
@@ -26,7 +26,7 @@ else:
26
27
logger = init_logger(__name__)
28
29
-USE_TPU_COMMONS = False
29
+USE_TPU_INFERENCE = False
30
31
32
class TpuPlatform(Platform):
@@ -254,10 +254,10 @@ class TpuPlatform(Platform):
254
255
256
try:
257
- from tpu_commons.platforms import TpuPlatform as TpuCommonsPlatform
257
+ from tpu_inference.platforms import TpuPlatform as TpuInferencePlatform
258
259
- TpuPlatform = TpuCommonsPlatform # type: ignore
260
- USE_TPU_COMMONS = True
259
+ TpuPlatform = TpuInferencePlatform # type: ignore
260
+ USE_TPU_INFERENCE = True
261
except ImportError:
262
- logger.info("tpu_commons not found, using vLLM's TpuPlatform")
262
+ logger.info("tpu_inference not found, using vLLM's TpuPlatform")
263
pass
vllm/v1/attention/backends/pallas.py
+1
-1
@@ -35,7 +35,7 @@ TPU_STR_DTYPE_TO_TORCH_DTYPE = {
35
}
36
37
try:
38
- import tpu_commons # noqa: F401
38
+ import tpu_inference # noqa: F401
39
except ImportError:
40
# Lazy import torch_xla
41
import torch_xla.core.xla_builder as xb
vllm/v1/worker/tpu_worker.py
+6
-6
@@ -23,7 +23,7 @@ from vllm.logger import init_logger
23
from vllm.lora.request import LoRARequest
24
from vllm.model_executor import set_random_seed
25
from vllm.platforms import current_platform
26
-from vllm.platforms.tpu import USE_TPU_COMMONS
26
+from vllm.platforms.tpu import USE_TPU_INFERENCE
27
from vllm.tasks import SupportedTask
28
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
29
from vllm.v1.core.sched.output import SchedulerOutput
@@ -36,8 +36,8 @@ logger = init_logger(__name__)
36
37
_R = TypeVar("_R")
38
39
-if not USE_TPU_COMMONS:
40
- logger.info("tpu_commons not found, using vLLM's TPUWorker.")
39
+if not USE_TPU_INFERENCE:
40
+ logger.info("tpu_inference not found, using vLLM's TPUWorker.")
41
import torch_xla.core.xla_model as xm
42
import torch_xla.debug.profiler as xp
43
import torch_xla.runtime as xr
@@ -346,7 +346,7 @@ class TPUWorker:
346
return fn(self.get_model())
347
348
349
-if USE_TPU_COMMONS:
350
- from tpu_commons.worker import TPUWorker as TPUCommonsWorker
349
+if USE_TPU_INFERENCE:
350
+ from tpu_inference.worker import TPUWorker as TpuInferenceWorker
351
352
- TPUWorker = TPUCommonsWorker # type: ignore
352
+ TPUWorker = TpuInferenceWorker # type: ignore