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 15, 13:30 UTC
[V1] Move usage stats to worker and start logging TPU hardware (#16211)
Daniel Li committed
Apr 25, 2025 at 13:06 UTC
48cb2109b61676ebc0e7e76022a0be51a41a08b8
6 files changed
+22
-10
vllm/usage/usage_lib.py
+9
@@ -174,6 +174,15 @@ class UsageMessage:
174
cuda_get_device_properties(0, ("name", "total_memory")))
175
if current_platform.is_cuda():
176
self.cuda_runtime = torch.version.cuda
177
+ if current_platform.is_tpu():
178
+ try:
179
+ import torch_xla
180
+ self.gpu_count = torch_xla.runtime.world_size()
181
+ self.gpu_type = torch_xla.tpu.get_tpu_type()
182
+ self.gpu_memory_per_device = (
183
+ torch_xla.core.xla_model.get_memory_info()["bytes_limit"])
184
+ except Exception:
185
+ pass
186
self.provider = _detect_cloud_provider()
187
self.architecture = platform.machine()
188
self.platform = platform.platform()
vllm/v1/engine/async_llm.py
-4
@@ -36,7 +36,6 @@ from vllm.v1.executor.abstract import Executor
36
from vllm.v1.metrics.loggers import (LoggingStatLogger, PrometheusStatLogger,
37
StatLoggerBase)
38
from vllm.v1.metrics.stats import IterationStats, SchedulerStats
39
-from vllm.v1.utils import report_usage_stats
39
40
logger = init_logger(__name__)
41
@@ -113,9 +112,6 @@ class AsyncLLM(EngineClient):
112
except RuntimeError:
113
pass
114
116
- # If usage stat is enabled, collect relevant info.
117
- report_usage_stats(vllm_config, usage_context)
118
-
115
@classmethod
116
def from_vllm_config(
117
cls,
vllm/v1/engine/llm_engine.py
-4
@@ -28,7 +28,6 @@ from vllm.v1.engine.output_processor import OutputProcessor
28
from vllm.v1.engine.parallel_sampling import ParentRequest
29
from vllm.v1.engine.processor import Processor
30
from vllm.v1.executor.abstract import Executor
31
-from vllm.v1.utils import report_usage_stats
31
32
logger = init_logger(__name__)
33
@@ -97,9 +96,6 @@ class LLMEngine:
96
# for v0 compatibility
97
self.model_executor = self.engine_core.engine_core.model_executor # type: ignore
98
100
- # If usage stat is enabled, collect relevant info.
101
- report_usage_stats(vllm_config, usage_context)
102
-
99
@classmethod
100
def from_vllm_config(
101
cls,
vllm/v1/utils.py
+3
-1
@@ -205,7 +205,9 @@ def copy_slice(from_tensor: torch.Tensor, to_tensor: torch.Tensor,
205
return to_tensor[:length].copy_(from_tensor[:length], non_blocking=True)
206
207
208
-def report_usage_stats(vllm_config, usage_context: UsageContext) -> None:
208
+def report_usage_stats(
209
+ vllm_config,
210
+ usage_context: UsageContext = UsageContext.ENGINE_CONTEXT) -> None:
211
"""Report usage statistics if enabled."""
212
213
if not is_usage_stats_enabled():
vllm/v1/worker/gpu_worker.py
+5
@@ -23,6 +23,7 @@ from vllm.platforms import current_platform
23
from vllm.utils import GiB_bytes
24
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
25
from vllm.v1.outputs import ModelRunnerOutput
26
+from vllm.v1.utils import report_usage_stats
27
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
28
from vllm.v1.worker.worker_base import WorkerBase
29
@@ -141,6 +142,10 @@ class Worker(WorkerBase):
142
self.model_runner: GPUModelRunner = GPUModelRunner(
143
self.vllm_config, self.device)
144
145
+ if self.rank == 0:
146
+ # If usage stat is enabled, collect relevant info.
147
+ report_usage_stats(self.vllm_config)
148
+
149
# FIXME(youkaichao & ywang96): Use TorchDispatchMode instead of memory pool
150
# to hijack tensor allocation.
151
def load_model(self) -> None:
vllm/v1/worker/tpu_worker.py
+5
-1
@@ -21,7 +21,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
21
from vllm.v1.kv_cache_interface import (AttentionSpec, KVCacheConfig,
22
KVCacheSpec)
23
from vllm.v1.outputs import ModelRunnerOutput
24
-from vllm.v1.utils import bind_kv_cache
24
+from vllm.v1.utils import bind_kv_cache, report_usage_stats
25
from vllm.v1.worker.tpu_model_runner import TPUModelRunner
26
27
logger = init_logger(__name__)
@@ -133,6 +133,10 @@ class TPUWorker:
133
# Init ModelRunner here, so that we have access to self.device.
134
self.model_runner = TPUModelRunner(self.vllm_config, self.device)
135
136
+ if rank == 0:
137
+ # If usage stat is enabled, collect relevant info.
138
+ report_usage_stats(self.vllm_config)
139
+
140
def determine_available_memory(self) -> int:
141
kv_caches: dict[str, torch.Tensor] = {}
142
kv_cache_spec = self.model_runner.get_kv_cache_spec()