Skip to content

Commit bdb2cdd

Browse files
authored
[Misc]Use a platform independent interface to obtain the device attributes (#17100)
1 parent ebb3930 commit bdb2cdd

File tree

3 files changed

+8
-4
lines changed

3 files changed

+8
-4
lines changed

tests/conftest.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -293,7 +293,8 @@ class HfRunner:
293293
def get_default_device(self):
294294
from vllm.platforms import current_platform
295295

296-
return ("cpu" if current_platform.is_cpu() else "cuda")
296+
return ("cpu"
297+
if current_platform.is_cpu() else current_platform.device_type)
297298

298299
def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
299300
if x is None or isinstance(x, (bool, )):

tests/v1/sample/test_sampler.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,16 @@
66
import pytest
77
import torch
88

9+
from vllm.platforms import current_platform
910
from vllm.utils import make_tensor_with_pad
1011
from vllm.v1.sample.metadata import SamplingMetadata
1112
from vllm.v1.sample.sampler import Sampler
1213

1314
VOCAB_SIZE = 1024
1415
NUM_OUTPUT_TOKENS = 20
1516
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)
1719
]
1820
MAX_NUM_PROMPT_TOKENS = 64
1921

vllm/worker/multi_step_model_runner.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
SamplerOutput,
1515
SamplingMetadata, get_logprobs,
1616
get_pythonized_sample_results)
17+
from vllm.platforms import current_platform
1718
from vllm.sequence import (CompletionSequenceGroupOutput, IntermediateTensors,
1819
Logprob, SequenceGroupMetadata, SequenceOutput)
1920
from vllm.utils import PyObjectCache, async_tensor_h2d, current_stream
@@ -158,8 +159,8 @@ class StatefulModelInput(BroadcastableModelInput):
158159
is_first_multi_step: bool = False
159160
base_output_proc_callback: Optional[Callable] = None
160161
# 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)
163164
num_seqs: int = -1
164165
num_queries: int = -1
165166
num_single_step_prefills: int = 0

0 commit comments

Comments
 (0)