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
[Misc] Add TPU usage report when using tpu_inference. (#27423)
Signed-off-by: Hongmin Fan <fanhongmin@google.com>
hfan committed
Oct 23, 2025 at 23:29 UTC
8dbe0c527fa76cd908bb6287f9f501df44e04473
1 file changed
+30
-10
vllm/usage/usage_lib.py
+30
-10
@@ -176,6 +176,32 @@ class UsageMessage:
176
self._report_usage_once(model_architecture, usage_context, extra_kvs)
177
self._report_continuous_usage()
178
179
+ def _report_tpu_inference_usage(self) -> bool:
180
+ try:
181
+ from tpu_inference import tpu_info, utils
182
+
183
+ self.gpu_count = tpu_info.get_num_chips()
184
+ self.gpu_type = tpu_info.get_tpu_type()
185
+ self.gpu_memory_per_device = utils.get_device_hbm_limit()
186
+ self.cuda_runtime = "tpu_inference"
187
+ return True
188
+ except Exception:
189
+ return False
190
+
191
+ def _report_torch_xla_usage(self) -> bool:
192
+ try:
193
+ import torch_xla
194
+
195
+ self.gpu_count = torch_xla.runtime.world_size()
196
+ self.gpu_type = torch_xla.tpu.get_tpu_type()
197
+ self.gpu_memory_per_device = torch_xla.core.xla_model.get_memory_info()[
198
+ "bytes_limit"
199
+ ]
200
+ self.cuda_runtime = "torch_xla"
201
+ return True
202
+ except Exception:
203
+ return False
204
+
205
def _report_usage_once(
206
self,
207
model_architecture: str,
@@ -192,16 +218,10 @@ class UsageMessage:
218
)
219
if current_platform.is_cuda():
220
self.cuda_runtime = torch.version.cuda
195
- if current_platform.is_tpu():
196
- try:
197
- import torch_xla
198
-
199
- self.gpu_count = torch_xla.runtime.world_size()
200
- self.gpu_type = torch_xla.tpu.get_tpu_type()
201
- self.gpu_memory_per_device = torch_xla.core.xla_model.get_memory_info()[
202
- "bytes_limit"
203
- ]
204
- except Exception:
221
+ if current_platform.is_tpu(): # noqa: SIM102
222
+ if (not self._report_tpu_inference_usage()) and (
223
+ not self._report_torch_xla_usage()
224
+ ):
225
logger.exception("Failed to collect TPU information")
226
self.provider = _detect_cloud_provider()
227
self.architecture = platform.machine()