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, 12:15 UTC
[Misc]Use a platform independent interface to obtain the device attributes (#17100)
ponix-j committed
Apr 29, 2025 at 14:59 UTC
bdb2cddafc524380e0d04cf1aa32b41d9411d6fd
3 files changed
+8
-4
tests/conftest.py
+2
-1
@@ -293,7 +293,8 @@ class HfRunner:
293
def get_default_device(self):
294
from vllm.platforms import current_platform
295
296
- return ("cpu" if current_platform.is_cpu() else "cuda")
296
+ return ("cpu"
297
+ if current_platform.is_cpu() else current_platform.device_type)
298
299
def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
300
if x is None or isinstance(x, (bool, )):
tests/v1/sample/test_sampler.py
+3
-1
@@ -6,6 +6,7 @@ import numpy as np
6
import pytest
7
import torch
8
9
+from vllm.platforms import current_platform
10
from vllm.utils import make_tensor_with_pad
11
from vllm.v1.sample.metadata import SamplingMetadata
12
from vllm.v1.sample.sampler import Sampler
@@ -13,7 +14,8 @@ from vllm.v1.sample.sampler import Sampler
14
VOCAB_SIZE = 1024
15
NUM_OUTPUT_TOKENS = 20
16
CUDA_DEVICES = [
16
- f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
17
+ f"{current_platform.device_type}:{i}"
18
+ for i in range(1 if current_platform.device_count() == 1 else 2)
19
]
20
MAX_NUM_PROMPT_TOKENS = 64
21
vllm/worker/multi_step_model_runner.py
+3
-2
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.sampler import (PromptLogprobs, SampleLogprobs,
14
SamplerOutput,
15
SamplingMetadata, get_logprobs,
16
get_pythonized_sample_results)
17
+from vllm.platforms import current_platform
18
from vllm.sequence import (CompletionSequenceGroupOutput, IntermediateTensors,
19
Logprob, SequenceGroupMetadata, SequenceOutput)
20
from vllm.utils import PyObjectCache, async_tensor_h2d, current_stream
@@ -158,8 +159,8 @@ class StatefulModelInput(BroadcastableModelInput):
159
is_first_multi_step: bool = False
160
base_output_proc_callback: Optional[Callable] = None
161
# ping-pong data structures for multi-step to wait on the previous step
161
- step_cuda_events: List[torch.cuda.Event] = field(
162
- default_factory=lambda: [torch.cuda.Event(blocking=True)] * 2)
162
+ step_cuda_events: List[current_platform.Event] = field(
163
+ default_factory=lambda: [current_platform.Event(blocking=True)] * 2)
164
num_seqs: int = -1
165
num_queries: int = -1
166
num_single_step_prefills: int = 0