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
[Misc] Refactor platform to get device specific stream and event (#14411)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
Apr 21, 2025 at 21:25 UTC
7272bfae77310c417e0b747a0590d9dd782c9f6b
2 files changed
+15
-5
vllm/platforms/interface.py
+9
@@ -404,6 +404,15 @@ class Platform:
404
) -> None:
405
"""Raises if this request is unsupported on this platform"""
406
407
+ def __getattr__(self, key: str):
408
+ device = getattr(torch, self.device_name, None)
409
+ if device is not None and hasattr(device, key):
410
+ return getattr(device, key)
411
+ else:
412
+ logger.warning("Current platform %s doesn't has '%s' attribute.",
413
+ self.device_name, key)
414
+ return None
415
+
416
417
class UnspecifiedPlatform(Platform):
418
_enum = PlatformEnum.UNSPECIFIED
vllm/spec_decode/metrics.py
+6
-5
@@ -8,6 +8,7 @@ import torch
8
9
from vllm.model_executor.layers.spec_decode_base_sampler import (
10
SpecDecodeBaseSampler)
11
+from vllm.platforms import current_platform
12
from vllm.utils import is_pin_memory_available
13
14
@@ -89,14 +90,14 @@ class AsyncMetricsCollector:
90
self._rank = rank
91
if isinstance(device_type, torch.device):
92
device_type = device_type.type
92
- if device_type == 'cuda':
93
- self._copy_stream = torch.cuda.Stream()
93
+ stream = current_platform.Stream
94
+ if stream is not None:
95
+ self._copy_stream = stream()
96
97
def maybe_collect_rejsample_metrics(
98
self, k: int) -> Optional[SpecDecodeWorkerMetrics]:
97
- # currently using cuda.Event, skip for any non_cuda_alike platform
98
- from vllm.platforms import current_platform
99
- if not current_platform.is_cuda_alike():
99
+ # Skip for any platform that doesn't have device Event
100
+ if current_platform.Event is None:
101
return None
102
103
# If a copy was initiated in the previous call, collect and return.