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, 22:35 UTC
[TPU] Support Pathways in vLLM (#21417)
Signed-off-by: wenxindongwork <wenxindong@google.com>
wenxindongwork committed
Jul 30, 2025 at 10:02 UTC
8f0d5167155247934d247eb10ae086108db8d473
2 files changed
+17
-6
vllm/envs.py
+5
@@ -124,6 +124,7 @@ if TYPE_CHECKING:
124
VLLM_V1_USE_OUTLINES_CACHE: bool = False
125
VLLM_TPU_BUCKET_PADDING_GAP: int = 0
126
VLLM_TPU_MOST_MODEL_LEN: Optional[int] = None
127
+ VLLM_TPU_USING_PATHWAYS: bool = False
128
VLLM_USE_DEEP_GEMM: bool = False
129
VLLM_USE_FLASHINFER_MOE_FP8: bool = False
130
VLLM_USE_FLASHINFER_MOE_FP4: bool = False
@@ -900,6 +901,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
901
"VLLM_TPU_MOST_MODEL_LEN":
902
lambda: maybe_convert_int(os.environ.get("VLLM_TPU_MOST_MODEL_LEN", None)),
903
904
+ # Whether using Pathways
905
+ "VLLM_TPU_USING_PATHWAYS":
906
+ lambda: bool("proxy" in os.getenv("JAX_PLATFORMS", "").lower()),
907
+
908
# Allow use of DeepGemm kernels for fused moe ops.
909
"VLLM_USE_DEEP_GEMM":
910
lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM", "0"))),
vllm/platforms/__init__.py
+12
-6
@@ -1,11 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
import logging
4
import traceback
5
from itertools import chain
6
from typing import TYPE_CHECKING, Optional
7
8
+from vllm import envs
9
from vllm.plugins import load_plugins_by_group
10
from vllm.utils import resolve_obj_by_qualname, supports_xccl
11
@@ -31,20 +31,26 @@ def vllm_version_matches_substr(substr: str) -> bool:
31
32
33
def tpu_platform_plugin() -> Optional[str]:
34
- is_tpu = False
34
logger.debug("Checking if TPU platform is available.")
35
+
36
+ # Check for Pathways TPU proxy
37
+ if envs.VLLM_TPU_USING_PATHWAYS:
38
+ logger.debug("Confirmed TPU platform is available via Pathways proxy.")
39
+ return "tpu_commons.platforms.tpu_jax.TpuPlatform"
40
+
41
+ # Check for libtpu installation
42
try:
43
# While it's technically possible to install libtpu on a
44
# non-TPU machine, this is a very uncommon scenario. Therefore,
39
- # we assume that libtpu is installed if and only if the machine
45
+ # we assume that libtpu is installed only if the machine
46
# has TPUs.
47
+
48
import libtpu # noqa: F401
42
- is_tpu = True
49
logger.debug("Confirmed TPU platform is available.")
50
+ return "vllm.platforms.tpu.TpuPlatform"
51
except Exception as e:
52
logger.debug("TPU platform is not available because: %s", str(e))
46
-
47
- return "vllm.platforms.tpu.TpuPlatform" if is_tpu else None
53
+ return None
54
55
56
def cuda_platform_plugin() -> Optional[str]: